"""CPU "Fluoddity-lite": render a compiled Rule to a phenotype image.

A faithful-in-spirit NumPy port of the Fluoddity inner loop (the lab manual's
Week 8 capstone). It exists so the evolutionary loop has a phenotype to select
on without touching the GPU app.

Pipeline per step (cf. shaders/entity_update.glsl + shaders/canvas.frag):
  1. body frame from velocity; two sensors offset by +/- angle at distance d
  2. sample the vector trail field at each sensor, project into the body frame
  3. evaluate the symmetrized Rule -> (force, strafe) in body coords
  4. drag + force -> velocity; velocity + strafe -> position; wrap boundary
  5. deposit velocity into the field; diffuse + persist the field
The phenotype image is the trail field rendered as HSV (hue = flow direction).
"""

from __future__ import annotations

import colorsys

import numpy as np

from fourier_rule import eval_rule


def _bilinear_wrap(field: np.ndarray, gx: np.ndarray, gy: np.ndarray) -> np.ndarray:
    """Sample field (G, G, C) at fractional grid coords (wrap). Returns (N, C)."""
    G = field.shape[0]
    x0 = np.floor(gx).astype(np.int64)
    y0 = np.floor(gy).astype(np.int64)
    fx = (gx - x0)[:, None]
    fy = (gy - y0)[:, None]
    x0m, y0m = x0 % G, y0 % G
    x1m, y1m = (x0 + 1) % G, (y0 + 1) % G
    f00 = field[y0m, x0m]
    f10 = field[y0m, x1m]
    f01 = field[y1m, x0m]
    f11 = field[y1m, x1m]
    return (f00 * (1 - fx) * (1 - fy) + f10 * fx * (1 - fy)
            + f01 * (1 - fx) * fy + f11 * fx * fy)


def _diffuse(field: np.ndarray, k: float) -> np.ndarray:
    """5-tap blur with wrap, matching canvas.frag getBlur (center weight k)."""
    nb = (np.roll(field, 1, 0) + np.roll(field, -1, 0)
          + np.roll(field, 1, 1) + np.roll(field, -1, 1))
    return (field * k + nb) / (4.0 + k)


def _hsv_image(field: np.ndarray) -> np.ndarray:
    """Render vector field (G, G, 2) as RGB: hue=direction, value=magnitude."""
    G = field.shape[0]
    ang = (np.arctan2(field[..., 1], field[..., 0]) / (2 * np.pi)) % 1.0
    mag = np.sqrt((field ** 2).sum(-1))
    val = np.clip(mag / (np.percentile(mag, 99) + 1e-9), 0, 1)
    sat = np.full_like(val, 0.85)
    rgb = np.zeros((G, G, 3))
    # vectorized HSV->RGB
    h6 = ang * 6.0
    i = np.floor(h6).astype(int) % 6
    f = h6 - np.floor(h6)
    p = val * (1 - sat)
    q = val * (1 - sat * f)
    t = val * (1 - sat * (1 - f))
    for idx, (r, g, b) in enumerate([(val, t, p), (q, val, p), (p, val, t),
                                     (p, q, val), (t, p, val), (val, p, q)]):
        m = i == idx
        rgb[m] = np.stack([r[m], g[m], b[m]], axis=-1)
    return (rgb * 255).astype(np.uint8)


DEFAULTS = dict(
    sensor_dist=0.03, sensor_angle=0.25, drag=0.85,
    axial=1.0, lateral=1.0, force_mult=0.012, strafe_power=0.006,
    sensor_gain=6.0, persistence=0.94, diffusion=2.0,
)


def _y_reflect(v):  # mirror across the body's forward axis: (a, b) -> (a, -b)
    out = v.copy()
    out[..., 1] *= -1
    return out


class Sim:
    """Steppable Fluoddity-lite. Hold state, advance one frame at a time, read
    out the current field as an image. Params live in self.params and may be
    changed between steps (the interactive viewer does exactly that)."""

    def __init__(self, rule, *, seed=0, n_particles=20000, grid=160, **params):
        self.rule = rule
        self.grid = grid
        self.n_particles = n_particles
        self.params = {**DEFAULTS, **params}
        self.reset(seed)

    def reset(self, seed=None):
        if seed is not None:
            self.seed = seed
        rng = np.random.default_rng(self.seed)
        self.pos = rng.uniform(-1, 1, size=(self.n_particles, 2))
        self.vel = rng.normal(0, 1e-3, size=(self.n_particles, 2))
        self.field = np.zeros((self.grid, self.grid, 2))
        self.frame = 0

    def set_rule(self, rule, seed=None):
        self.rule = rule
        self.reset(seed)

    def step(self):
        p = self.params
        grid = self.grid
        pos, vel, field = self.pos, self.vel, self.field

        speed = np.linalg.norm(vel, axis=1, keepdims=True)
        forward = np.where(speed > 1e-12, vel / np.maximum(speed, 1e-12), 0.0)
        left = np.stack([forward[:, 1], -forward[:, 0]], axis=1)

        # sensor offsets: rotate forward by +/- angle*pi, scale by distance
        a = p["sensor_angle"] * np.pi
        ca, sa = np.cos(a), np.sin(a)
        rotL = np.stack([forward[:, 0] * ca - forward[:, 1] * sa,
                         forward[:, 0] * sa + forward[:, 1] * ca], axis=1)
        rotR = np.stack([forward[:, 0] * ca + forward[:, 1] * sa,
                         -forward[:, 0] * sa + forward[:, 1] * ca], axis=1)
        pL = pos + rotL * p["sensor_dist"]
        pR = pos + rotR * p["sensor_dist"]

        def sample(pt):
            gx = (pt[:, 0] * 0.5 + 0.5) * grid
            gy = (pt[:, 1] * 0.5 + 0.5) * grid
            return _bilinear_wrap(field, gx, gy)
        g = p["sensor_gain"]
        Lw, Rw = sample(pL) * g, sample(pR) * g
        L = np.stack([(Lw * forward).sum(1), (Lw * left).sum(1)], axis=1)
        R = np.stack([(Rw * forward).sum(1), (Rw * left).sum(1)], axis=1)

        # symmetrized rule eval (cf. calculate_entity_behavior)
        inp = np.concatenate([L, R], axis=1)
        base = eval_rule(self.rule, inp)
        mirror_inp = np.concatenate([_y_reflect(R), _y_reflect(L)], axis=1)
        mirror = eval_rule(self.rule, mirror_inp)
        force_local = base[:, :2] + _y_reflect(mirror[:, :2])
        strafe_local = base[:, 2:] + _y_reflect(mirror[:, 2:])

        force = (forward * (force_local[:, :1] * p["axial"])
                 + left * (force_local[:, 1:] * p["lateral"])) * p["force_mult"]
        strafe = (forward * (strafe_local[:, :1] * p["axial"])
                  + left * (strafe_local[:, 1:] * p["lateral"])) * p["strafe_power"]

        vel = vel * p["drag"] + force
        pos = pos + vel + strafe
        pos = (pos + 1.0) % 2.0 - 1.0   # wrap to [-1, 1]

        gx = ((pos[:, 0] * 0.5 + 0.5) * grid).astype(np.int64) % grid
        gy = ((pos[:, 1] * 0.5 + 0.5) * grid).astype(np.int64) % grid
        deposit = np.zeros_like(field)
        np.add.at(deposit, (gy, gx), vel)

        self.field = (_diffuse(field, p["diffusion"]) * p["persistence"]
                      + (1 - p["persistence"]) * deposit * 60.0)
        self.pos, self.vel = pos, vel
        self.frame += 1

    def image(self) -> np.ndarray:
        return _hsv_image(self.field)


def render(rule, *, seed=0, n_particles=20000, grid=160, steps=140,
           return_field=False, **params):
    """Batch helper: run `steps` and return an RGB image. Used by evolve.py."""
    sim = Sim(rule, seed=seed, n_particles=n_particles, grid=grid, **params)
    for _ in range(steps):
        sim.step()
    if return_field:
        return sim.image(), sim.field
    return sim.image()


def viability(img: np.ndarray) -> float:
    """Structure score: variance that survives blurring. Coherent patterns
    (trails, bands, blobs) keep variance after a blur; pixel noise averages
    out. Near 0 = dead or noise; higher = organized phenotype. Used to filter
    duds and as an automatic selection signal."""
    v = img.mean(axis=2) / 255.0
    # 4x4 box downsample (blur) then measure remaining spatial variance
    G = (v.shape[0] // 4) * 4
    small = v[:G, :G].reshape(G // 4, 4, G // 4, 4).mean(axis=(1, 3))
    return float(small.var())


if __name__ == "__main__":
    from cppn import CPPN
    from fourier_rule import compile_cppn
    from PIL import Image

    net = CPPN(4, 4, rng=np.random.default_rng(3))
    for _ in range(25):
        net = net.mutate()
    rule, err = compile_cppn(net, seed=7)
    print(f"fit error {err:.3f}")
    img = render(rule, seed=0)
    print("viability:", round(viability(img), 5))
    Image.fromarray(img).save("smoke_phenotype.png")
    print("wrote smoke_phenotype.png")