#!/usr/bin/env python3
"""Push-Box point-flow overlay: GT CoTracker tracks vs 3-PoinTr predicted flow.

GREEN = actual moving-points (CoTracker on the demo, the GT tracking points).
RED   = 3-PoinTr predicted future trajectory, rolled out from frame-0's GT points.

The push-box 3-PoinTr (push_box_2d_150pts_384) was trained on ANISOTROPICALLY
squashed 384x384 tracks, so we run inference in 384-space (rescale frame-0 GT
1280x720 -> 384), then rescale the predicted flow back to native 1280x720 to draw
a crisp overlay on the original demo video.  horizon L=249 (point_flow_head=498=249*2).

Usage:
  python gen_pred_flow_pushbox.py --ep ep000 [--ep ep005 ...]   # or --all
"""
import argparse, os, sys
import cv2, imageio
import numpy as np
import torch

POINTR_REPO = "/home/hongyi/PWM/reward_model/3PoinTr_repo"
if POINTR_REPO not in sys.path:
    sys.path.insert(0, POINTR_REPO)
from util.training_utils.arch import TrajectoryModel
from util.geometry_utils.geom_utils import (
    compute_start_sphere_params, sim_transform, sim_transform_inverse_points,
)

CKPT = "/home/hongyi/PWM/reward_model/3pointr_models/push_box_2d_150pts_384_flow_backbone_final.pth"
TRACK_CACHE = "/data0/hongyi/reward_model_data/track_cache/demos_push_box"
PROMPT_SLUG = "cheez-it_box"
SQUASH = 384
MAX_POINTS = 150          # match the model's "150pts" training subsample


def infer_horizon(state):
    w = state.get("decoder_only.point_flow_head.weight")
    return None if w is None else int(w.shape[0] // 2)   # out = L*point_dim


def load_model(ckpt_path, device):
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    cfg = ckpt.get("args", {}) or {}
    state = ckpt.get("flow_backbone_state_dict") or ckpt.get("model_state_dict") or ckpt
    horizon = infer_horizon(state) or int(cfg.get("horizon", 167))
    model = TrajectoryModel(
        arch=cfg.get("arch", "flow_conditioned"),
        dim=int(cfg.get("arch_dim", 256)), depth=int(cfg.get("arch_depth", 2)),
        query_depth=int(cfg.get("arch_query_depth", 1)),
        num_heads=int(cfg.get("arch_heads", 4)),
        query_heads=int(cfg.get("arch_query_heads", 1)),
        use_task_conditioning=False, num_tasks=1, horizon=horizon,
        num_queries=int(cfg.get("num_queries", 10)), point_dim=2,
    ).to(device)
    missing, unexpected = model.load_state_dict(state, strict=False)
    print(f"[predict] horizon={horizon}  missing={len(missing)} unexpected={len(unexpected)}")
    model.eval()
    return model, horizon


@torch.no_grad()
def predict(model, horizon, x0_px, device):
    """x0_px (N,2) in 384-space -> (L,N,2) predicted positions in 384-space."""
    x0 = torch.from_numpy(x0_px.astype(np.float32)).to(device).unsqueeze(0)   # (1,N,2)
    scales, center = compute_start_sphere_params(x0)
    inv_scales = 1.0 / scales
    x0n = sim_transform(x0, scales=inv_scales, rotations=None, translations=-center)
    out = model({"x0": x0n, "num_steps": horizon,
                 "task_ids": torch.zeros(1, dtype=torch.long, device=device),
                 "scales": scales, "center": center}, mode="eval")
    pred = out["pred_pos"]                                  # (1, L, N, 2)
    L, N = pred.shape[1], pred.shape[2]
    un = sim_transform_inverse_points(pred.reshape(1, L * N, 2),
                                      scales=inv_scales, rotations=None, translations=-center)
    return un.view(L, N, 2).cpu().numpy()


def render(ep, model, horizon, device, out_dir):
    cache = f"{TRACK_CACHE}/tableHF_{ep}"
    cap = cv2.VideoCapture(f"{cache}/input_short.mp4")
    frames = []
    while True:
        ok, f = cap.read()
        if not ok:
            break
        frames.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB))
    cap.release()
    Td = len(frames)
    H, W = frames[0].shape[:2]
    z = np.load(f"{cache}/{PROMPT_SLUG}/tracks_filtered.npz")
    actual = z["tracks"].astype(np.float32)                 # (T,N,2) native
    vis = z["vis"] if "vis" in z.files else np.ones(actual.shape[:2], bool)
    N = actual.shape[1]
    if N > MAX_POINTS:                                       # even subsample to ~150
        idx = np.linspace(0, N - 1, MAX_POINTS).astype(int)
        actual, vis = actual[:, idx], vis[:, idx]
        N = actual.shape[1]

    sx, sy = SQUASH / W, SQUASH / H
    x0_384 = actual[0] * np.array([sx, sy], np.float32)     # frame-0 GT -> 384 space
    pred_384 = predict(model, horizon, x0_384, device)      # (L,N,2) 384
    pred = pred_384 / np.array([sx, sy], np.float32)        # back to native
    L = pred.shape[0]

    # resample predicted (1..L) onto demo timeline (align pred step 0 -> demo frame 0)
    tt = np.linspace(0, L - 1, Td, dtype=np.float32)
    t0 = np.floor(tt).astype(int).clip(0, L - 1)
    t1 = np.minimum(t0 + 1, L - 1)
    w = (tt - t0)[:, None, None]
    pred_al = (1 - w) * pred[t0] + w * pred[t1]             # (Td,N,2)

    out_frames = []
    for ti, f in enumerate(frames):
        cur = f.copy()
        for pi in range(N):
            if vis[ti, pi]:
                x, y = actual[ti, pi]
                cv2.circle(cur, (int(x), int(y)), 2, (0, 255, 0), -1)   # GT green
        for pi in range(N):
            x, y = pred_al[ti, pi]
            cv2.circle(cur, (int(x), int(y)), 2, (255, 0, 0), -1)       # pred red
        err = np.linalg.norm(actual[ti] - pred_al[ti], axis=-1)
        m = vis[ti].astype(np.float32)
        me = float((err * m).sum() / max(m.sum(), 1))
        cv2.putText(cur, f"t={ti}  mean px err={me:5.1f}", (12, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 0), 2)
        cv2.putText(cur, "GT (CoTracker)", (12, H - 40), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.putText(cur, "3PoinTr pred", (12, H - 14), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
        out_frames.append(cur)

    os.makedirs(out_dir, exist_ok=True)
    out = f"{out_dir}/pred_flow_{ep}.mp4"
    imageio.mimwrite(out, out_frames, fps=15, macro_block_size=1)
    err = np.linalg.norm(actual - pred_al, axis=-1) * vis
    denom = vis.sum(1).clip(min=1)
    me = (err.sum(1) / denom)
    print(f"[{ep}] N={N} L={L} Td={Td}  px err first={me[0]:.1f} mid={me[Td//2]:.1f} "
          f"last={me[-1]:.1f} mean={me.mean():.1f}  -> {out}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ep", action="append", default=[])
    ap.add_argument("--all", action="store_true")
    ap.add_argument("--ckpt", default=CKPT)
    ap.add_argument("--out_dir", default="/home/hongyi/PWM/external/pred_flow_pushbox")
    ap.add_argument("--device", default="cuda:0")
    a = ap.parse_args()
    eps = a.ep
    if a.all:
        import glob
        eps = sorted(os.path.basename(p).replace("tableHF_", "")
                     for p in glob.glob(f"{TRACK_CACHE}/tableHF_ep*"))
    if not eps:
        eps = ["ep000"]
    device = torch.device(a.device)
    model, horizon = load_model(a.ckpt, device)
    for ep in eps:
        try:
            render(ep, model, horizon, device, a.out_dir)
        except Exception as e:
            print(f"[{ep}] FAILED: {e}")


if __name__ == "__main__":
    main()