#!/usr/bin/env python3
"""Launch the resident Fast-WAM runtime and run the action-expert denoise loop on the NPU
WITHOUT reload, verify parity vs the PyTorch golden, and measure latency.

Fast-WAM's on-device unit here is the MoT **action expert** (`action_step`), which cross-
attends to the video KV cache produced by the video world-model DiT (`video_prefill`).
The video prefill runs host-side for now (its context binary needs a MolmoAct2-style
layer-split, see README), so the host:
  1. loads the video KV cache + context (here: from golden/, i.e. the video_prefill output),
  2. runs the flow-matching Euler loop: for each of num_inference_steps, feed the current
     latents_action + timestep -> pred_noise (NPU), then update latents with the Wan
     continuous flow-match step  latents += pred * delta,
  3. verifies the single-step output vs golden (parity gate), reports latency.

Pipeline (on-device portion): action_step ×N (resident). No per-inference reload.
"""
import sys, os, socket, struct, subprocess, time
import numpy as np

WS = os.environ.get("FASTWAM_WS", "/root/fastwam_workspace")
RT = f"{WS}/runtime"
ROLES = ["action_step"]
PORT0 = 5801
STEPS = int(os.environ.get("FASTWAM_STEPS", "10"))          # num_inference_steps
SHIFT = float(os.environ.get("FASTWAM_SHIFT", "5.0"))        # action_scheduler infer_shift
NUM_TRAIN = int(os.environ.get("FASTWAM_NUM_TRAIN_TIMESTEPS", "1000"))


def recvall(s, n):
    b = bytearray()
    while len(b) < n:
        c = s.recv(min(n - len(b), 8 << 20))
        if not c:
            raise ConnectionError()
        b += c
    return bytes(b)


def recv_tensors(s):
    n = struct.unpack("<I", recvall(s, 4))[0]
    out = {}
    for _ in range(n):
        nl = struct.unpack("<I", recvall(s, 4))[0]; name = recvall(s, nl).decode()
        dl = struct.unpack("<I", recvall(s, 4))[0]; dt = recvall(s, dl).decode()
        nd = struct.unpack("<I", recvall(s, 4))[0]; shape = struct.unpack(f"<{nd}I", recvall(s, nd * 4))
        nb = struct.unpack("<Q", recvall(s, 8))[0]; data = recvall(s, nb)
        out[name] = np.frombuffer(data, dtype=dt).reshape(shape).copy()
    return out


def send_tensors(s, d):
    parts = [struct.pack("<I", len(d))]
    for name, a in d.items():
        a = np.ascontiguousarray(a, np.float32)
        nb = name.encode(); dt = str(a.dtype).encode()
        parts.append(struct.pack("<I", len(nb)) + nb)
        parts.append(struct.pack("<I", len(dt)) + dt)
        parts.append(struct.pack("<I", a.ndim) + struct.pack(f"<{a.ndim}I", *a.shape))
        parts.append(struct.pack("<Q", a.nbytes) + a.tobytes())
    s.sendall(b"".join(parts))


def call(sock, d):
    send_tensors(sock, d)
    return recv_tensors(sock)


def cos(r, g):
    r = np.asarray(r, np.float64).ravel(); g = np.asarray(g, np.float64).ravel()
    return float(r @ g / (np.linalg.norm(r) * np.linalg.norm(g) + 1e-12))


def build_schedule(steps, shift, num_train):
    """Wan continuous flow-match inference schedule (vendored from video_dit.py):
    sigma = linspace(1,0,steps+1)[:steps]; sigma = shift*sigma/(1+(shift-1)*sigma).
    timesteps = sigma*num_train; deltas = sigma_next - sigma (sigma_next[-1]=0)."""
    sigma = np.linspace(1, 0, steps + 1)[:steps]
    sigma = shift * sigma / (1 + (shift - 1) * sigma)
    timesteps = (sigma * num_train).astype(np.float32)
    sigma_next = np.concatenate([sigma[1:], np.zeros(1)])
    deltas = (sigma_next - sigma).astype(np.float32)
    return timesteps, deltas


def main():
    os.makedirs(f"{WS}/tmp", exist_ok=True)
    for r in ROLES:
        try: os.remove(f"{WS}/tmp/w_{r}.ready")
        except OSError: pass
    env = os.environ.copy()
    env["LD_LIBRARY_PATH"] = f"{WS}/qnn_libs:" + env.get("LD_LIBRARY_PATH", "")
    env["ADSP_LIBRARY_PATH"] = f"{WS}/dsp_libs"

    procs, ports = {}, {}
    t0 = time.time()
    for i, r in enumerate(ROLES):
        ports[r] = PORT0 + i
        logf = open(f"{WS}/tmp/w_{r}.log", "w")
        procs[r] = subprocess.Popen(["python3", "-u", f"{RT}/resident_worker.py", r, str(ports[r])],
                                    env=env, stdout=logf, stderr=subprocess.STDOUT)
    print(f"[launch] {len(ROLES)} worker spawning ...", flush=True)
    for r in ROLES:
        for _ in range(240):
            if os.path.exists(f"{WS}/tmp/w_{r}.ready"):
                break
            if procs[r].poll() is not None:
                print(f"[ERR] worker {r} died during load (rc={procs[r].returncode}); see tmp/w_{r}.log")
                return 3
            time.sleep(0.5)
        else:
            print(f"[ERR] worker {r} not ready"); return 3
    print(f"[launch] worker RESIDENT in {time.time()-t0:.1f}s (1 NPU session)", flush=True)

    socks = {}
    for r in ROLES:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect(("127.0.0.1", ports[r])); s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        socks[r] = s

    g = np.load(f"{WS}/golden/action_step_io.npz")
    vk = g["video_keys"].astype(np.float32)          # [30,1,98,3072] (video world-model KV cache)
    vv = g["video_values"].astype(np.float32)
    ctx = g["context"].astype(np.float32)            # [1,129,4096]
    cmask = g["context_mask"].astype(np.float32)     # [1,129]
    la0 = g["latents_action"].astype(np.float32)     # [1,32,7] initial noise
    ts0 = g["timestep"].astype(np.float32)           # [1] step-0 timestep
    pred_gold = g["pred"].astype(np.float32)         # [1,32,7] one-step reference

    const = {"video_keys": vk, "video_values": vv, "context": ctx, "context_mask": cmask}

    # ---- parity: single step with golden inputs vs golden pred ----
    out1 = call(socks["action_step"], {**const, "latents_action": la0, "timestep": ts0})
    key = "pred_noise" if "pred_noise" in out1 else list(out1)[0]
    c = cos(pred_gold, out1[key])
    print(f"[verify] action_step single-step pred cos={c:.8f}  (PASS={c>=0.999})", flush=True)

    # ---- full N-step flow-match loop (latency) ----
    timesteps, deltas = build_schedule(STEPS, SHIFT, NUM_TRAIN)

    def run_loop(timing=None):
        la = la0.copy()
        for i in range(STEPS):
            ts = np.array([timesteps[i]], np.float32)
            s = time.time()
            o = call(socks["action_step"], {**const, "latents_action": la, "timestep": ts})
            if timing is not None:
                timing["action_step"] = timing.get("action_step", 0) + time.time() - s
            pred = np.ascontiguousarray(o[key], np.float32).reshape(la.shape)
            la = la + pred * deltas[i]     # Wan continuous flow-match step: x += pred*delta
        return la

    run_loop()  # warmup
    N = 5
    timing = {}
    t0 = time.time()
    for _ in range(N):
        run_loop(timing)
    total = (time.time() - t0) / N
    print(f"\n=== RESIDENT LATENCY (avg of {N}, 1 session, NO reload, {STEPS}-step flow-match) ===", flush=True)
    print(f"  action_step (x{STEPS})  {timing['action_step']/N*1000:7.1f} ms  (incl. video-KV TCP each step)", flush=True)
    print(f"  {'per-step (wall)':18}  {timing['action_step']/N/STEPS*1000:7.1f} ms", flush=True)
    print(f"  {'TOTAL':18}  {total*1000:7.1f} ms", flush=True)
    print("  (pure NPU ctx.Inference per step -> tmp/w_action_step.log [infer-ms])", flush=True)

    for s in socks.values():
        try: s.close()
        except OSError: pass
    time.sleep(1)
    for p in procs.values():
        try: p.terminate()
        except OSError: pass
    return 0


if __name__ == "__main__":
    sys.exit(main())