# demo/preview.py
import numpy as np
from visualizer import draw_box_on_frame, draw_trajectory_on_frame
from utils.box_utils import interpolate_boxes

def preview_trajectory(
    first_frame: np.ndarray,   # [H, W, 3]
    start_box: list,           # [x1, y1, x2, y2]
    end_box:   list,           # [x1, y1, x2, y2]
    num_frames: int = 81
) -> np.ndarray:
    """
    Shows the planned trajectory on the first frame BEFORE running.
    User sees this immediately after drawing boxes — fast feedback.
    """
    keyboxes = {0: start_box, num_frames - 1: end_box}
    boxes = interpolate_boxes(keyboxes, num_frames)

    frame = first_frame.copy()

    # Draw full trajectory path (center points)
    centers = np.stack([
        (boxes[:, 0] + boxes[:, 2]) / 2,
        (boxes[:, 1] + boxes[:, 3]) / 2
    ], axis=1).astype(int)

    for i in range(1, len(centers)):
        alpha = i / len(centers)
        color = (
            int(255 * (1 - alpha)),
            int(200 * alpha),
            255
        )
        import cv2
        cv2.line(frame,
                 tuple(centers[i-1]),
                 tuple(centers[i]),
                 color, 2)

    # Draw start box (solid yellow)
    frame = draw_box_on_frame(
        frame, start_box,
        color=(255, 220, 0),
        label="START",
        dashed=False
    )

    # Draw end box (dashed yellow)
    frame = draw_box_on_frame(
        frame, end_box,
        color=(255, 220, 0),
        label="END",
        dashed=True
    )

    # Draw a few intermediate boxes (faded)
    for i in [20, 40, 60]:
        if i < len(boxes):
            frame = draw_box_on_frame(
                frame, boxes[i],
                color=(200, 200, 200),
                label=f"t={i}",
                dashed=True,
                thickness=1
            )

    return frame


def preview_trajectory_strip(
    frames: np.ndarray,     # [T, H, W, 3]
    start_box: list,
    end_box:   list,
) -> np.ndarray:
    """
    Shows predicted box overlaid on 5 sampled frames.
    Gives sense of how box moves through the video.
    """
    T = len(frames)
    keyboxes = {0: start_box, T - 1: end_box}
    boxes = interpolate_boxes(keyboxes, T)

    sample_ts = [0, T//4, T//2, 3*T//4, T-1]
    previews = []

    for t in sample_ts:
        frame = frames[t].copy()
        frame = draw_box_on_frame(
            frame, boxes[t],
            color=(0, 255, 255),
            label=f"t={t}",
            dashed=(t > 0)
        )
        # Add small frame counter
        import cv2
        H, W = frame.shape[:2]
        progress = f"{t}/{T-1}"
        cv2.putText(frame, progress, (W-80, H-10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    (200, 200, 200), 1)
        previews.append(frame)

    return np.concatenate(previews, axis=1)  # horizontal strip