#!/usr/bin/env python3
"""DGX remote-inference server for the CUAS RawFCN (stateful, batched). By Claude on 06/20/2026.

Java uploads the SUBAVG+LoG-conditioned (optionally synth-mixed) stack ONCE; the DGX builds the
temporal pyramid (0.5*(now+prev), as Java's temporalAverageLReLU), then serves BATCHED full-res
inference: one INFER does a whole range of scenes of a level (chunked for GPU memory), shift-and-
stitch -> on-GPU GHOSTBUSTER (== CuasDetectRT.dnnGhostbust) -> decode -> returns the full-frame
offset {dx,dy,s,Vx,Vy} + the ROI-only 121-cell softmax*s, per scene. A debug READBACK returns one
conditioned pyramid frame. Runs inside nvcr.io/nvidia/pytorch:25.10-py3.

Protocol (big-endian, matches Java DataInput/OutputStream). Each request: int32 cmd, then:
  UPLOAD  (1): int32 T,H,W ; T*H*W float32 (conditioned stack)
              -> reply: int32 n_levels ; n_levels x int32 frames_per_level ; int32 N ; float64 build_ms
  INFER   (2): int32 level, start, count, stride, roi_x, roi_y, roi_w, roi_h ; float64 rmax_cells
              (count scenes: newest = start + j*stride, j=0..count-1; rmax_cells<=0 disables ghostbuster)
              -> reply: float64 gpu_ms ; int32 H,W,count,nvel,rh,rw
                        count*5*H*W float32 (offset5 dx,dy,s,Vx,Vy, plane-major per scene)
                        count*rh*rw*nvel float32 (ROI softmax*s, pixel-major per scene)
  READBACK(3): int32 level, frame -> reply: int32 H,W ; H*W float32
  BYE     (0): close
"""
import argparse, os, socket, struct, time
from datetime import datetime
import numpy as np
import torch
import torch.nn.functional as F
from model import RawFCN

CMD_BYE, CMD_UPLOAD, CMD_INFER, CMD_READBACK = 0, 1, 2, 3
CMD_STATUS = 4     # report loaded L1/L2 model paths so the client can detect a model change. By Claude on 06/24/2026
GPU_CHUNK = 16     # scenes processed per batched GPU pass (memory vs utilization)
VEL_DECIMATE = 4   # velocity-grid cells per px/level-frame (Java curt_vel_decimate); L2 was trained on Vx,Vy in px/frame (cells/4). By Claude on 06/22/2026
AGE_THR = 0.2      # L2 track-age death threshold: a cell with det<=AGE_THR "dies" (age 0). Raised 0.01->0.2 so the
                   # weak noise halo dies and the 5x5 max-pool can't dilate age across gaps. By Claude on 06/24/2026
AGE_K   = 0.5      # ancestor gate: a 5x5 previous-frame neighbor may pass its age only if its det >= AGE_K * (local
                   # max det in that 5x5) - blocks a weak-but-old straggler from seeding age. By Claude on 06/24/2026
NOISE_REF_LEVEL = 3  # the net is calibrated to ~LEV3's absolute noise (low-contrast signals tested mainly on LEV3).
                   # The pyramid averages 2 frames/level so sigma drops sqrt(2)/level; scale each level's L1 input by
                   # sqrt(2)^(level-REF) to put every level at LEV3's absolute noise (uniform FP). By Claude on 06/24/2026


def load_l2(run_dir, device):
    # Optional Layer-2 (track-before-detect) recurrent net; FCN so it runs on any H,W. By Claude on 06/22/2026
    from layer2 import Layer2Net
    ck = torch.load(os.path.join(run_dir, "model.pt"), map_location="cpu", weights_only=False)
    a = ck.get("args", {}) or {}
    m = Layer2Net(ch_in=3, ch_hidden=a.get("ch", 24), grid=a.get("G", 32), vmax=a.get("vmax", 1.4))
    m.load_state_dict(ck["model"]); m.eval().to(device)
    print(f"loaded L2 {run_dir}/model.pt: ch_hidden={a.get('ch',24)} vmax={a.get('vmax',1.4)}", flush=True)
    return m


def load_model(run_dir, device):
    ck = torch.load(os.path.join(run_dir, "model.pt"), map_location="cpu", weights_only=False)
    a = ck.get("args", {}) or {}
    kw = dict(n_frames=a.get("nframes", 8), vel_radius=a.get("vel_radius", 5),
              patch=a.get("patch", 24), velocity_mode=a.get("velocity_mode", "grid"),
              vmax=a.get("vmax", 1.4))
    if a.get("ch") is not None:
        kw["ch"] = tuple(a["ch"])
    m = RawFCN(**kw)
    m.load_state_dict(ck["model"])
    m.eval().to(device)
    print(f"loaded {run_dir}/model.pt: N={kw['n_frames']} patch={kw['patch']} vr={kw['vel_radius']} "
          f"mode={kw['velocity_mode']} out_ch={m.out_ch}", flush=True)
    return m, kw["n_frames"], kw["patch"], kw["vel_radius"]


def recvall(conn, n):
    buf = bytearray()
    while len(buf) < n:
        chunk = conn.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("short read")
        buf += chunk
    return bytes(buf)


def build_pyramid(log, n_levels_max=8):
    # Replicates Java's pyramid (CuasDetectRT temporalAverageLReLU, linear). log: [T,H,W].
    levels = [0.5 * (log[1:] + log[:-1])]                 # [T-1,H,W]
    while len(levels) < n_levels_max:
        prev = levels[-1]
        nl = len(prev) // 2 - 1
        if nl < 1:
            break
        idx = torch.arange(nl, device=log.device)
        levels.append(0.5 * (prev[2 * idx + 2] + prev[2 * idx]))
    return levels


@torch.no_grad()
def shift_stitch(m, x, P, S=4):
    # x: [B,N,H,W] -> full-res field [B,C,H,W] (validated == per-pixel in fp64, dense_check.py).
    B, N, H, W = x.shape
    half = P // 2
    xp = F.pad(x, (half, half, half, half))               # [B,N,H+P,W+P]
    full = torch.zeros(B, m.out_ch, H, W, device=x.device, dtype=x.dtype)
    for sy in range(S):
        for sx in range(S):
            y = m(xp[:, :, sy:, sx:])                     # [B,C,oH,oW]
            oh = min(y.shape[2], (H - sy + S - 1) // S)
            ow = min(y.shape[3], (W - sx + S - 1) // S)
            full[:, :, sy::S, sx::S] = y[:, :, :oh, :ow]
    return full                                           # [B,C,H,W]


@torch.no_grad()
def decode(field, vr, roi, rmax_cells):
    # field: [B,C,H,W]. On-GPU ghostbuster (== CuasDetectRT.dnnGhostbust) + decode.
    vdim = 2 * vr + 1
    nvel = vdim * vdim
    B, C, H, W = field.shape
    x0, y0, rw, rh = roi
    s = field[:, 0].sigmoid()                             # [B,H,W]
    p = field[:, 1:1 + nvel].softmax(1)                   # [B,nvel,H,W]
    k = torch.arange(nvel, device=field.device)
    cx = (k % vdim - vr).to(field.dtype)                  # [nvel] vx cell coord
    cy = (k // vdim - vr).to(field.dtype)                 # [nvel] vy cell coord
    if rmax_cells > 0:                                    # ghostbuster
        corner = (cx * cx + cy * cy) > (rmax_cells * rmax_cells)   # [nvel] untrained corner cells
        ghost = corner[p.argmax(1)]                       # [B,H,W] peak lands in a corner -> whole pixel is a ghost
        p = p * (~corner).to(p.dtype)[None, :, None, None]   # zero corner cells everywhere
        keep = (~ghost).to(p.dtype)                       # [B,H,W]
        p = p * keep[:, None]                             # zero all cells at ghost pixels
        s = s * keep                                      # s=0 at ghost pixels
    psum = p.sum(1).clamp_min(1e-12)                      # [B,H,W] (normalize the centroid)
    vx = (p * cx[None, :, None, None]).sum(1) / psum      # [B,H,W] velocity centroid (cells)
    vy = (p * cy[None, :, None, None]).sum(1) / psum
    dx, dy = field[:, 1 + nvel], field[:, 1 + nvel + 1]   # [B,H,W]
    offset5 = torch.stack([dx, dy, s, vx, vy], 1)         # [B,5,H,W]  -> -OFFSET
    roi_field = (p * s[:, None])[:, :, y0:y0 + rh, x0:x0 + rw]   # [B,nvel,rh,rw] softmax*s over ROI
    return offset5, roi_field.permute(0, 2, 3, 1).contiguous(), nvel  # [B,5,H,W], [B,rh,rw,nvel]


def serve(run_dir, host, port, l2_run=None):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.backends.cudnn.benchmark = True
    torch.set_grad_enabled(False)   # inference-only server; L2 recurrence (m2.cell/decode) isn't @no_grad'd. By Claude 06/22/2026
    m, N, P, vr = load_model(run_dir, device)
    m2 = load_l2(l2_run, device) if l2_run else None    # optional Layer-2; None -> L1-only (old way). By Claude 06/22/2026
    print(f"device={device} gpu={torch.cuda.get_device_name(0) if device=='cuda' else 'cpu'} "
          f"patch={P} N={N} vr={vr} L2={'on('+l2_run+')' if m2 is not None else 'off'}", flush=True)
    _ = shift_stitch(m, torch.zeros(1, N, 64, 64, device=device), P)   # warm-up
    if device == "cuda":
        torch.cuda.synchronize()
    print("warm-up done", flush=True)

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((host, port))
    srv.listen(4)
    print(f"listening on {host}:{port} (N={N}, batched full-res shift-and-stitch + ghostbuster)", flush=True)
    pyr = None
    while True:
        conn, addr = srv.accept()
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        print(f"{datetime.now():%H:%M:%S} client {addr}", flush=True)
        h_l2 = None     # Layer-2 recurrent hidden state [1,ch,H,W]; persists across INFER chunks, reset on l2_reset. By Claude 06/22/2026
        age_l2 = None   # L2 track-age field [1,1,H,W]; sprev_l2 = previous-frame L2 det; carried+reset like h_l2. By Claude 06/24/2026
        sprev_l2 = None
        try:
            while True:
                cmd = struct.unpack(">i", recvall(conn, 4))[0]
                if cmd == CMD_BYE:
                    break
                if cmd == CMD_STATUS:
                    # Reply with loaded L1 + L2 model paths (len-prefixed UTF-8); empty L2 = L1-only.
                    # Lets the Java client detect a model change and relaunch. By Claude on 06/24/2026
                    b1 = run_dir.encode("utf-8")
                    b2 = (l2_run or "").encode("utf-8")
                    conn.sendall(struct.pack(">i", len(b1)) + b1 + struct.pack(">i", len(b2)) + b2)
                    continue
                if cmd == CMD_UPLOAD:
                    T, H, W = struct.unpack(">iii", recvall(conn, 12))
                    data = recvall(conn, T * H * W * 4)
                    log = torch.from_numpy(np.frombuffer(data, dtype=">f4").astype(np.float32)
                                           .reshape(T, H, W)).to(device)
                    t0 = time.perf_counter()
                    pyr = build_pyramid(log)
                    if device == "cuda":
                        torch.cuda.synchronize()
                    bms = (time.perf_counter() - t0) * 1e3
                    nl = len(pyr)
                    print(f"{datetime.now():%H:%M:%S} UPLOAD T={T} {H}x{W} -> {nl} levels "
                          f"{[len(l) for l in pyr]}  build={bms:.1f}ms ({T*H*W*4/1e6:.1f}MB)", flush=True)
                    conn.sendall(struct.pack(">i", nl) + b"".join(struct.pack(">i", len(l)) for l in pyr)
                                 + struct.pack(">id", N, bms))
                elif cmd == CMD_INFER:
                    level, start, count, stride, rx, ry, rw, rh = struct.unpack(">iiiiiiii", recvall(conn, 32))
                    rmax = struct.unpack(">d", recvall(conn, 8))[0]
                    l2_enable, l2_reset = struct.unpack(">ii", recvall(conn, 8))   # By Claude 06/22/2026
                    noise_scale = struct.unpack(">d", recvall(conn, 8))[0]         # per-level L1-input noise scale from Java (single source of truth); <=0 -> server fallback. By Claude 06/24/2026
                    use_l2 = bool(l2_enable) and (m2 is not None)
                    # Per-level noise normalization: scale this level's L1 input to LEV3's absolute noise so all
                    # levels sit in the net's trained regime (uniform FP across levels). LEV3 -> 1.0, lower/noisier
                    # levels scale down, higher levels up. Independent of the age filter. By Claude on 06/24/2026
                    if noise_scale <= 0.0:                # fallback only: Java didn't send one -> theoretical sqrt(2)^(level-ref)
                        noise_scale = 2.0 ** ((level - NOISE_REF_LEVEL) / 2.0)
                    lev = pyr[level] * noise_scale        # [Tl,H,W]
                    H, W = lev.shape[1], lev.shape[2]
                    nvel = (2 * vr + 1) ** 2
                    o5_gpu, rf_gpu = [], []
                    # Time PURE GPU compute (shift-and-stitch + decode), continuous over the whole range -
                    # the production throughput. Results stay on-GPU (prod feeds Layer 2 there); the D2H copy
                    # below is dev-only and NOT timed. By Claude on 06/20/2026
                    if device == "cuda":
                        ev0 = torch.cuda.Event(enable_timing=True); ev1 = torch.cuda.Event(enable_timing=True)
                        torch.cuda.synchronize(); ev0.record()
                    else:
                        t0 = time.perf_counter()
                    for c0 in range(0, count, GPU_CHUNK):
                        b = min(GPU_CHUNK, count - c0)
                        # newest-first windows (channel 0 = newest), matching the Java order
                        wins = torch.stack([lev[(start + (c0 + j) * stride) - N + 1:
                                                (start + (c0 + j) * stride) + 1].flip(0) for j in range(b)])  # [b,N,H,W]
                        field = shift_stitch(m, wins, P)  # [b,C,H,W]
                        o5, rf, nv = decode(field, vr, (rx, ry, rw, rh), rmax)   # L1: ghostbusted offset5 + ROI
                        nvel = nv
                        if use_l2:
                            # Layer-2 (track-before-detect) over the scene/time axis. Feed the FULL
                            # (non-ghostbusted) field as (s, Vx/vd, Vy/vd) px/level-frame; carry the recurrent
                            # hidden state across chunks (reset on l2_reset at the level's first chunk). Output
                            # replaces offset5 with {L1 dx, L1 dy, L2 det, L2 Vx*vd, L2 Vy*vd} (vel back to cells
                            # so Java's existing /vel_decimate viz scaling -> px/level-frame). By Claude 06/22/2026
                            ong, _, _ = decode(field, vr, (rx, ry, rw, rh), 0.0)        # no ghostbuster (L2 gets full field)
                            l2in = torch.stack([ong[:, 2], ong[:, 3] / VEL_DECIMATE, ong[:, 4] / VEL_DECIMATE], 1)  # [b,3,H,W]
                            # FPN-bad margins arrive as NaN; the recurrent circular conv would otherwise spread
                            # NaN inward by the kernel radius every frame ("eating" the borders). Sanitize the
                            # input so NaN can never seed/propagate through the hidden state. By Claude 06/22/2026
                            l2in = torch.nan_to_num(l2in, nan=0.0, posinf=0.0, neginf=0.0)
                            Hf, Wf = l2in.shape[2], l2in.shape[3]
                            if (h_l2 is None) or (h_l2.shape[2] != Hf) or (h_l2.shape[3] != Wf) or (l2_reset and c0 == 0):
                                h_l2 = torch.zeros(1, m2.ch_hidden, Hf, Wf, device=device, dtype=field.dtype)
                                age_l2 = torch.zeros(1, 1, Hf, Wf, device=device, dtype=field.dtype)   # track age, carried+reset like h_l2
                                sprev_l2 = torch.zeros(1, 1, Hf, Wf, device=device, dtype=field.dtype)  # previous-frame L2 det
                            dets, vxs, vys, ages = [], [], [], []
                            for j in range(b):                                          # forward in time, carry hidden + age
                                h_l2 = m2.cell(l2in[j:j+1], h_l2)
                                dlog, vel = m2.decode(h_l2)                             # [1,1,H,W],[1,2,H,W]
                                s = torch.sigmoid(dlog[:, 0:1])                         # [1,1,H,W] current L2 det
                                # AGE (track-before-detect persistence): die where det<=AGE_THR, else 1 + oldest age among
                                # 5x5 PREVIOUS-frame neighbors that are themselves STRONG (det >= AGE_K * local-max det) -
                                # so a weak-but-old straggler can't seed age; the raised AGE_THR stops the noise halo from
                                # dilating age across gaps. Level-uniform 5x5 (pyramid keeps ~const px/level-frame). By Claude 06/24/2026
                                maxS = F.max_pool2d(sprev_l2, 5, 1, 2)                          # local max prev-det in 5x5
                                elig = (sprev_l2 >= AGE_K * maxS) & (sprev_l2 > AGE_THR)        # strong AND alive ancestors
                                prev = torch.where(elig, age_l2, torch.zeros_like(age_l2))      # only strong ancestors pass age
                                age_l2 = torch.where(s > AGE_THR, F.max_pool2d(prev, 5, 1, 2) + 1.0, torch.zeros_like(age_l2))
                                sprev_l2 = s
                                dets.append(s[:, 0]); ages.append(age_l2[:, 0]); vxs.append(vel[:, 0]); vys.append(vel[:, 1])
                            l2vx = torch.cat(vxs, 0) * VEL_DECIMATE; l2vy = torch.cat(vys, 0) * VEL_DECIMATE
                            o5 = torch.stack([o5[:, 0], o5[:, 1], torch.cat(dets, 0), l2vx, l2vy, torch.cat(ages, 0)], 1)  # +L2 age (6th); keep L1 dx,dy
                        o5_gpu.append(o5); rf_gpu.append(rf)   # keep on GPU (rf = L1 ROI reference even when L2 on)
                    if device == "cuda":
                        ev1.record(); torch.cuda.synchronize(); gms = ev0.elapsed_time(ev1)
                    else:
                        gms = (time.perf_counter() - t0) * 1e3
                    allo = torch.cat(o5_gpu, 0).cpu().numpy().astype(">f4")  # [count,nch,H,W] nch=5 (L1) or 6 (L2:+age) D2H untimed
                    nch = allo.shape[1]                                      # channel count sent in the header (was hardcoded 5)
                    allr = torch.cat(rf_gpu, 0).cpu().numpy().astype(">f4")  # [count,rh,rw,nvel]
                    print(f"{datetime.now():%H:%M:%S} INFER lev={level} {count} scenes (f{start}..,stride {stride}) "
                          f"ROI={rw}x{rh} ghost={rmax:.1f} nscale={noise_scale:.3f} L2={'on' if use_l2 else 'off'}{'(reset)' if (use_l2 and l2_reset) else ''} "
                          f"gpu={gms:.1f}ms ({(allo.nbytes+allr.nbytes)/1e6:.1f}MB out)", flush=True)
                    conn.sendall(struct.pack(">diiiiiii", gms, H, W, count, nch, nvel, rh, rw))  # +nch (offset channels). By Claude 06/24/2026
                    conn.sendall(allo.tobytes())
                    conn.sendall(allr.tobytes())
                elif cmd == CMD_READBACK:
                    level, frame = struct.unpack(">ii", recvall(conn, 8))
                    fr = pyr[level][frame].cpu().numpy().astype(">f4")
                    print(f"{datetime.now():%H:%M:%S} READBACK lev={level} f={frame}", flush=True)
                    conn.sendall(struct.pack(">ii", fr.shape[0], fr.shape[1]) + fr.tobytes())
        except (ConnectionError, struct.error, IndexError) as ex:
            print(f"client {addr} closed/err: {ex}", flush=True)
        finally:
            conn.close()


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", default="runs/weighted9_pm_s")
    ap.add_argument("--l2run", default=None, help="optional Layer-2 run dir (model.pt); omit for L1-only")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=5577)
    args = ap.parse_args()
    serve(args.run, args.host, args.port, l2_run=args.l2run)
