%matplotlib inline
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchmatch

plt.rcParams.update({"figure.dpi": 120, "font.size": 11})
rng = np.random.default_rng(7)

def box_iou(boxes_a: torch.Tensor, boxes_b: torch.Tensor) -> torch.Tensor:
    """Compute pairwise IoU between two sets of (x1, y1, x2, y2) boxes.

    Args:
        boxes_a: (N, 4) float tensor
        boxes_b: (M, 4) float tensor

    Returns:
        (N, M) float tensor of IoU values.
    """
    # Intersection
    lo = torch.maximum(boxes_a[:, None, :2], boxes_b[None, :, :2])  # (N, M, 2)
    hi = torch.minimum(boxes_a[:, None, 2:], boxes_b[None, :, 2:])  # (N, M, 2)
    inter = (hi - lo).clamp_min(0).prod(-1)  # (N, M)

    area_a = (boxes_a[:, 2:] - boxes_a[:, :2]).prod(-1)  # (N,)
    area_b = (boxes_b[:, 2:] - boxes_b[:, :2]).prod(-1)  # (M,)

    union = area_a[:, None] + area_b[None, :] - inter
    return inter / union.clamp_min(1e-9)


# Quick sanity check
b1 = torch.tensor([[0.0, 0.0, 2.0, 2.0]])  # 2×2 box at origin
b2 = torch.tensor([[1.0, 1.0, 3.0, 3.0]])  # 2×2 box shifted by (1, 1)
iou = box_iou(b1, b2)
# Intersection = 1×1 = 1; union = 4 + 4 - 1 = 7; IoU = 1/7
print(f"IoU (expected ≈ 1/7 ≈ 0.143): {iou.item():.3f}")

T = 20        # frames
N_obj = 6     # number of objects
IMG = 512.0   # image width/height (pixels)
BOX_W = 40.0  # box half-width

# Each object moves along a random linear trajectory
starts = rng.uniform(60, IMG - 60, size=(N_obj, 2))   # (x_c, y_c)
velocities = rng.uniform(-8, 8, size=(N_obj, 2))       # pixels per frame

def object_boxes(t):
    """Ground-truth boxes at frame t, shape (N_obj, 4)."""
    centers = starts + velocities * t
    centers = centers.clip(BOX_W + 1, IMG - BOX_W - 1)
    x1y1 = centers - BOX_W
    x2y2 = centers + BOX_W
    return np.concatenate([x1y1, x2y2], axis=1).astype(np.float32)

def noisy_detections(t, noise_std=8.0):
    """Simulated detector output: noisy boxes with some missed detections."""
    boxes = object_boxes(t)
    noise = rng.normal(0, noise_std, size=boxes.shape)
    return (boxes + noise).clip(0, IMG).astype(np.float32)

# Visualise first 4 frames
fig, axes = plt.subplots(1, 4, figsize=(14, 3.5))
colors = plt.cm.tab10(np.linspace(0, 0.9, N_obj))

for col, t in enumerate([0, 5, 10, 15]):
    ax = axes[col]
    ax.set_xlim(0, IMG)
    ax.set_ylim(IMG, 0)  # image coords: y increases downward
    ax.set_aspect("equal")
    ax.set_title(f"Frame {t}")
    ax.set_facecolor("#111111")

    for i, (box, color) in enumerate(zip(object_boxes(t), colors)):
        x1, y1, x2, y2 = box
        rect = patches.Rectangle(
            (x1, y1), x2 - x1, y2 - y1,
            linewidth=2, edgecolor=color, facecolor="none",
        )
        ax.add_patch(rect)
        ax.text(x1 + 2, y1 + 14, f"obj{i}", color=color, fontsize=8)

plt.suptitle("Simulated video — ground-truth boxes", y=1.02)
plt.tight_layout()
plt.show()

class SimpleTracker:
    """SORT-style tracker using torchmatch.assignment.solve."""

    GATE = 0.1  # detections with IoU < GATE are forbidden (set to +inf)

    def __init__(self):
        self.tracks = {}          # track_id → last box (np.ndarray, shape (4,))
        self.next_id = 0
        self.history = {}         # track_id → list of center positions

    def update(self, detections: np.ndarray) -> dict:
        """
        Args:
            detections: (M, 4) array of (x1, y1, x2, y2) boxes

        Returns:
            dict mapping track_id → matched detection box
        """
        M = len(detections)
        track_ids = list(self.tracks.keys())
        N = len(track_ids)

        matched_out = {}

        if N == 0 or M == 0:
            # No tracks or no detections: birth all detections as new tracks
            for box in detections:
                self._birth(box)
            return matched_out

        track_boxes = np.stack([self.tracks[tid] for tid in track_ids])
        track_t = torch.from_numpy(track_boxes)
        det_t = torch.from_numpy(detections)

        iou = box_iou(track_t, det_t)          # (N, M)
        cost = 1.0 - iou                        # lower = better match
        cost[iou < self.GATE] = float("inf")    # gate implausible pairs

        # Solve the assignment
        row_to_col = torchmatch.assignment.solve(cost)

        matched_dets = set()
        for i, j in enumerate(row_to_col.tolist()):
            if j == -1 or cost[i, j] == float("inf"):
                continue  # track i unmatched
            tid = track_ids[i]
            box = detections[j]
            self.tracks[tid] = box
            self.history[tid].append(_center(box))
            matched_out[tid] = box
            matched_dets.add(j)

        # Birth new tracks for unmatched detections
        for j in range(M):
            if j not in matched_dets:
                self._birth(detections[j])

        return matched_out

    def _birth(self, box):
        tid = self.next_id
        self.next_id += 1
        self.tracks[tid] = box
        self.history[tid] = [_center(box)]

def _center(box):
    return ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)

tracker = SimpleTracker()

for t in range(T):
    dets = noisy_detections(t)
    tracker.update(dets)

print(f"Tracker assigned {tracker.next_id} unique IDs over {T} frames")
print(f"Active tracks at end: {list(tracker.tracks.keys())}")

fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xlim(0, IMG)
ax.set_ylim(IMG, 0)
ax.set_facecolor("#1a1a2e")
ax.set_aspect("equal")
ax.set_title(f"Tracked trajectories  ({T} frames)", color="white", fontsize=13)
ax.tick_params(colors="white")
fig.patch.set_facecolor("#1a1a2e")

cmap = plt.cm.rainbow
track_colors = {tid: cmap(tid / max(tracker.next_id, 1))
                for tid in tracker.history}

for tid, centers in tracker.history.items():
    if len(centers) < 2:
        continue
    xs, ys = zip(*centers)
    color = track_colors[tid]
    ax.plot(xs, ys, "-o", color=color, markersize=4, lw=1.5, alpha=0.85)
    ax.text(xs[-1], ys[-1], f" {tid}", color=color, fontsize=9)

# Also draw ground-truth trajectories (dashed)
for i in range(N_obj):
    gt_centers = [_center(object_boxes(t)[i]) for t in range(T)]
    xs, ys = zip(*gt_centers)
    ax.plot(xs, ys, "--", color="white", alpha=0.2, lw=1)

legend_elems = [
    plt.Line2D([0], [0], color="white", ls="--", alpha=0.4, label="Ground truth"),
    plt.Line2D([0], [0], color="cyan", lw=2, label="Tracker output"),
]
ax.legend(handles=legend_elems, facecolor="#333", labelcolor="white")
plt.tight_layout()
plt.show()

# Stack all frames into a batch
all_track_boxes = np.stack([object_boxes(t) for t in range(T)])  # (T, N_obj, 4)
all_det_boxes = np.stack([noisy_detections(t) for t in range(T)])  # (T, N_obj, 4)

all_iou = torch.stack([
    box_iou(torch.from_numpy(all_track_boxes[t]),
            torch.from_numpy(all_det_boxes[t]))
    for t in range(T)
])  # (T, N_obj, N_obj)

cost_batch = 1.0 - all_iou
cost_batch[all_iou < SimpleTracker.GATE] = float("inf")

# Solve all T frames at once
assignments_batch = torchmatch.assignment.solve(cost_batch)  # (T, N_obj)
print("Batched input shape :", cost_batch.shape)
print("Output shape        :", assignments_batch.shape)

# Verify: assignments should be close to the identity permutation when
# objects don't cross (they mostly shouldn't for our slow velocities)
identity_hits = (assignments_batch == torch.arange(N_obj).unsqueeze(0)).float().mean()
print(f"Frames where all objects matched correctly: {identity_hits:.1%}")