# demo/visualizer.py
import numpy as np
import cv2
from typing import Optional

def draw_box_on_frame(
    frame: np.ndarray,   # [H, W, 3] uint8 RGB
    box: list,           # [x1, y1, x2, y2]
    color: tuple = (255, 255, 0),
    label: str = "",
    thickness: int = 2,
    dashed: bool = False
) -> np.ndarray:
    """Draw a single bounding box on a frame"""
    frame = frame.copy()
    x1, y1, x2, y2 = [int(v) for v in box]

    if dashed:
        # Draw dashed rectangle manually
        dash_len = 10
        gap_len  = 5
        pts = [
            ((x1, y1), (x2, y1)),  # top
            ((x2, y1), (x2, y2)),  # right
            ((x2, y2), (x1, y2)),  # bottom
            ((x1, y2), (x1, y1)),  # left
        ]
        for (px1, py1), (px2, py2) in pts:
            dx = px2 - px1
            dy = py2 - py1
            dist = max(abs(dx), abs(dy))
            if dist == 0:
                continue
            for i in range(0, dist, dash_len + gap_len):
                s = i / dist
                e = min(i + dash_len, dist) / dist
                sx = int(px1 + s * dx)
                sy = int(py1 + s * dy)
                ex = int(px1 + e * dx)
                ey = int(py1 + e * dy)
                cv2.line(frame, (sx, sy), (ex, ey), color, thickness)
    else:
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness)

    if label:
        cv2.putText(
            frame, label,
            (x1, max(y1 - 8, 12)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6, color, 2
        )

    return frame


def draw_trajectory_on_frame(
    frame: np.ndarray,
    boxes: np.ndarray,    # [T, 4] — full trajectory
    current_t: int,
    color: tuple = (255, 200, 0)
) -> np.ndarray:
    """
    Draw the motion path (center points) up to current frame.
    Gives a visual "trail" showing where the object came from.
    """
    frame = frame.copy()
    centers = np.stack([
        (boxes[:, 0] + boxes[:, 2]) / 2,
        (boxes[:, 1] + boxes[:, 3]) / 2
    ], axis=1).astype(int)

    # Draw path line
    for i in range(1, current_t + 1):
        alpha = i / (current_t + 1)  # fade older points
        c = tuple(int(v * alpha) for v in color)
        cv2.line(
            frame,
            tuple(centers[i-1]),
            tuple(centers[i]),
            c, 2
        )

    # Draw current center dot
    cv2.circle(frame, tuple(centers[current_t]), 5, color, -1)

    return frame


def create_comparison_strip(
    original: np.ndarray,   # [T, H, W, 3]
    result:   np.ndarray,   # [T, H, W, 3]
    pred_boxes: np.ndarray, # [T, 4]
    sample_ts: list = None  # which frames to show
) -> np.ndarray:
    """
    Creates a horizontal strip for visual comparison.
    Shows: Original | Result | Diff for N sampled frames.
    """
    T = len(original)
    if sample_ts is None:
        sample_ts = [0, T//4, T//2, 3*T//4, T-1]

    rows = []
    for t in sample_ts:
        orig_t = original[t].copy()
        res_t  = result[t].copy()

        # Draw box on result
        res_t = draw_box_on_frame(
            res_t, pred_boxes[t],
            color=(0, 255, 0),
            label=f"t={t}"
        )

        # Amplified diff
        diff_t = np.abs(
            orig_t.astype(np.int32) - result[t].astype(np.int32)
        )
        diff_t = (diff_t * 4).clip(0, 255).astype(np.uint8)

        # Add labels
        def add_label(img, text):
            img = img.copy()
            cv2.putText(img, text, (10, 25),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                        (255, 255, 255), 2)
            cv2.putText(img, text, (10, 25),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                        (0, 0, 0), 1)
            return img

        orig_t = add_label(orig_t, "Original")
        res_t  = add_label(res_t,  "Result")
        diff_t = add_label(diff_t, "Diff x4")

        row = np.concatenate([orig_t, res_t, diff_t], axis=1)
        rows.append(row)

    return np.concatenate(rows, axis=0)