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

Java uploads the SUBAVG+LoG-conditioned stack ONCE; the DGX builds the temporal pyramid
(the 0.5*(now+prev) averaging, exactly as Java's temporalAverageLReLU), then serves
full-resolution dense inference per (level, scene): shift-and-stitch -> decode on GPU ->
return only the FULL-FRAME offset {dx,dy,s,Vx,Vy} + the ROI-only 121-cell softmax*s field.
A debug READBACK returns one conditioned pyramid frame for inspection. 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 ; float64 build_ms
  INFER   (2): int32 level, newest, roi_x, roi_y, roi_w, roi_h
              -> reply: float64 gpu_ms ; int32 H,W ; 5*H*W float32 (offset5: dx,dy,s,Vx,Vy, plane-major)
                        int32 rh,rw,nvel ; rh*rw*nvel float32 (ROI softmax*s, pixel-major)
  READBACK(3): int32 level, frame  -> reply: int32 H,W ; H*W float32 (one conditioned frame)
  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


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.java:824-883, temporalAverageLReLU linear).
    # log: [T,H,W] LoG-conditioned frames.
    #   level0[i]   = 0.5*(log[i+1] + log[i])                 i=0..T-2          (slice-as-shifted-view add)
    #   level L+1[i]= 0.5*(levelL[2i+2] + levelL[2i])         i=0..len(L)//2-2
    levels = [0.5 * (log[1:] + log[:-1])]                     # [T-1, H, W]
    while len(levels) < n_levels_max:
        prev = levels[-1]                                     # [Tl, H, W]
        nl = len(prev) // 2 - 1
        if nl < 1:
            break
        idx = torch.arange(nl, device=log.device)            # [nl]
        levels.append(0.5 * (prev[2 * idx + 2] + prev[2 * idx]))  # [nl, H, W]
    return levels


@torch.no_grad()
def shift_stitch(m, x, P, S=4):
    # FULL-RES field via shift-and-stitch (validated == per-pixel in fp64, dense_check.py).
    # x: [N,H,W]. Pad by half so dense cell (oi,oj) aligns to input pixel (S*oi,S*oj); the
    # S*S shifts interleave into [C,H,W].
    N, H, W = x.shape
    half = P // 2
    xp = F.pad(x, (half, half, half, half))                   # [N, H+P, W+P]
    full = torch.zeros(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:][None])[0]                   # [C, oH, oW]
            oh = min(y.shape[1], (H - sy + S - 1) // S)
            ow = min(y.shape[2], (W - sx + S - 1) // S)
            full[:, sy::S, sx::S] = y[:, :oh, :ow]
    return full                                               # [C, H, W]


@torch.no_grad()
def decode(field, vr, roi):
    # field: [C,H,W] full-res. C = 1(det) + nvel + 2(dx,dy). Decode on GPU.
    vdim = 2 * vr + 1
    nvel = vdim * vdim
    C, H, W = field.shape
    s = field[0].sigmoid()                                    # [H,W]   detection confidence
    p = field[1:1 + nvel].softmax(0)                          # [nvel,H,W] velocity softmax
    # cell k -> (vx,vy) = (k%vdim - vr, k//vdim - vr); centroid = sum_k p_k * coord_k
    k = torch.arange(nvel, device=field.device)
    cellx = (k % vdim - vr).to(field.dtype)                   # [nvel]
    celly = (k // vdim - vr).to(field.dtype)
    vx = (p * cellx[:, None, None]).sum(0)                    # [H,W]  velocity centroid (cells)
    vy = (p * celly[:, None, None]).sum(0)
    dx, dy = field[1 + nvel], field[1 + nvel + 1]            # [H,W]
    offset5 = torch.stack([dx, dy, s, vx, vy])               # [5,H,W]  -> -OFFSET (full frame)
    x0, y0, rw, rh = roi
    roi_field = (p * s[None])[:, y0:y0 + rh, x0:x0 + rw]     # [nvel,rh,rw]  softmax*s over ROI
    return offset5, roi_field.permute(1, 2, 0).contiguous(), nvel  # offset5 [5,H,W], roi [rh,rw,nvel]


def serve(run_dir, host, port):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.backends.cudnn.benchmark = True
    m, N, P, vr = load_model(run_dir, device)
    print(f"device={device} gpu={torch.cuda.get_device_name(0) if device=='cuda' else 'cpu'} "
          f"patch={P} N={N} vr={vr}", flush=True)
    _ = shift_stitch(m, torch.zeros(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}, full-res shift-and-stitch, stateful)", flush=True)
    pyr = None                                               # uploaded pyramid (per connection)
    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)
        try:
            while True:
                cmd = struct.unpack(">i", recvall(conn, 4))[0]
                if cmd == CMD_BYE:
                    break
                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))   # +N so the Java client can size scene windows
                elif cmd == CMD_INFER:
                    level, newest, rx, ry, rw, rh = struct.unpack(">iiiiii", recvall(conn, 24))
                    lev = pyr[level]                          # [Tl,H,W]
                    # newest-FIRST window (channel 0 = newest), matching Java inferROI window[h]=frames[newest-h]
                    # and the training order; .flip(0) reverses the ascending slice. By Claude on 06/20/2026
                    win = lev[newest - N + 1:newest + 1].flip(0)  # [N,H,W]
                    if device == "cuda":
                        ev0 = torch.cuda.Event(enable_timing=True); ev1 = torch.cuda.Event(enable_timing=True)
                        torch.cuda.synchronize(); ev0.record()
                    field = shift_stitch(m, win, P)           # [C,H,W]
                    offset5, roi_field, nvel = decode(field, vr, (rx, ry, rw, rh))
                    if device == "cuda":
                        ev1.record(); torch.cuda.synchronize(); gms = ev0.elapsed_time(ev1)
                    else:
                        gms = 0.0
                    H, W = field.shape[1], field.shape[2]
                    o5 = offset5.cpu().numpy().astype(">f4")          # [5,H,W]
                    rf = roi_field.cpu().numpy().astype(">f4")        # [rh,rw,nvel]
                    print(f"{datetime.now():%H:%M:%S} INFER lev={level} f={newest} ROI={rw}x{rh} "
                          f"gpu_exec={gms:.1f}ms ({(o5.nbytes+rf.nbytes)/1e6:.1f}MB out)", flush=True)
                    conn.sendall(struct.pack(">dii", gms, H, W) + o5.tobytes()
                                 + struct.pack(">iii", rh, rw, nvel) + rf.tobytes())
                elif cmd == CMD_READBACK:
                    level, frame = struct.unpack(">ii", recvall(conn, 8))
                    fr = pyr[level][frame].cpu().numpy().astype(">f4")  # [H,W]
                    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("--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)
