from __future__ import annotations

import argparse
import io
import json
import os
import tempfile
import time
from collections import defaultdict
from pathlib import Path

os.environ.setdefault("MPLCONFIGDIR", str(Path(tempfile.gettempdir()) / "matplotlib"))
os.makedirs(os.environ["MPLCONFIGDIR"], exist_ok=True)

import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch

from plot_artifacts import GradCam, attention_boxes, evenly_sample_rows, load_manifest
from train import build_model, read_labels
from triangulation_3d import az_el_from_points, triangulate_from_az_el


DEFAULT_MODULE_A = np.array([0.0, 0.0, 1.4])
DEFAULT_MODULE_B = np.array([18.0, 0.0, 1.4])
DEFAULT_TARGET = np.array([8.5, 28.0, 11.0])
RECEPTOR_SPECS = [
    {"name": "RF receptor 1", "module": "Module A", "axis": "horizontal pair", "time_offset": 0},
    {"name": "RF receptor 2", "module": "Module A", "axis": "vertical pair", "time_offset": 2},
    {"name": "RF receptor 3", "module": "Module B", "axis": "horizontal pair", "time_offset": 4},
    {"name": "RF receptor 4", "module": "Module B", "axis": "vertical pair", "time_offset": 6},
]


def normalize_for_model(spec: np.ndarray) -> torch.Tensor:
    x = spec.astype(np.float32)
    x = (x - x.mean()) / (x.std() + 1e-6)
    return torch.from_numpy(x).unsqueeze(0).unsqueeze(0)


def load_resnet18(processed_dir: Path, checkpoint_dir: Path, device: torch.device):
    labels = read_labels(processed_dir)
    model = build_model("resnet18", len(labels)).to(device)
    checkpoint_path = checkpoint_dir / "resnet18" / "best_model.pt"
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint["model"])
    model.eval()
    return model, labels, checkpoint_path


def classify_with_gradcam(
    model: torch.nn.Module,
    gradcam: GradCam,
    labels: list[str],
    spec: np.ndarray,
    device: torch.device,
) -> dict[str, object]:
    model_input = normalize_for_model(spec).to(device)
    if device.type == "cuda":
        torch.cuda.synchronize()
    started = time.perf_counter()
    logits = model(model_input)
    if device.type == "cuda":
        torch.cuda.synchronize()
    latency_ms = (time.perf_counter() - started) * 1000.0
    pred_idx = int(logits.argmax(dim=1).item())
    cam = gradcam(logits, pred_idx, spec.shape)
    boxes = attention_boxes(cam, threshold=0.70, min_area=16, max_boxes=2)
    peak_y, peak_x = np.unravel_index(int(np.argmax(spec)), spec.shape)
    return {
        "prediction": labels[pred_idx],
        "latency_ms": latency_ms,
        "cam": cam,
        "boxes": boxes,
        "peak_bin": [int(peak_y), int(peak_x)],
        "peak_power_db": float(spec[peak_y, peak_x]),
    }


def make_receptor_view(spec: np.ndarray, time_offset: int) -> np.ndarray:
    """Use a circular offset so receptor views stay real-looking without padded edge artifacts."""
    if time_offset <= 0:
        return spec.copy()
    return np.roll(spec, shift=-time_offset, axis=1).copy()


def draw_left_telemetry_spectrogram(
    spec: np.ndarray,
    cam: np.ndarray,
    boxes: list[tuple[int, int, int, int, float]],
    out_path: Path,
    title: str,
    telemetry: list[tuple[str, str]],
) -> None:
    fig = plt.figure(figsize=(14, 6), facecolor="#05070a")
    gs = fig.add_gridspec(1, 2, width_ratios=[1.45, 3.55], wspace=0.16)
    side = fig.add_subplot(gs[0, 0])
    ax = fig.add_subplot(gs[0, 1])

    side.set_facecolor("#05070a")
    side.axis("off")
    side.set_xlim(0, 1)
    side.set_ylim(0, 1)
    side.text(0.05, 0.95, "CLASSIFIER TELEMETRY", color="white", fontsize=13.5, weight="bold", va="top")
    y = 0.82
    for key, value in telemetry:
        side.text(0.05, y, key.upper(), color="#6f7890", fontsize=8.5, weight="bold", va="top")
        side.text(0.05, y - 0.04, value, color="#f1f5f9", fontsize=11, va="top", linespacing=1.15)
        line_count = str(value).count("\n") + 1
        y -= 0.098 + (line_count - 1) * 0.035

    ax.set_facecolor("#05070a")
    height, width = spec.shape
    image_extent = (0, width, 0, height)
    vmin, vmax = np.percentile(spec, [2, 98])
    ax.imshow(spec, aspect="auto", origin="lower", extent=image_extent, cmap="turbo", vmin=vmin, vmax=vmax)
    ax.imshow(
        cam,
        aspect="auto",
        origin="lower",
        extent=image_extent,
        cmap="cool",
        alpha=np.clip(cam, 0.0, 0.65) * 0.42,
    )

    for idx, (x1, y1, x2, y2, score) in enumerate(boxes, start=1):
        x1 = max(0, min(width, int(x1)))
        x2 = max(0, min(width, int(x2)))
        y1 = max(0, min(height, int(y1)))
        y2 = max(0, min(height, int(y2)))
        edge = "#00f5ff" if idx == 1 else "#ffe66d"
        rect = patches.Rectangle(
            (x1, y1),
            x2 - x1,
            y2 - y1,
            linewidth=2.2,
            edgecolor=edge,
            facecolor="none",
        )
        ax.add_patch(rect)
        ax.text(
            x1,
            min(height - 2, y2 + 2),
            f"A{idx} {score:.2f}",
            color="#05070a",
            fontsize=8,
            weight="bold",
            bbox=dict(facecolor=edge, edgecolor="#05070a", boxstyle="square,pad=0.15"),
        )

    peak_y, peak_x = np.unravel_index(int(np.argmax(spec)), spec.shape)
    ax.scatter([peak_x + 0.5], [peak_y + 0.5], marker="x", s=48, c="white", linewidths=1.8)
    ax.set_title(title, color="white", fontsize=17, pad=12)
    ax.set_xlabel("Time frame", color="white", fontsize=12)
    ax.set_ylabel("Frequency bin", color="white", fontsize=12, labelpad=8)
    ax.tick_params(colors="#9aa4b2")
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    for spine in ax.spines.values():
        spine.set_color("#1f2937")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=160, facecolor=fig.get_facecolor(), bbox_inches="tight")
    plt.close(fig)


def feature_angle_offsets(classifications: list[dict[str, object]], spec_shape: tuple[int, int]) -> tuple[float, float, float, float]:
    height, width = spec_shape
    peaks = [item["peak_bin"] for item in classifications]
    x_norm = [(peak[1] / max(1, width - 1)) - 0.5 for peak in peaks]
    y_norm = [(peak[0] / max(1, height - 1)) - 0.5 for peak in peaks]
    az_a_offset = float(np.mean(x_norm[:2]) * 0.75)
    el_a_offset = float(np.mean(y_norm[:2]) * 0.50)
    az_b_offset = float(np.mean(x_norm[2:]) * 0.75)
    el_b_offset = float(np.mean(y_norm[2:]) * 0.50)
    return az_a_offset, el_a_offset, az_b_offset, el_b_offset


def draw_triangulation_3d(
    out_path: Path,
    result,
    angles: dict[str, float],
    label: str,
    predictions: list[str],
) -> None:
    fig = plt.figure(figsize=(11, 8), facecolor="#05070a")
    ax = fig.add_subplot(111, projection="3d", facecolor="#05070a")
    pane = (0.46, 0.46, 0.46, 1.0)
    grid = (0.72, 0.72, 0.72, 0.55)
    for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
        axis.set_pane_color(pane)
        axis._axinfo["grid"]["color"] = grid
        axis._axinfo["axisline"]["color"] = (0.9, 0.9, 0.9, 1.0)
        axis._axinfo["tick"]["color"] = (0.9, 0.9, 0.9, 1.0)

    ax.scatter(*result.module_a, s=95, c="#00f5ff", marker="^", label="Module A")
    ax.scatter(*result.module_b, s=95, c="#ffe66d", marker="^", label="Module B")
    ax.scatter(*result.estimated_position, s=130, c="#ff4d6d", marker="*", label="Estimated drone")

    for origin, direction, color, name in [
        (result.module_a, result.direction_a, "#00f5ff", "Ray A"),
        (result.module_b, result.direction_b, "#ffe66d", "Ray B"),
    ]:
        distances = np.linspace(0, 38, 2)
        points = origin[:, None] + direction[:, None] * distances
        ax.plot(points[0], points[1], points[2], color=color, linewidth=2.4, label=name)

    ax.plot(
        [result.closest_a[0], result.closest_b[0]],
        [result.closest_a[1], result.closest_b[1]],
        [result.closest_a[2], result.closest_b[2]],
        color="white",
        linestyle="--",
        linewidth=1.8,
        label="Closest segment",
    )

    ax.set_title(f"3D AoA triangulation | {label}", color="white", fontsize=17, pad=18)
    ax.set_xlabel("East (m)", color="white", labelpad=10)
    ax.set_ylabel("North (m)", color="white", labelpad=10)
    ax.set_zlabel("Up (m)", color="white", labelpad=10)
    ax.tick_params(colors="#d1d5db")
    ax.set_xlim(-3, 22)
    ax.set_ylim(-2, 34)
    ax.set_zlim(0, 16)
    ax.view_init(elev=24, azim=-58)
    ax.legend(loc="upper left", facecolor="#111827", edgecolor="#374151", labelcolor="white")
    ax.text2D(
        0.03,
        0.04,
        (
            f"A: az {angles['az_a']:.1f} deg, el {angles['el_a']:.1f} deg\n"
            f"B: az {angles['az_b']:.1f} deg, el {angles['el_b']:.1f} deg\n"
            f"Estimated ENU: ({result.estimated_position[0]:.1f}, {result.estimated_position[1]:.1f}, {result.estimated_position[2]:.1f}) m\n"
            f"Residual: {result.residual_m:.2f} m | predictions: {', '.join(sorted(set(predictions)))}"
        ),
        transform=ax.transAxes,
        color="#e5e7eb",
        fontsize=10,
        bbox=dict(facecolor="#111827", edgecolor="#374151", alpha=0.88, boxstyle="round,pad=0.45"),
    )

    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=170, facecolor=fig.get_facecolor(), bbox_inches="tight")
    plt.close(fig)


def target_position_for_frame(frame_idx: int, total_frames: int) -> np.ndarray:
    if total_frames <= 1:
        return DEFAULT_TARGET.copy()
    t = frame_idx / (total_frames - 1)
    phase = 2.0 * np.pi * t
    return np.array(
        [
            4.5 + 11.0 * t + 2.4 * np.sin(2.2 * phase),
            14.5 + 16.0 * t + 3.2 * np.sin(1.1 * phase + 0.6),
            8.8 + 2.7 * np.sin(1.5 * phase) + 1.2 * np.cos(0.7 * phase),
        ],
        dtype=np.float64,
    )


def classify_receptor_views(
    base_spec: np.ndarray,
    source_row: dict[str, str],
    model: torch.nn.Module,
    gradcam: GradCam,
    labels: list[str],
    device: torch.device,
) -> list[dict[str, object]]:
    classifications = []
    for receptor in RECEPTOR_SPECS:
        spec = make_receptor_view(base_spec, receptor["time_offset"])
        result = classify_with_gradcam(model, gradcam, labels, spec, device)
        result.update(receptor)
        result["spec"] = spec
        result["source_sample"] = source_row
        result["time_offset"] = receptor["time_offset"]
        result["part_window"] = f"circular time offset {receptor['time_offset']}"
        classifications.append(result)
    return classifications


def load_classified_frame(
    row: dict[str, str],
    sample_dir: Path,
    model: torch.nn.Module,
    gradcam: GradCam,
    labels: list[str],
    device: torch.device,
) -> tuple[np.ndarray, list[dict[str, object]]]:
    base_spec = np.load(sample_dir / row["path"])["x"].astype(np.float32)
    classifications = classify_receptor_views(base_spec, row, model, gradcam, labels, device)
    return base_spec, classifications


def find_agreeing_frame(
    candidate_rows: list[dict[str, str]],
    start_idx: int,
    sample_dir: Path,
    model: torch.nn.Module,
    gradcam: GradCam,
    labels: list[str],
    device: torch.device,
    target_label: str,
    max_attempts: int = 80,
) -> tuple[dict[str, str], np.ndarray, list[dict[str, object]]]:
    attempts = min(max_attempts, len(candidate_rows))
    fallback: tuple[dict[str, str], np.ndarray, list[dict[str, object]]] | None = None
    for attempt in range(attempts):
        row = candidate_rows[(start_idx + attempt) % len(candidate_rows)]
        base_spec, classifications = load_classified_frame(row, sample_dir, model, gradcam, labels, device)
        predictions = [str(item["prediction"]) for item in classifications]
        if fallback is None:
            fallback = (row, base_spec, classifications)
        if all(pred == target_label for pred in predictions):
            return row, base_spec, classifications
    if fallback is None:
        raise ValueError("No candidate rows available")
    return fallback


def triangulate_for_classifications(
    classifications: list[dict[str, object]],
    spec_shape: tuple[int, int],
    target_position: np.ndarray,
):
    az_a_base, el_a_base = az_el_from_points(DEFAULT_MODULE_A, target_position)
    az_b_base, el_b_base = az_el_from_points(DEFAULT_MODULE_B, target_position)
    az_a_offset, el_a_offset, az_b_offset, el_b_offset = feature_angle_offsets(classifications, spec_shape)
    az_a = az_a_base + az_a_offset
    el_a = el_a_base + el_a_offset
    az_b = az_b_base + az_b_offset
    el_b = el_b_base + el_b_offset
    tri = triangulate_from_az_el(DEFAULT_MODULE_A, az_a, el_a, DEFAULT_MODULE_B, az_b, el_b)
    return tri, {"az_a": az_a, "el_a": el_a, "az_b": az_b, "el_b": el_b}


def draw_spectrogram_panel(ax, item: dict[str, object], frame_idx: int, total_frames: int) -> None:
    spec = item["spec"]
    cam = item["cam"]
    boxes = item["boxes"]
    height, width = spec.shape
    extent = (0, width, 0, height)
    vmin, vmax = np.percentile(spec, [2, 98])
    ax.imshow(spec, aspect="auto", origin="lower", extent=extent, cmap="turbo", vmin=vmin, vmax=vmax)
    ax.imshow(cam, aspect="auto", origin="lower", extent=extent, cmap="cool", alpha=np.clip(cam, 0.0, 0.65) * 0.35)
    for box_idx, (x1, y1, x2, y2, score) in enumerate(boxes, start=1):
        x1 = max(0, min(width, int(x1)))
        x2 = max(0, min(width, int(x2)))
        y1 = max(0, min(height, int(y1)))
        y2 = max(0, min(height, int(y2)))
        edge = "#00f5ff" if box_idx == 1 else "#ffe66d"
        ax.add_patch(
            patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=1.6, edgecolor=edge, facecolor="none")
        )
    peak_y, peak_x = item["peak_bin"]
    ax.scatter([peak_x + 0.5], [peak_y + 0.5], marker="x", s=28, c="white", linewidths=1.4)
    ax.set_title(
        f"{item['name']} | {item['prediction']}",
        color="white",
        fontsize=10,
        pad=5,
    )
    ax.text(
        0.01,
        0.02,
        (
            f"frame {frame_idx + 1:02d}/{total_frames:02d}  latency {item['latency_ms']:.1f} ms\n"
            f"peak bin f={peak_y}, t={peak_x}  max {item['peak_power_db']:.1f} dB"
        ),
        transform=ax.transAxes,
        color="#e5e7eb",
        fontsize=7.0,
        bbox=dict(facecolor="#05070a", edgecolor="#374151", alpha=0.78, boxstyle="square,pad=0.25"),
    )
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    ax.tick_params(colors="#9aa4b2", labelsize=7)
    for spine in ax.spines.values():
        spine.set_color("#1f2937")


def draw_triangulation_panel(ax, tri) -> None:
    pane = (0.46, 0.46, 0.46, 1.0)
    grid = (0.72, 0.72, 0.72, 0.55)
    ax.set_facecolor("#05070a")
    for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
        axis.set_pane_color(pane)
        axis._axinfo["grid"]["color"] = grid
        axis._axinfo["axisline"]["color"] = (0.9, 0.9, 0.9, 1.0)
        axis._axinfo["tick"]["color"] = (0.9, 0.9, 0.9, 1.0)
    ax.scatter(*tri.module_a, s=65, c="#00f5ff", marker="^", label="Module A")
    ax.scatter(*tri.module_b, s=65, c="#ffe66d", marker="^", label="Module B")
    ax.scatter(*tri.estimated_position, s=95, c="#ff4d6d", marker="*", label="Estimated drone")
    for origin, direction, color in [
        (tri.module_a, tri.direction_a, "#00f5ff"),
        (tri.module_b, tri.direction_b, "#ffe66d"),
    ]:
        distances = np.linspace(0, 38, 2)
        points = origin[:, None] + direction[:, None] * distances
        ax.plot(points[0], points[1], points[2], color=color, linewidth=2.0)
    ax.plot(
        [tri.closest_a[0], tri.closest_b[0]],
        [tri.closest_a[1], tri.closest_b[1]],
        [tri.closest_a[2], tri.closest_b[2]],
        color="white",
        linestyle="--",
        linewidth=1.4,
    )
    ax.set_title("3D AoA triangulation", color="white", fontsize=12, pad=10)
    ax.set_xlabel("East (m)", color="white", labelpad=4)
    ax.set_ylabel("North (m)", color="white", labelpad=4)
    ax.set_zlabel("Up (m)", color="white", labelpad=4)
    ax.tick_params(colors="#d1d5db", labelsize=7)
    ax.set_xlim(-3, 22)
    ax.set_ylim(-2, 34)
    ax.set_zlim(0, 16)
    ax.view_init(elev=24, azim=-58)


def draw_triangulation_telemetry(ax, tri, angles: dict[str, float], predictions: list[str]) -> None:
    ax.set_facecolor("#05070a")
    ax.axis("off")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.text(0.02, 0.82, "3D AoA telemetry", color="white", fontsize=11, weight="bold", va="top")
    lines = [
        f"Module A az/el: {angles['az_a']:.1f} / {angles['el_a']:.1f} deg",
        f"Module B az/el: {angles['az_b']:.1f} / {angles['el_b']:.1f} deg",
        f"Estimated ENU: {tri.estimated_position[0]:.1f}, {tri.estimated_position[1]:.1f}, {tri.estimated_position[2]:.1f} m",
        f"Residual: {tri.residual_m:.2f} m",
        f"Class: {', '.join(sorted(set(predictions)))}",
    ]
    ax.text(
        0.02,
        0.56,
        "\n".join(lines),
        color="#e5e7eb",
        fontsize=8.2,
        va="top",
        linespacing=1.25,
        bbox=dict(facecolor="#111827", edgecolor="#374151", alpha=0.86, boxstyle="round,pad=0.45"),
    )


def figure_to_image(fig) -> Image.Image:
    buffer = io.BytesIO()
    fig.savefig(buffer, format="png", dpi=120, facecolor=fig.get_facecolor())
    plt.close(fig)
    buffer.seek(0)
    return Image.open(buffer).convert("RGB")


def render_sequence_frame(
    classifications: list[dict[str, object]],
    tri,
    angles: dict[str, float],
    target_label: str,
    frame_idx: int,
    total_frames: int,
) -> Image.Image:
    fig = plt.figure(figsize=(15, 8.5), facecolor="#05070a")
    gs = fig.add_gridspec(2, 3, width_ratios=[1.0, 1.0, 1.18], height_ratios=[1.0, 1.0], wspace=0.18, hspace=0.22)
    fig.suptitle(
        f"{target_label} | four RF receptor views + 3D triangulation | frame {frame_idx + 1:02d}/{total_frames:02d}",
        color="white",
        fontsize=15,
        y=0.985,
    )
    for idx, item in enumerate(classifications):
        ax = fig.add_subplot(gs[idx // 2, idx % 2])
        draw_spectrogram_panel(ax, item, frame_idx, total_frames)
        if idx // 2 == 1:
            ax.set_xlabel("Time frame", color="white", fontsize=9)
        if idx % 2 == 0:
            ax.set_ylabel("Frequency bin", color="white", fontsize=9)
    predictions = [str(item["prediction"]) for item in classifications]
    right = gs[:, 2].subgridspec(2, 1, height_ratios=[4.0, 1.15], hspace=0.08)
    tri_ax = fig.add_subplot(right[0], projection="3d")
    telemetry_ax = fig.add_subplot(right[1])
    draw_triangulation_panel(tri_ax, tri)
    draw_triangulation_telemetry(telemetry_ax, tri, angles, predictions)
    return figure_to_image(fig)


def write_sequence_gif(
    candidate_rows: list[dict[str, str]],
    frame_count: int,
    sample_dir: Path,
    model: torch.nn.Module,
    gradcam: GradCam,
    labels: list[str],
    device: torch.device,
    target_label: str,
    out_path: Path,
    fps: int,
) -> list[dict[str, object]]:
    gif_frames: list[Image.Image] = []
    frame_summaries: list[dict[str, object]] = []
    total_frames = min(frame_count, len(candidate_rows))
    for frame_idx in range(total_frames):
        if total_frames <= 1:
            start_idx = 0
        else:
            start_idx = int(round(frame_idx * (len(candidate_rows) - 1) / (total_frames - 1)))
        row, base_spec, classifications = find_agreeing_frame(
            candidate_rows,
            start_idx,
            sample_dir,
            model,
            gradcam,
            labels,
            device,
            target_label,
        )
        spec_shape = classifications[0]["spec"].shape
        target_position = target_position_for_frame(frame_idx, total_frames)
        tri, angles = triangulate_for_classifications(classifications, spec_shape, target_position)
        predictions = [str(item["prediction"]) for item in classifications]
        gif_frames.append(render_sequence_frame(classifications, tri, angles, target_label, frame_idx, total_frames))
        frame_summaries.append(
            {
                "frame": frame_idx + 1,
                "source_sample": row,
                "all_predictions_same": len(set(predictions)) == 1,
                "predictions": predictions,
                "target_position_m": target_position.tolist(),
                "estimated_position_m": tri.estimated_position.tolist(),
                "residual_m": tri.residual_m,
                "azimuth_a_deg": angles["az_a"],
                "elevation_a_deg": angles["el_a"],
                "azimuth_b_deg": angles["az_b"],
                "elevation_b_deg": angles["el_b"],
            }
        )

    out_path.parent.mkdir(parents=True, exist_ok=True)
    duration_ms = int(1000 / max(1, fps))
    gif_frames[0].save(
        out_path,
        save_all=True,
        append_images=gif_frames[1:],
        duration=duration_ms,
        loop=0,
        optimize=True,
    )
    return frame_summaries


def select_candidate(
    rows: list[dict[str, str]],
    sample_dir: Path,
    target_label: str,
    frames: int,
    frame_index: int,
) -> tuple[dict[str, str], np.ndarray, int]:
    by_label: dict[str, list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        by_label[row["label"]].append(row)
    if target_label not in by_label:
        raise ValueError(f"No rows found for target label: {target_label}")
    selected = evenly_sample_rows(by_label[target_label], min(frames, len(by_label[target_label])))
    selected_idx = max(0, min(frame_index - 1, len(selected) - 1))
    row = selected[selected_idx]
    spec = np.load(sample_dir / row["path"])["x"].astype(np.float32)
    return row, spec, selected_idx + 1


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate four-receptor RF classification and 3D triangulation artifacts.")
    parser.add_argument("--processed-dir", type=Path, default=Path("/data/RFUAV_processed"))
    parser.add_argument("--checkpoint-dir", type=Path, default=Path("/data/checkpoints"))
    parser.add_argument("--plots-dir", type=Path, default=Path("/data/plots"))
    parser.add_argument("--target-label", default="DJI MINI4 PRO")
    parser.add_argument("--frames", type=int, default=60)
    parser.add_argument("--frame-index", type=int, default=1)
    parser.add_argument("--gif-fps", type=int, default=3)
    args = parser.parse_args()

    rows = load_manifest(args.processed_dir)
    sample_dir = args.processed_dir / "samples"
    by_label: dict[str, list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        by_label[row["label"]].append(row)
    if args.target_label not in by_label:
        raise ValueError(f"No rows found for target label: {args.target_label}")
    candidate_rows = sorted(by_label[args.target_label], key=lambda row: (row["split"], row["path"]))
    selected_rows = evenly_sample_rows(candidate_rows, min(args.frames, len(candidate_rows)))
    base_row, _base_spec, actual_frame = select_candidate(
        rows,
        sample_dir,
        args.target_label,
        args.frames,
        args.frame_index,
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, labels, checkpoint_path = load_resnet18(args.processed_dir, args.checkpoint_dir, device)
    gradcam = GradCam(model, model.layer4[-1])
    out_dir = args.plots_dir / "triangulation"
    out_dir.mkdir(parents=True, exist_ok=True)

    try:
        static_frame_idx = max(0, min(args.frame_index - 1, len(selected_rows) - 1))
        static_start_idx = int(round(static_frame_idx * (len(candidate_rows) - 1) / max(1, len(selected_rows) - 1)))
        base_row, _base_spec, classifications = find_agreeing_frame(
            candidate_rows,
            static_start_idx,
            sample_dir,
            model,
            gradcam,
            labels,
            device,
            args.target_label,
        )
        for idx, receptor in enumerate(RECEPTOR_SPECS, start=1):
            result = classifications[idx - 1]
            spec = result["spec"]

            telemetry = [
                ("receptor", receptor["name"]),
                ("frame", f"{actual_frame:02d}/{len(selected_rows):02d}"),
                ("window", result["part_window"]),
                ("module", f"{receptor['module']}\n{receptor['axis']}"),
                ("prediction", str(result["prediction"])),
                ("latency", f"{result['latency_ms']:.2f} ms"),
                ("peak bin", f"f={result['peak_bin'][0]}, t={result['peak_bin'][1]}"),
                ("peak power", f"{result['peak_power_db']:.1f} dB"),
            ]
            draw_left_telemetry_spectrogram(
                spec=spec,
                cam=result["cam"],
                boxes=result["boxes"],
                out_path=out_dir / f"receptor_{idx}_spectrogram.png",
                title=f"{receptor['name']} | {args.target_label}",
                telemetry=telemetry,
            )

        target_position = target_position_for_frame(static_frame_idx, len(selected_rows))
        tri, angles = triangulate_for_classifications(classifications, classifications[0]["spec"].shape, target_position)
        predictions = [str(item["prediction"]) for item in classifications]
        draw_triangulation_3d(
            out_path=out_dir / "triangulation_3d.png",
            result=tri,
            angles=angles,
            label=args.target_label,
            predictions=predictions,
        )

        sequence_frames = write_sequence_gif(
            candidate_rows=candidate_rows,
            frame_count=len(selected_rows),
            sample_dir=sample_dir,
            model=model,
            gradcam=gradcam,
            labels=labels,
            device=device,
            target_label=args.target_label,
            out_path=out_dir / "triangulation_sequence.gif",
            fps=args.gif_fps,
        )
    finally:
        gradcam.close()

    serializable = []
    for idx, item in enumerate(classifications, start=1):
        serializable.append(
            {
                "receptor": item["name"],
                "module": item["module"],
                "axis": item["axis"],
                "time_offset": item["time_offset"],
                "source_sample": item["source_sample"],
                "prediction": item["prediction"],
                "latency_ms": item["latency_ms"],
                "peak_bin": item["peak_bin"],
                "peak_power_db": item["peak_power_db"],
                "attention_boxes": [
                    {"x1": x1, "y1": y1, "x2": x2, "y2": y2, "score": score}
                    for x1, y1, x2, y2, score in item["boxes"]
                ],
                "plot": f"triangulation/receptor_{idx}_spectrogram.png",
            }
        )

    summary = {
        "target_label": args.target_label,
        "source_sample": base_row,
        "frame": actual_frame,
        "frames": len(selected_rows),
        "model": "resnet18",
        "checkpoint": str(checkpoint_path),
        "all_predictions_same": len(set(predictions)) == 1,
        "predictions": predictions,
        "note": (
            "RFUAV is single-channel RF data. Four receptor views are circular time-offset views of the same real "
            "RFUAV spectrogram window, avoiding padded/generated edge artifacts while demonstrating a multi-receptor "
            "battlefield pipeline. 3D AoA geometry is simulated with a moving target path."
        ),
        "receptors": serializable,
        "triangulation": {
            "module_a_m": DEFAULT_MODULE_A.tolist(),
            "module_b_m": DEFAULT_MODULE_B.tolist(),
            "target_position_m": target_position.tolist(),
            "azimuth_a_deg": angles["az_a"],
            "elevation_a_deg": angles["el_a"],
            "azimuth_b_deg": angles["az_b"],
            "elevation_b_deg": angles["el_b"],
            "estimated_position_m": tri.estimated_position.tolist(),
            "closest_a_m": tri.closest_a.tolist(),
            "closest_b_m": tri.closest_b.tolist(),
            "residual_m": tri.residual_m,
            "plot": "triangulation/triangulation_3d.png",
        },
        "sequence_gif": {
            "path": "triangulation/triangulation_sequence.gif",
            "frames": len(sequence_frames),
            "fps": args.gif_fps,
            "frame_summaries": sequence_frames,
        },
        "plots": [
            "triangulation/receptor_1_spectrogram.png",
            "triangulation/receptor_2_spectrogram.png",
            "triangulation/receptor_3_spectrogram.png",
            "triangulation/receptor_4_spectrogram.png",
            "triangulation/triangulation_3d.png",
            "triangulation/triangulation_sequence.gif",
        ],
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

    print(f"Wrote triangulation artifacts to {out_dir}")
    for item in summary["plots"]:
        print(item)
    print(out_dir / "summary.json")


if __name__ == "__main__":
    main()