"""C5P Layer-2 (track-before-detect) — minimal circular-ConvGRU on a torus. By Claude on 06/21/2026

Layer 1 (frozen RawFCN) emits, per level-frame, a dense stride-4 field {s, Vx, Vy, dx, dy}.
Layer 2 is a RECURRENT net whose hidden state is the running 4D track memory (x, y, vx, vy),
fed a target-following 32x32 slice of that field one frame at a time. This first cut is the
SIMPLEST viable version (per Andrey 06/21):
  - plain circular ConvGRU  (NO explicit velocity-advection warp yet — added as a 2nd step;
    the conv recurrence still learns local motion implicitly),
  - dense Gaussian-bump readout (det map + Vx,Vy maps; supervise with a bump at truth),
  - single target, free-orbit (absolute position = torus-local + winding offset, tracked
    OUTSIDE the net; not needed for this module's forward/backward).

Torus rationale: xy is a PERIODIC 32x32 grid (Conv2d padding_mode='circular'). With the target
drift over a window staying << 32 cells, the single target "lives in infinite space" on a tiny
fixed array — no border code, translation-equivariant everywhere, trivial to batch. vx,vy are
NOT periodic (bounded by vmax; velocity does not wrap).

UNITS: the field grid is stride-4, so one torus cell = 4 scene px. Vx,Vy channels and the
velocity readout are kept in Layer-1 units (px/level-frame); vmax≈1.4 px/frame => ~0.35 cells/
frame => ~2.8 cells over N=8 (<< 32, the R<<G condition the torus relies on). The /4 conversion
to cells only matters once we add the advection warp.

Run the smoke test:  python layer2.py
"""

import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# Recurrent cell
# ---------------------------------------------------------------------------
class ConvGRUCellTorus(nn.Module):
    """One ConvGRU step with circular (toroidal) padding on the xy grid. By Claude on 06/21/2026

    Standard ConvGRU:
        z = sigmoid(Wz . [x, h])            update gate     [B, Ch, G, G]
        r = sigmoid(Wr . [x, h])            reset gate      [B, Ch, G, G]
        n = tanh   (Wn . [x, r*h])          candidate state [B, Ch, G, G]
        h'= (1 - z) * h + z * n             new hidden      [B, Ch, G, G]
    All convs are k x k with padding_mode='circular' so the 32x32 grid wraps both axes.
    """
    def __init__(self, ch_in, ch_hidden, k=3):
        super().__init__()
        pad = k // 2
        cat = ch_in + ch_hidden                                  # concat of input + hidden along channels
        # one conv per gate; circular pad makes the receptive field wrap the torus edges
        self.conv_z = nn.Conv2d(cat, ch_hidden, k, padding=pad, padding_mode='circular')
        self.conv_r = nn.Conv2d(cat, ch_hidden, k, padding=pad, padding_mode='circular')
        self.conv_n = nn.Conv2d(cat, ch_hidden, k, padding=pad, padding_mode='circular')

    def forward(self, x, h):
        # x: [B, Cin, G, G]   h: [B, Ch, G, G]  ->  h_new: [B, Ch, G, G]
        xh = torch.cat([x, h], dim=1)                            # [B, Cin+Ch, G, G]
        z = torch.sigmoid(self.conv_z(xh))                       # [B, Ch, G, G] update gate
        r = torch.sigmoid(self.conv_r(xh))                       # [B, Ch, G, G] reset gate
        xrh = torch.cat([x, r * h], dim=1)                       # [B, Cin+Ch, G, G] reset-masked hidden
        n = torch.tanh(self.conv_n(xrh))                         # [B, Ch, G, G] candidate
        return (1.0 - z) * h + z * n                             # [B, Ch, G, G] new hidden


# ---------------------------------------------------------------------------
# Layer-2 net
# ---------------------------------------------------------------------------
class Layer2Net(nn.Module):
    """Recurrent track-before-detect over a torus field sequence. By Claude on 06/21/2026

    forward(seq) consumes T frames of the Layer-1 field slice and returns, per frame, a dense
    det logit + (Vx,Vy) over the torus. Hidden state starts at 0 (no track) and accumulates
    evidence across frames — the recurrence IS the track filter.
    """
    def __init__(self, ch_in=3, ch_hidden=24, grid=32, vmax=1.4, k=3):
        super().__init__()
        self.ch_in = ch_in           # field channels fed in: s, Vx, Vy
        self.ch_hidden = ch_hidden   # hidden track-memory channels
        self.grid = grid             # torus side G (cells); one cell = 4 scene px
        self.vmax = vmax             # velocity readout bound, px/level-frame (matches Layer-1 training vmax)
        self.cell = ConvGRUCellTorus(ch_in, ch_hidden, k=k)
        # readout head: hidden -> det(1) + raw Vx,Vy(2); 1x1 conv = per-cell decode
        self.head = nn.Conv2d(ch_hidden, 1 + 2, 1)

    def init_hidden(self, B, device, dtype):
        # zero hidden = "no track yet"; [B, Ch, G, G]
        return torch.zeros(B, self.ch_hidden, self.grid, self.grid, device=device, dtype=dtype)

    def decode(self, h):
        # h: [B, Ch, G, G] -> det_logit [B, 1, G, G], vel [B, 2, G, G] bounded to +-vmax
        o = self.head(h)                                         # [B, 3, G, G]
        det = o[:, 0:1]                                          # [B, 1, G, G] raw logit
        vel = self.vmax * torch.tanh(o[:, 1:3])                  # [B, 2, G, G] px/level-frame
        return det, vel

    def forward(self, seq, h=None):
        # seq: [B, T, Cin, G, G]  ->  det [B, T, 1, G, G], vel [B, T, 2, G, G]
        B, T = seq.shape[0], seq.shape[1]
        if h is None:
            h = self.init_hidden(B, seq.device, seq.dtype)
        dets, vels = [], []
        for t in range(T):                                       # BPTT unrolls this loop
            h = self.cell(seq[:, t], h)                          # [B, Ch, G, G] recurrent update
            det, vel = self.decode(h)                            # per-frame readout
            dets.append(det)
            vels.append(vel)
        det = torch.stack(dets, dim=1)                           # [B, T, 1, G, G]
        vel = torch.stack(vels, dim=1)                           # [B, T, 2, G, G]
        return det, vel


# ---------------------------------------------------------------------------
# Dense Gaussian-bump supervision (single target)
# ---------------------------------------------------------------------------
def bump_target(pos_xy, grid, sigma=1.0, device="cpu"):
    """Toroidal Gaussian bump at (sub-cell) position pos_xy. By Claude on 06/21/2026
    pos_xy: [B, T, 2] (x, y) in torus cells (may be fractional / out of [0,G) — wraps).
    Returns det bump [B, T, 1, G, G] in [0,1]. Distance uses the WRAPPED (toroidal) metric so
    a target near the edge still gets a single round bump that straddles the seam.
    """
    B, T = pos_xy.shape[0], pos_xy.shape[1]
    coord = torch.arange(grid, device=device).float()           # [G]
    gy, gx = torch.meshgrid(coord, coord, indexing='ij')        # [G, G] each
    gx = gx[None, None]; gy = gy[None, None]                     # [1,1,G,G] broadcast over B,T
    px = pos_xy[..., 0][..., None, None]                         # [B, T, 1, 1]
    py = pos_xy[..., 1][..., None, None]                         # [B, T, 1, 1]
    # wrapped (toroidal) coordinate difference: nearest image around the G-periodic grid
    dx = (gx - px + grid / 2) % grid - grid / 2                  # [B, T, G, G] in (-G/2, G/2]
    dy = (gy - py + grid / 2) % grid - grid / 2
    g = torch.exp(-(dx * dx + dy * dy) / (2 * sigma * sigma))    # [B, T, G, G]
    return g[:, :, None]                                         # [B, T, 1, G, G]


def layer2_loss(det_logit, vel, det_t, vel_t, support=0.3, pos_weight=20.0):
    """Detection BCE (sparse bump -> pos_weight) + velocity MSE on the bump support. By Claude 06/21
    det_logit: [B,T,1,G,G] raw   det_t: [B,T,1,G,G] in [0,1]
    vel:       [B,T,2,G,G]       vel_t: [B,T,2,G,G]  (px/level-frame; only used where det_t>support)
    """
    pw = torch.tensor(pos_weight, device=det_logit.device)
    l_det = F.binary_cross_entropy_with_logits(det_logit, det_t, pos_weight=pw)
    m = (det_t > support)                                        # [B,T,1,G,G] bump core mask
    if m.any():
        m2 = m.expand_as(vel)                                    # [B,T,2,G,G]
        l_vel = F.mse_loss(vel[m2], vel_t[m2])
    else:
        l_vel = vel.sum() * 0.0
    return l_det + 0.3 * l_vel, {"det": float(l_det.detach()), "vel": float(l_vel.detach() if torch.is_tensor(l_vel) else l_vel)}


# ---------------------------------------------------------------------------
# Smoke test: fake Layer-1-like field, single target on a wrapping straight line.
# Verifies the module trains end-to-end (forward + BPTT + loss) BEFORE real Layer-1 fields.
# This is NOT the real training data — that comes in the next step (trajectory-sequence gen).
# ---------------------------------------------------------------------------
def fake_field_batch(rng, B, T, grid, vmax, sigma=1.0, snr=4.0, device="cpu"):
    """Build a toy 'Layer-1 field' sequence + truth. By Claude on 06/21/2026
    A single target starts at a random torus cell, moves at constant (vx,vy) px/frame
    (=> (vx,vy)/4 cells/frame), wrapping. The s-channel is a noisy Gaussian bump at the target;
    Vx,Vy channels carry the true velocity over the bump (+ noise), 0 elsewhere. Returns:
      seq    [B,T,3,G,G]  (s, Vx, Vy)
      pos    [B,T,2]      target (x,y) in cells
      veltru [B,T,2]      true (Vx,Vy) px/level-frame
    """
    seq = torch.zeros(B, T, 3, grid, grid, device=device)
    pos = torch.zeros(B, T, 2, device=device)
    veltru = torch.zeros(B, T, 2, device=device)
    for b in range(B):
        x0 = rng.uniform(0, grid); y0 = rng.uniform(0, grid)
        ang = rng.uniform(0, 2 * np.pi); spd = rng.uniform(0.3, 1.0) * vmax
        vx = spd * np.cos(ang); vy = spd * np.sin(ang)          # px/level-frame
        for t in range(T):
            cx = (x0 + vx / 4.0 * t)                             # cells (stride-4 => /4)
            cy = (y0 + vy / 4.0 * t)
            pos[b, t, 0] = cx % grid; pos[b, t, 1] = cy % grid
            veltru[b, t, 0] = vx; veltru[b, t, 1] = vy
        # s channel: noisy toroidal bump at the target; vel channels: truth over the bump
        bump = bump_target(pos[b:b+1].unsqueeze(0).reshape(1, T, 2), grid, sigma, device)  # [1,T,1,G,G]
        bump = bump[0, :, 0]                                     # [T,G,G]
        noise = torch.from_numpy(rng.standard_normal((T, grid, grid)).astype(np.float32)).to(device)
        seq[b, :, 0] = (snr * bump + noise).clamp(min=0.0)       # s >= 0, SNR-scaled signal in noise
        core = (bump > 0.3).float()                              # [T,G,G]
        seq[b, :, 1] = vx * core; seq[b, :, 2] = vy * core
    return seq, pos, veltru


def smoke_test(steps=400, B=16, T=8, grid=32, vmax=1.4, device=None):
    """Overfit the toy generator a few hundred steps; det peak should sharpen, vel MSE drop."""
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    rng = np.random.default_rng(0)
    net = Layer2Net(ch_in=3, ch_hidden=24, grid=grid, vmax=vmax).to(device)
    opt = torch.optim.Adam(net.parameters(), 2e-3)
    nparams = sum(p.numel() for p in net.parameters())
    print(f"Layer2Net: {nparams} params, grid={grid}, ch_hidden=24, device={device}", flush=True)
    for step in range(1, steps + 1):
        seq, pos, veltru = fake_field_batch(rng, B, T, grid, vmax, device=device)
        det_t = bump_target(pos, grid, sigma=1.0, device=device)         # [B,T,1,G,G]
        vel_t = torch.zeros(B, T, 2, grid, grid, device=device)
        core = (det_t[:, :, 0] > 0.3)                                     # [B,T,G,G]
        for c in range(2):
            vel_t[:, :, c][core] = veltru[..., c][..., None, None].expand(-1, -1, grid, grid)[core]
        det_logit, vel = net(seq)
        loss, comp = layer2_loss(det_logit, vel, det_t, vel_t)
        opt.zero_grad(); loss.backward(); opt.step()
        if step % 50 == 0 or step == 1:
            with torch.no_grad():
                p = torch.sigmoid(det_logit)
                peak = float(p[det_t > 0.3].mean())
                bg = float(p[det_t < 0.05].max())
            print(f"step {step:4d}  det {comp['det']:.4f}  vel {comp['vel']:.4f}  "
                  f"peak(s@truth) {peak:.3f}  max-bg {bg:.3f}", flush=True)
    print("smoke test done.", flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--steps", type=int, default=400)
    ap.add_argument("--grid", type=int, default=32)
    ap.add_argument("--vmax", type=float, default=1.4)
    a = ap.parse_args()
    smoke_test(steps=a.steps, grid=a.grid, vmax=a.vmax)
