"""Visualize a tracks.npz on top of source frames: point IDs are color-coded,
visibility shown via marker style, and trails up to the current frame are drawn."""
import argparse, os
import numpy as np
import cv2
import imageio


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


def colormap(N):
    rng = np.random.default_rng(0)
    return (rng.uniform(0, 255, (N, 3))).astype(np.uint8)


def draw_frame(frame, tracks_so_far, vis_so_far, colors, trail_alpha=0.6, radius=4):
    """tracks_so_far: (t+1, N, 2). Draw trail and current-pt for each id."""
    img = frame.copy()
    overlay = img.copy()
    T_now, N, _ = tracks_so_far.shape

    # Trails (semi-transparent)
    for n in range(N):
        pts = tracks_so_far[:, n, :].astype(np.int32)
        v   = vis_so_far[:, n]
        col = tuple(int(c) for c in colors[n])
        # draw segments where both endpoints visible
        for i in range(1, T_now):
            if v[i] and v[i-1]:
                cv2.line(overlay, tuple(pts[i-1]), tuple(pts[i]), col, 1, cv2.LINE_AA)
    img = cv2.addWeighted(overlay, trail_alpha, img, 1 - trail_alpha, 0)

    # Current points: filled circle if visible, hollow ring otherwise
    cur = tracks_so_far[-1].astype(np.int32)
    cur_v = vis_so_far[-1]
    for n in range(N):
        col = tuple(int(c) for c in colors[n])
        x, y = int(cur[n, 0]), int(cur[n, 1])
        if cur_v[n]:
            cv2.circle(img, (x, y), radius, col, -1, cv2.LINE_AA)
            cv2.circle(img, (x, y), radius, (0, 0, 0), 1, cv2.LINE_AA)
        else:
            cv2.circle(img, (x, y), radius, col, 1, cv2.LINE_AA)
    return img


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--video", required=True)
    ap.add_argument("--tracks_npz", required=True)
    ap.add_argument("--out_dir", required=True)
    ap.add_argument("--max_frames", type=int, default=0)
    ap.add_argument("--keyframes", type=int, nargs="+", default=None,
                    help="frames to dump as PNG (defaults to evenly spaced 4)")
    ap.add_argument("--make_video", action="store_true",
                    help="also dump a track-with-trails MP4")
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    frames = load_video(args.video, args.max_frames or None)
    z = np.load(args.tracks_npz)
    tracks, vis = z["tracks"], z["vis"].astype(bool)
    T = min(len(frames), tracks.shape[0])
    frames, tracks, vis = frames[:T], tracks[:T], vis[:T]
    N = tracks.shape[1]
    colors = colormap(N)
    print(f"T={T}, N={N}")
    print(f"vis frame 0: {int(vis[0].sum())}/{N}")
    print(f"vis frame -1: {int(vis[-1].sum())}/{N}")

    keyframes = args.keyframes or list(np.linspace(0, T - 1, 4).astype(int))
    panels = []
    for k in keyframes:
        img = draw_frame(frames[k], tracks[:k+1], vis[:k+1], colors)
        cv2.putText(img, f"frame {k}  vis={int(vis[k].sum())}/{N}", (16, 36),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 3, cv2.LINE_AA)
        cv2.putText(img, f"frame {k}  vis={int(vis[k].sum())}/{N}", (16, 36),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 1, cv2.LINE_AA)
        out_p = os.path.join(args.out_dir, f"trail_frame{k:03d}.png")
        imageio.imwrite(out_p, img)
        print(f"  saved {out_p}")
        panels.append(img)

    # Side-by-side panel of key frames
    H, W, _ = panels[0].shape
    panel = np.concatenate(panels, axis=1)
    panel_p = os.path.join(args.out_dir, "panel.png")
    imageio.imwrite(panel_p, panel)
    print(f"  saved {panel_p}  ({panel.shape[1]}x{panel.shape[0]})")

    if args.make_video:
        out = []
        for t in range(T):
            out.append(draw_frame(frames[t], tracks[:t+1], vis[:t+1], colors))
        outp = os.path.join(args.out_dir, "tracks_with_trails.mp4")
        imageio.mimwrite(outp, out, fps=15, macro_block_size=1)
        print(f"  saved {outp}")


if __name__ == "__main__":
    main()