"""Overlay 3-PoinTr object tracks (3D positions, projected via stored intrinsics)
onto the existing hand-action _pred.mp4 frames, producing _pred_handobject.mp4.

The hand silhouettes (black) are preserved; object tracks are drawn only on the
white background. Using small colored dots avoids visually colliding with the
hand silhouettes the WM is already conditioned on.
"""
import argparse, os, json
import numpy as np
import torch
import cv2
import imageio.v2 as imageio


def load_video(path):
    cap = cv2.VideoCapture(path); frames = []
    while True:
        r, f = cap.read()
        if not r: break
        frames.append(cv2.cvtColor(f, cv2.COLOR_BGR2RGB))
    cap.release()
    return frames


def load_episode(ep_dir):
    meta = json.load(open(os.path.join(ep_dir, "metadata.json")))
    T = meta["T"]
    pos, vis = [], []
    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"), weights_only=False).numpy())
        vis.append(torch.load(os.path.join(ts, f"{k}_visibility.pth"), weights_only=False).numpy())
    pos = np.stack(pos, 0); vis = np.stack(vis, 0)
    K = meta["intrinsics"]["camera"]
    return 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 is_white(pixel, thresh=235):
    return (pixel >= thresh).all()


def overlay_object_tracks(pred_frames, pos_3d, vis, K, *, color="red",
                           dot_radius=3, stride=1, draw_only_on_white=True,
                           white_thresh=230, trail_len=0):
    """Draw colored dots on the pred frames at projected object positions.

    `draw_only_on_white=True` ensures hand silhouettes (black) are never overwritten,
    so the WM can still see the same hand instruction.
    """
    fx, fy, cx, cy = K
    Tp = len(pred_frames)
    Tt = pos_3d.shape[0]
    T = min(Tp, Tt)
    H, W = pred_frames[0].shape[:2]

    color_map = {
        "red": (220, 30, 30),
        "blue": (30, 120, 255),
        "green": (40, 200, 40),
        "magenta": (220, 30, 200),
    }
    col_rgb = color_map[color]

    out = []
    proj = project(pos_3d, fx, fy, cx, cy)  # (T, N, 2)
    for t in range(T):
        f = pred_frames[t].copy()
        # current-frame dots
        for n in range(0, pos_3d.shape[1], stride):
            if not vis[t, n]:
                continue
            x, y = proj[t, n]
            if not (0 <= x < W and 0 <= y < H):
                continue
            xi, yi = int(round(x)), int(round(y))
            if draw_only_on_white and not is_white(f[yi, xi], white_thresh):
                # skip drawing on top of hand silhouette pixels
                continue
            cv2.circle(f, (xi, yi), dot_radius, col_rgb, -1, cv2.LINE_AA)
        # short trail
        if trail_len > 0 and t > 0:
            t0 = max(0, t - trail_len)
            for n in range(0, pos_3d.shape[1], stride):
                for ti in range(t0 + 1, t + 1):
                    if not (vis[ti, n] and vis[ti - 1, n]):
                        continue
                    x1, y1 = proj[ti - 1, n]
                    x2, y2 = proj[ti, n]
                    if not (0 <= x1 < W and 0 <= y1 < H and 0 <= x2 < W and 0 <= y2 < H):
                        continue
                    cv2.line(f, (int(x1), int(y1)), (int(x2), int(y2)), col_rgb, 1, cv2.LINE_AA)
        out.append(f)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pred_video", required=True)
    ap.add_argument("--ep_dir", required=True,
                    help="3-PoinTr episode dir, e.g. .../no_actions/stack_can/0001")
    ap.add_argument("--out_video", required=True)
    ap.add_argument("--color", default="red", choices=["red", "blue", "green", "magenta"])
    ap.add_argument("--stride", type=int, default=1)
    ap.add_argument("--dot_radius", type=int, default=3)
    ap.add_argument("--trail_len", type=int, default=0)
    ap.add_argument("--no_white_only", action="store_true",
                    help="Allow drawing on hand silhouette too (default protects hands).")
    args = ap.parse_args()

    print(f"[in]  pred  : {args.pred_video}")
    print(f"[in]  ep_dir: {args.ep_dir}")
    pred_frames = load_video(args.pred_video)
    pos, vis, K, meta = load_episode(args.ep_dir)
    print(f"  pred T={len(pred_frames)}  tracks T={pos.shape[0]}  N={pos.shape[1]}")
    print(f"  intrinsics: fx={K[0]:.1f} fy={K[1]:.1f} cx={K[2]:.1f} cy={K[3]:.1f}")

    out = overlay_object_tracks(
        pred_frames, pos, vis, K,
        color=args.color, dot_radius=args.dot_radius, stride=args.stride,
        draw_only_on_white=(not args.no_white_only),
        trail_len=args.trail_len,
    )
    imageio.mimwrite(args.out_video, out, fps=15, macro_block_size=1)
    print(f"[out] {args.out_video}  ({len(out)} frames)")


if __name__ == "__main__":
    main()