"""Load trained 3-PoinTr backbone, predict trajectories on a val episode,
project predictions back to image, overlay GT and predicted points side-by-side."""
import os, sys, json
sys.path.insert(0, "/home/hongyi/PWM/reward_model/3PoinTr_repo")

import numpy as np
import torch
import imageio.v2 as imageio
import cv2

from util.training_utils.arch import TrajectoryModel
from util.training_utils.train_util import (
    build_common_arg_parser, prepare_datasets_and_loaders, normalize_and_augment,
    extract_diffusion_kwargs,
)


DATA_ROOT = "/home/hongyi/PWM/external/3pointr_data"
SPLIT_FILE = os.path.join(DATA_ROOT, "stack_can_flow.json")
CKPT = "/home/hongyi/PWM/reward_model/3pointr_models/stack_can_flowonly_512pts_flow_backbone.pth"
EP_DIR = "/home/hongyi/PWM/external/3pointr_data/no_actions/stack_can/0036"
OUT_VIDEO = "/home/hongyi/PWM/reward_model/3pointr_models/val_ep0036_pred_vs_gt.mp4"


def load_ep(ep_dir):
    meta = json.load(open(os.path.join(ep_dir, "metadata.json")))
    T = meta["T"]
    pos, vis, rgbs = [], [], []
    for k in range(1, T + 1):
        ts = os.path.join(ep_dir, f"timestep_{k}")
        pos.append(torch.load(os.path.join(ts, f"{k}_pos.pth")).numpy())
        vis.append(torch.load(os.path.join(ts, f"{k}_visibility.pth")).numpy())
        rgbs.append(imageio.imread(os.path.join(ts, f"{k}_rgb.png")))
    pos = np.stack(pos, 0); vis = np.stack(vis, 0)
    K = meta["intrinsics"]["camera"]
    return rgbs, pos, vis, (K["fx"], K["fy"], K["cx"], K["cy"]), meta


def project(p, fx, fy, cx, cy):
    z = np.maximum(p[..., 2], 1e-6)
    u = p[..., 0] / z * fx + cx
    v = p[..., 1] / z * fy + cy
    return np.stack([u, v], -1)


def main():
    parser = build_common_arg_parser("predict")
    args = parser.parse_args([
        "--dataset", DATA_ROOT + "/", "--split_file", SPLIT_FILE,
        "--subsample", "512", "--batch_size", "1",
        "--num_queries", "10",
        "--arch_dim", "256", "--arch_depth", "2", "--arch_heads", "4",
        "--arch_query_depth", "1", "--arch_query_heads", "1",
        "--flow_data_only", "--tasks", "stack_can",
    ])

    arts = prepare_datasets_and_loaders(args)
    traj_len = arts.train_set.traj_len
    num_tasks = len(arts.tasks)
    print(f"traj_len={traj_len}, num_tasks={num_tasks}")

    diff_kwargs = extract_diffusion_kwargs(args)
    diff_kwargs["n_action_steps"] = traj_len - 1
    model = TrajectoryModel(
        dim=args.arch_dim, depth=args.arch_depth, num_heads=args.arch_heads,
        query_depth=args.arch_query_depth, query_heads=args.arch_query_heads,
        horizon=traj_len - 1, traj_len=traj_len,
        use_task_conditioning=False, num_tasks=num_tasks,
        decoder_only_modules=("decoder_only.point_flow_head",),
        **diff_kwargs,
    ).cuda().eval()
    sd = torch.load(CKPT, map_location="cuda")
    if isinstance(sd, dict) and "model" in sd: sd = sd["model"]
    if isinstance(sd, dict) and "state_dict" in sd: sd = sd["state_dict"]
    missing, unexpected = model.load_state_dict(sd, strict=False)
    print(f"loaded ckpt: missing={len(missing)} unexpected={len(unexpected)}")

    rgbs, pos_gt, vis_gt, (fx, fy, cx, cy), meta = load_ep(EP_DIR)
    T_ep = pos_gt.shape[0]; N = pos_gt.shape[1]; H, W = rgbs[0].shape[:2]
    print(f"ep T={T_ep}, N={N}")

    # Build a fake batch matching the dataloader format. Use the episode's actual values.
    x0 = torch.from_numpy(pos_gt[0]).float().unsqueeze(0)            # (1, N, 3)
    cum = torch.from_numpy(pos_gt - pos_gt[0:1]).float()             # (T, N, 3)
    # Pad cum to traj_len-1 steps (drop t=0 since cum[0]==0)
    L = traj_len - 1
    cum_pad = torch.zeros(L, N, 3)
    pmask = torch.zeros(L, N, dtype=torch.bool)
    smask = torch.zeros(L, dtype=torch.bool)
    nsteps = min(L, T_ep - 1)
    cum_pad[:nsteps] = torch.from_numpy(pos_gt[1:1+nsteps] - pos_gt[0:1])
    pmask[:nsteps] = torch.from_numpy(vis_gt[1:1+nsteps]).bool()
    smask[:nsteps] = True
    cum = cum_pad.unsqueeze(0); point_mask = pmask.unsqueeze(0); step_mask = smask.unsqueeze(0)
    action_mask = torch.zeros(1, dtype=torch.bool)

    norm = normalize_and_augment(args, x0, cum, None, None, None,
                                 step_mask, action_mask, point_mask, "cuda",
                                 apply_noise=False)
    feed = {"x0": norm["x0n"], "step_mask": norm["step_mask"],
            "point_mask": norm["point_mask"], "num_steps": L,
            "flow_loss_w": 0.0, "cum_target": norm["cum_gtn"]}
    with torch.no_grad():
        out = model(feed, mode="eval")
    pred_pos_norm = out["pred_pos"][0].cpu().numpy()    # (L, N, 3) normalized

    # Un-normalize: pred_pos_norm * scale + center  (inverse of x0 normalization)
    scales = norm["scales"][0].cpu().numpy()
    center = (x0[0].numpy()).mean(axis=0)
    # Re-derive normalization: x0n = (x0 - center) / scales — inverse: x = x0n * scales + center
    # But normalize_and_augment uses 'compute_start_sphere_params(x0)' — need exact inverse.
    # Easier: read scales/center directly from norm output and apply the documented formula:
    # x0n = sim_transform(x0, scales=inv_scales, translations=-center)
    # → x0n = (x0 - center) * inv_scales = (x0 - center) / scales
    # So: x0_world = x0n * scales + center
    pred_pos_world = pred_pos_norm * scales + center      # (L, N, 3)
    print(f"pred range world: x={pred_pos_world[:,:,0].min():.3f}..{pred_pos_world[:,:,0].max():.3f}  "
          f"z={pred_pos_world[:,:,2].min():.3f}..{pred_pos_world[:,:,2].max():.3f}")
    print(f"GT range world:   x={pos_gt[:,:,0].min():.3f}..{pos_gt[:,:,0].max():.3f}  "
          f"z={pos_gt[:,:,2].min():.3f}..{pos_gt[:,:,2].max():.3f}")

    # Project both
    pred_uv = project(pred_pos_world, fx, fy, cx, cy)
    gt_uv   = project(pos_gt[1:1+nsteps], fx, fy, cx, cy)
    init_uv = project(pos_gt[0:1], fx, fy, cx, cy)[0]      # (N, 2)

    out_frames = []
    for t in range(T_ep):
        f = rgbs[t].copy()
        # Initial points (blue) — show always
        for n in range(0, N, 4):
            x, y = init_uv[n]
            if 0 <= x < W and 0 <= y < H:
                cv2.circle(f, (int(x), int(y)), 3, (50, 100, 255), -1, cv2.LINE_AA)
        # GT (green) at this timestep
        if t > 0 and (t - 1) < gt_uv.shape[0]:
            for n in range(0, N, 4):
                if not vis_gt[t, n]: continue
                x, y = gt_uv[t-1, n]
                if 0 <= x < W and 0 <= y < H:
                    cv2.circle(f, (int(x), int(y)), 3, (0, 255, 0), -1, cv2.LINE_AA)
        # Pred (red) at this timestep
        if t > 0 and (t - 1) < pred_uv.shape[0]:
            for n in range(0, N, 4):
                x, y = pred_uv[t-1, n]
                if 0 <= x < W and 0 <= y < H:
                    cv2.circle(f, (int(x), int(y)), 3, (255, 50, 50), -1, cv2.LINE_AA)
        cv2.putText(f, f"t={t}  blue=init green=GT red=PRED", (16, 32),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2, cv2.LINE_AA)
        out_frames.append(f)

    imageio.mimwrite(OUT_VIDEO, out_frames, fps=15, macro_block_size=1)
    print(f"saved {OUT_VIDEO}")


if __name__ == "__main__":
    main()