"""Rainbow-keypoint + polyline trajectory overlay.

One function used by both:
- ``vis_dataset.py``: plot recorded EE trajectory on the start frame.
- training viz (wandb) + deploy viz (rerun): plot predicted EE
  trajectory.

Caller provides 2D pixel coords (already projected). Color = rainbow by
time so the temporal ordering is visible. Out-of-frustum points get a
faded color and are clipped to the image bounds.
"""
from __future__ import annotations

from typing import Optional

import cv2
import numpy as np


def _rainbow_colors(n: int) -> np.ndarray:
    """(n, 3) BGR uint8 — HSV rainbow from red→violet."""
    hsv = np.zeros((n, 1, 3), dtype=np.uint8)
    hsv[:, 0, 0] = np.linspace(0, 170, n, dtype=np.uint8)   # OpenCV hue [0, 180)
    hsv[:, 0, 1] = 255
    hsv[:, 0, 2] = 255
    return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[:, 0, :]


def keypoints_overlay(
    img_bgr: np.ndarray,         # (H, W, 3) uint8 background image
    kp_2d: np.ndarray,           # (T, 2) pixel coords (xy)
    in_frustum: Optional[np.ndarray] = None,    # (T,) bool; default all True
    draw_line: bool = True,
    radius_px: int = 6,
    line_thickness: int = 2,
    filled: bool = True,
    color_override: Optional[tuple] = None,
) -> np.ndarray:
    """Return a new uint8 BGR image with markers + optional polyline.

    Two convention combos used in practice:
    - **Pred** (rainbow, filled, with line):
      ``keypoints_overlay(img, pred_pix, draw_line=True)`` (defaults).
    - **GT** (white open circles, no line):
      ``keypoints_overlay(img, gt_pix, filled=False, draw_line=False,
                          color_override=(255, 255, 255))``.

    Out-of-frustum points fade toward gray and are clipped to the border.
    """
    out = img_bgr.copy()
    H, W = out.shape[:2]
    T = len(kp_2d)
    if T == 0:
        return out

    kp = kp_2d.astype(np.float32)
    if in_frustum is None:
        in_frustum = np.ones(T, dtype=bool)
    in_frustum = np.asarray(in_frustum, dtype=bool)
    kp_int = np.clip(np.round(kp), [0, 0], [W - 1, H - 1]).astype(np.int32)
    colors = (np.tile(np.asarray(color_override, dtype=np.uint8), (T, 1))
              if color_override is not None
              else _rainbow_colors(T))

    if draw_line and T >= 2:
        cv2.polylines(out, [kp_int.reshape(-1, 1, 2)], isClosed=False,
                      color=(255, 255, 255), thickness=line_thickness,
                      lineType=cv2.LINE_AA)

    for t in range(T):
        c = tuple(int(v) for v in colors[t])
        if not in_frustum[t]:
            c = tuple(int(0.5 * v + 0.5 * 128) for v in c)
        pt = tuple(kp_int[t])
        if filled:
            cv2.circle(out, pt, radius_px, c, -1, lineType=cv2.LINE_AA)
            cv2.circle(out, pt, radius_px, (0, 0, 0), 1, lineType=cv2.LINE_AA)
        else:
            cv2.circle(out, pt, radius_px, c, 2, lineType=cv2.LINE_AA)
    return out