scripts/plot_artifacts.py

from __future__ import annotations

import argparse
import csv
import io
import json
import shutil
import time
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from PIL import Image
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
from scipy import ndimage
import torch
import torch.nn.functional as F

from train import build_model, read_labels


CLASS_ORDER = ["DAUTEL EVO nano", "DJI MAVIC3 PRO", "DJI MINI3", "DJI MINI4 PRO"]
PLOT_3D_BG = "#050505"
PLOT_3D_FG = "#e8e8e8"
PLOT_3D_PANE = (0.46, 0.46, 0.46, 1)
PLOT_3D_GRID = (0.72, 0.72, 0.72, 0.65)


def load_manifest(processed_dir: Path) -> list[dict[str, str]]:
    manifest = processed_dir / "manifest.csv"
    if not manifest.exists():
        raise FileNotFoundError(f"Missing processed manifest: {manifest}")
    with manifest.open() as f:
        return list(csv.DictReader(f))


def plot_one_spectrogram(npz_path: Path, title: str, out_path: Path) -> None:
    data = np.load(npz_path)
    x = data["x"].astype(np.float32)
    fig, ax = plt.subplots(figsize=(7, 4.5))
    im = ax.imshow(x, aspect="auto", origin="lower", cmap="viridis")
    ax.set_title(title)
    ax.set_xlabel("Time frame")
    ax.set_ylabel("Frequency bin")
    fig.colorbar(im, ax=ax, label="Log magnitude")
    fig.tight_layout()
    fig.savefig(out_path, dpi=160)
    plt.close(fig)


def normalized_surface(
    x: np.ndarray,
    max_freq_bins: int = 160,
    max_time_frames: int = 120,
    front_crop_fraction: float = 0.12,
):
    front_crop = int(x.shape[0] * front_crop_fraction)
    if front_crop > 0 and x.shape[0] > front_crop:
        x = x[front_crop:, :]
    freq_step = max(1, int(np.ceil(x.shape[0] / max_freq_bins)))
    time_step = max(1, int(np.ceil(x.shape[1] / max_time_frames)))
    z = x[::freq_step, ::time_step].astype(np.float32)
    lo, hi = np.percentile(z, [2, 98])
    z = np.clip(z, lo, hi)
    z = (z - z.min()) / (z.max() - z.min() + 1e-6)
    t = np.arange(z.shape[1])
    f = np.arange(z.shape[0])
    tt, ff = np.meshgrid(t, f)
    return tt, ff, z


def plot_one_spectrogram_3d(npz_path: Path, title: str, out_path: Path, cmap: str = "turbo") -> None:
    data = np.load(npz_path)
    x = data["x"].astype(np.float32)
    tt, ff, z = normalized_surface(x)

    fig = plt.figure(figsize=(8, 5), facecolor=PLOT_3D_BG)
    ax = fig.add_subplot(111, projection="3d", facecolor=PLOT_3D_BG)
    ax.plot_surface(
        tt,
        ff,
        z,
        cmap=cmap,
        linewidth=0,
        antialiased=True,
        shade=True,
        rstride=1,
        cstride=1,
    )
    ax.view_init(elev=32, azim=-58)
    ax.set_box_aspect((1.65, 1.0, 0.45))
    ax.set_title(title, color=PLOT_3D_FG, pad=12)
    ax.set_xlabel("Time", color=PLOT_3D_FG)
    ax.set_ylabel("Frequency", color=PLOT_3D_FG)
    ax.set_zlabel("Power", color=PLOT_3D_FG)
    ax.tick_params(colors=PLOT_3D_FG)
    ax.xaxis.pane.set_facecolor(PLOT_3D_PANE)
    ax.yaxis.pane.set_facecolor(PLOT_3D_PANE)
    ax.zaxis.pane.set_facecolor(PLOT_3D_PANE)
    for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
        axis._axinfo["grid"]["color"] = PLOT_3D_GRID
        axis._axinfo["grid"]["linewidth"] = 0.8
    ax.grid(True)
    fig.tight_layout()
    fig.savefig(out_path, dpi=180, facecolor=fig.get_facecolor())
    plt.close(fig)


def style_3d_axis(ax, title: str) -> None:
    ax.set_facecolor(PLOT_3D_BG)
    ax.view_init(elev=32, azim=-58)
    ax.set_box_aspect((1.65, 1.0, 0.45))
    ax.set_title(title, color=PLOT_3D_FG, pad=10, fontsize=12)
    ax.set_xlabel("Time", color=PLOT_3D_FG, labelpad=6)
    ax.set_ylabel("Frequency", color=PLOT_3D_FG, labelpad=6)
    ax.set_zlabel("Power", color=PLOT_3D_FG, labelpad=6)
    ax.tick_params(colors=PLOT_3D_FG, labelsize=7)
    ax.xaxis.pane.set_facecolor(PLOT_3D_PANE)
    ax.yaxis.pane.set_facecolor(PLOT_3D_PANE)
    ax.zaxis.pane.set_facecolor(PLOT_3D_PANE)
    for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
        axis._axinfo["grid"]["color"] = PLOT_3D_GRID
        axis._axinfo["grid"]["linewidth"] = 0.8
    ax.grid(True)


def render_3d_spectrogram_frame(x: np.ndarray, title: str) -> Image.Image:
    tt, ff, z = normalized_surface(x)
    fig = plt.figure(figsize=(8, 5), facecolor=PLOT_3D_BG)
    ax = fig.add_subplot(111, projection="3d")
    ax.plot_surface(
        tt,
        ff,
        z,
        cmap="turbo",
        linewidth=0,
        antialiased=True,
        shade=True,
        rstride=1,
        cstride=1,
    )
    style_3d_axis(ax, title)
    fig.tight_layout()

    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 plot_spectrogram_previews(
    processed_dir: Path,
    plots_dir: Path,
    max_per_class: int,
    max_3d_per_class: int,
) -> list[dict[str, str]]:
    rows = load_manifest(processed_dir)
    sample_dir = processed_dir / "samples"
    out_dir = plots_dir / "spectrograms"
    out_3d_dir = plots_dir / "spectrograms_3d"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_3d_dir.mkdir(parents=True, exist_ok=True)

    by_label: dict[str, list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        by_label[row["label"]].append(row)

    summary = []
    grid_items = []
    for label, label_rows in sorted(by_label.items()):
        selected = label_rows[:max_per_class]
        safe_label = label.replace(" ", "_").replace("/", "_")
        for idx, row in enumerate(selected):
            npz_path = sample_dir / row["path"]
            out_path = out_dir / f"{safe_label}_{idx}.png"
            plot_one_spectrogram(npz_path, f"{label} ({row['split']})", out_path)
            item = {"label": label, "split": row["split"], "plot": str(out_path)}
            if idx < max_3d_per_class:
                out_3d_path = out_3d_dir / f"{safe_label}_{idx}_3d.png"
                plot_one_spectrogram_3d(npz_path, f"{label} ({row['split']})", out_3d_path)
                item["plot_3d"] = str(out_3d_path)
            summary.append(item)
            if idx == 0:
                grid_items.append((label, np.load(npz_path)["x"].astype(np.float32)))

    if grid_items:
        fig, axes = plt.subplots(1, len(grid_items), figsize=(5 * len(grid_items), 4), squeeze=False)
        for ax, (label, x) in zip(axes[0], grid_items):
            ax.imshow(x, aspect="auto", origin="lower", cmap="viridis")
            ax.set_title(label)
            ax.set_xlabel("Time")
            ax.set_ylabel("Frequency")
        fig.tight_layout()
        fig.savefig(out_dir / "spectrogram_grid.png", dpi=160)
        plt.close(fig)

    return summary


def plot_resnet18_spectrogram_previews(
    processed_dir: Path,
    plots_dir: Path,
    max_per_class: int,
    max_3d_per_class: int,
) -> list[dict[str, str]]:
    rows = load_manifest(processed_dir)
    sample_dir = processed_dir / "samples"
    out_root = plots_dir / "spectrograms_resnet18"
    out_2d_dir = out_root / "2d"
    out_3d_dir = out_root / "3d"
    out_2d_dir.mkdir(parents=True, exist_ok=True)
    out_3d_dir.mkdir(parents=True, exist_ok=True)
    for stale_plot in out_3d_dir.glob("*.png"):
        stale_plot.unlink()

    by_label: dict[str, list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        by_label[row["label"]].append(row)

    summary = []
    for label, label_rows in sorted(by_label.items()):
        selected = label_rows[:max_per_class]
        safe_label = label.replace(" ", "_").replace("/", "_")
        for idx, row in enumerate(selected):
            npz_path = sample_dir / row["path"]
            out_path = out_2d_dir / f"{safe_label}_{idx}.png"
            title = f"ResNet18 input: {label} ({row['split']})"
            plot_one_spectrogram(npz_path, title, out_path)
            item = {"label": label, "split": row["split"], "plot": str(out_path)}
            if idx < max_3d_per_class:
                out_3d_path = out_3d_dir / f"{safe_label}_{idx}_3d.png"
                plot_one_spectrogram_3d(npz_path, title, out_3d_path)
                item["plot_3d"] = str(out_3d_path)
            summary.append(item)

    return summary


def plot_resnet18_3d_grid(processed_dir: Path, plots_dir: Path) -> dict[str, object]:
    rows = load_manifest(processed_dir)
    sample_dir = processed_dir / "samples"
    out_root = plots_dir / "spectrograms_resnet18"
    out_grid_dir = out_root / "grid"
    out_grid_dir.mkdir(parents=True, exist_ok=True)

    by_label: dict[str, list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        by_label[row["label"]].append(row)

    selected = []
    for label in CLASS_ORDER:
        label_rows = sorted(by_label.get(label, []), key=lambda row: (row["split"], row["path"]))
        if label_rows:
            selected.append((label, label_rows[0]))

    if len(selected) < 4:
        return {"created": False, "reason": "Need one sample for each selected class"}

    out_path = out_grid_dir / "class_comparison_3d_grid.png"
    fig = plt.figure(figsize=(14, 9), facecolor=PLOT_3D_BG)
    fig.suptitle("RFUAV 3D Spectrogram Class Comparison", color=PLOT_3D_FG, fontsize=18)

    for idx, (label, row) in enumerate(selected[:4], start=1):
        npz_path = sample_dir / row["path"]
        data = np.load(npz_path)
        x = data["x"].astype(np.float32)
        tt, ff, z = normalized_surface(x)
        ax = fig.add_subplot(2, 2, idx, projection="3d")
        ax.plot_surface(
            tt,
            ff,
            z,
            cmap="turbo",
            linewidth=0,
            antialiased=True,
            shade=True,
            rstride=1,
            cstride=1,
        )
        style_3d_axis(ax, label)

    fig.tight_layout(rect=(0, 0, 1, 0.95))
    fig.savefig(out_path, dpi=180, facecolor=fig.get_facecolor())
    plt.close(fig)

    return {"created": True, "path": str(out_path), "labels": [label for label, _ in selected[:4]]}


def plot_mavic3_3d_spectrogram_gif(
    processed_dir: Path,
    plots_dir: Path,
    frame_count: int,
    fps: int,
) -> dict[str, object]:
    rows = load_manifest(processed_dir)
    sample_dir = processed_dir / "samples"
    out_dir = plots_dir / "gifs"
    out_dir.mkdir(parents=True, exist_ok=True)

    mavic_rows = sorted(
        [row for row in rows if row["label"] == "DJI MAVIC3 PRO"],
        key=lambda row: (row["split"], row["path"]),
    )
    if not mavic_rows:
        return {"created": False, "reason": "No DJI MAVIC3 PRO samples available"}

    selected_rows = evenly_sample_rows(mavic_rows, min(frame_count, len(mavic_rows)))
    gif_frames = []
    for idx, row in enumerate(selected_rows, start=1):
        x = np.load(sample_dir / row["path"])["x"].astype(np.float32)
        gif_frames.append(render_3d_spectrogram_frame(x, f"DJI MAVIC3 PRO | frame {idx}/{len(selected_rows)}"))

    out_path = out_dir / "dji_mavic3_pro_3d_spectrogram.gif"
    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 {
        "created": True,
        "path": str(out_path),
        "frames": len(gif_frames),
        "fps": fps,
        "label": "DJI MAVIC3 PRO",
    }


class GradCam:
    def __init__(self, model: torch.nn.Module, target_layer: torch.nn.Module):
        self.model = model
        self.activations: torch.Tensor | None = None
        self.gradients: torch.Tensor | None = None
        self.forward_handle = target_layer.register_forward_hook(self._save_activation)
        self.backward_handle = target_layer.register_full_backward_hook(self._save_gradient)

    def _save_activation(self, _module, _inputs, output) -> None:
        self.activations = output

    def _save_gradient(self, _module, _grad_input, grad_output) -> None:
        self.gradients = grad_output[0]

    def close(self) -> None:
        self.forward_handle.remove()
        self.backward_handle.remove()

    def __call__(self, logits: torch.Tensor, class_idx: int, out_shape: tuple[int, int]) -> np.ndarray:
        self.model.zero_grad(set_to_none=True)
        score = logits[0, class_idx]
        score.backward(retain_graph=True)
        if self.activations is None or self.gradients is None:
            raise RuntimeError("Grad-CAM hooks did not capture activations/gradients")

        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
        cam = (weights * self.activations).sum(dim=1, keepdim=True)
        cam = torch.relu(cam)
        cam = F.interpolate(cam, size=out_shape, mode="bilinear", align_corners=False)
        cam = cam[0, 0].detach().cpu().numpy()
        cam = cam - cam.min()
        cam = cam / (cam.max() + 1e-8)
        return cam


def load_resnet18_for_gradcam(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 attention_boxes(cam: np.ndarray, threshold: float = 0.72, min_area: int = 18, max_boxes: int = 5) -> list[tuple[int, int, int, int, float]]:
    mask = cam >= threshold
    labeled, count = ndimage.label(mask)
    regions = []
    for idx, region_slice in enumerate(ndimage.find_objects(labeled), start=1):
        if region_slice is None:
            continue
        y_slice, x_slice = region_slice
        area = int((labeled[region_slice] == idx).sum())
        if area < min_area:
            continue
        score = float(cam[region_slice][labeled[region_slice] == idx].mean())
        regions.append((x_slice.start, y_slice.start, x_slice.stop, y_slice.stop, score, area))

    if not regions and count:
        y, x = np.unravel_index(int(np.argmax(cam)), cam.shape)
        pad_y = max(4, cam.shape[0] // 18)
        pad_x = max(3, cam.shape[1] // 18)
        regions.append((
            max(0, x - pad_x),
            max(0, y - pad_y),
            min(cam.shape[1], x + pad_x),
            min(cam.shape[0], y + pad_y),
            float(cam[y, x]),
            1,
        ))

    regions = sorted(regions, key=lambda item: (item[4], item[5]), reverse=True)[:max_boxes]
    return [(x1, y1, x2, y2, score) for x1, y1, x2, y2, score, _area in regions]


def render_gradcam_frame(
    spec: np.ndarray,
    cam: np.ndarray,
    boxes: list[tuple[int, int, int, int, float]],
    frame_idx: int,
    total_frames: int,
    true_label: str,
    pred_label: str,
    latency_ms: float,
) -> Image.Image:
    peak_y, peak_x = np.unravel_index(int(np.argmax(spec)), spec.shape)
    fig = plt.figure(figsize=(12, 6), facecolor="#05070a")
    gs = fig.add_gridspec(1, 2, width_ratios=[3.2, 1.25], wspace=0.08)
    ax = fig.add_subplot(gs[0, 0])
    side = fig.add_subplot(gs[0, 1])

    ax.set_facecolor("#05070a")
    vmin, vmax = np.percentile(spec, [2, 98])
    height, width = spec.shape
    image_extent = (0, width, 0, height)
    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.52,
    )

    for box_idx, (x1, y1, x2, y2, score) in enumerate(boxes, start=1):
        x1 = max(0, min(width, x1))
        x2 = max(0, min(width, x2))
        y1 = max(0, min(height, y1))
        y2 = max(0, min(height, y2))
        rect = patches.Rectangle(
            (x1, y1),
            x2 - x1,
            y2 - y1,
            linewidth=2.0,
            edgecolor="#00f5ff" if box_idx == 1 else "#ffe66d",
            facecolor="none",
        )
        ax.add_patch(rect)
        ax.text(
            x1,
            min(height - 2, y2 + 2),
            f"A{box_idx} {score:.2f}",
            color="#05070a",
            fontsize=8,
            fontweight="bold",
            bbox={"facecolor": "#00f5ff" if box_idx == 1 else "#ffe66d", "alpha": 0.9, "pad": 1.5},
        )

    ax.scatter([peak_x + 0.5], [peak_y + 0.5], marker="x", s=50, linewidths=1.5, color="white")
    ax.set_title("DJI MINI4 PRO | RF spectrogram + ResNet18 Grad-CAM", color="white", fontsize=14, pad=10)
    ax.set_xlabel("Time frame", color="#d7e3f4")
    ax.set_ylabel("Frequency bin", color="#d7e3f4")
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    ax.tick_params(colors="#aeb9c8")

    side.set_facecolor("#0c1119")
    side.axis("off")
    panel_lines = [
        ("FRAME", f"{frame_idx:02d} / {total_frames:02d}"),
        ("TRUE CLASS", true_label),
        ("PREDICTION", pred_label),
        ("LATENCY", f"{latency_ms:.2f} ms"),
        ("PEAK BIN", f"f={peak_y}, t={peak_x}"),
        ("PEAK POWER", f"{float(spec[peak_y, peak_x]):.1f} dB"),
        ("ATTN BOXES", str(len(boxes))),
    ]
    side.text(0.04, 0.94, "CLASSIFIER TELEMETRY", color="white", fontsize=12, fontweight="bold", transform=side.transAxes)
    y = 0.84
    for label, value in panel_lines:
        side.text(0.05, y, label, color="#6f7f93", fontsize=8, fontweight="bold", transform=side.transAxes)
        side.text(0.05, y - 0.045, value, color="white", fontsize=12, transform=side.transAxes, wrap=True)
        y -= 0.105

    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 plot_mini4_gradcam_gif(
    processed_dir: Path,
    results_dir: Path,
    checkpoint_dir: Path,
    plots_dir: Path,
    frame_count: int,
    fps: int,
) -> dict[str, object]:
    rows = load_manifest(processed_dir)
    sample_dir = processed_dir / "samples"
    out_dir = plots_dir / "gifs"
    out_dir.mkdir(parents=True, exist_ok=True)

    mini4_rows = sorted(
        [row for row in rows if row["label"] == "DJI MINI4 PRO"],
        key=lambda row: (row["split"], row["path"]),
    )
    if not mini4_rows:
        return {"created": False, "reason": "No DJI MINI4 PRO samples available"}

    metrics_path = results_dir / "resnet18" / "metrics.json"
    test_accuracy = None
    if metrics_path.exists():
        metrics = json.loads(metrics_path.read_text())
        test_accuracy = metrics.get("test_accuracy")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, labels, checkpoint_path = load_resnet18_for_gradcam(processed_dir, checkpoint_dir, device)
    gradcam = GradCam(model, model.layer4[-1])
    selected_rows = evenly_sample_rows(mini4_rows, min(frame_count, len(mini4_rows)))
    gif_frames: list[Image.Image] = []

    try:
        for idx, row in enumerate(selected_rows, start=1):
            spec = np.load(sample_dir / row["path"])["x"].astype(np.float32)
            normalized = (spec - spec.mean()) / (spec.std() + 1e-6)
            x = torch.from_numpy(normalized).unsqueeze(0).unsqueeze(0).to(device)

            start = time.perf_counter()
            logits = model(x)
            if device.type == "cuda":
                torch.cuda.synchronize()
            latency_ms = (time.perf_counter() - start) * 1000.0

            probs = torch.softmax(logits.detach(), dim=1)[0]
            pred_id = int(torch.argmax(probs).item())
            cam = gradcam(logits, pred_id, spec.shape)
            boxes = attention_boxes(cam)
            frame = render_gradcam_frame(
                spec=spec,
                cam=cam,
                boxes=boxes,
                frame_idx=idx,
                total_frames=len(selected_rows),
                true_label=row["label"],
                pred_label=labels[pred_id],
                latency_ms=latency_ms,
            )
            gif_frames.append(frame)
    finally:
        gradcam.close()

    out_path = out_dir / "dji_mini4_pro_gradcam_boxes.gif"
    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 {
        "created": True,
        "path": str(out_path),
        "frames": len(gif_frames),
        "fps": fps,
        "label": "DJI MINI4 PRO",
        "model": "resnet18",
        "checkpoint": str(checkpoint_path),
    }


def evenly_sample_rows(rows: list[dict[str, str]], count: int) -> list[dict[str, str]]:
    rows = sorted(rows, key=lambda row: (row["split"], row["path"]))
    if len(rows) <= count:
        return rows
    indices = np.linspace(0, len(rows) - 1, count, dtype=int)
    return [rows[int(idx)] for idx in indices]


def plot_time_sweep_comparison_gif(
    processed_dir: Path,
    plots_dir: Path,
    frame_count: int,
    fps: int,
) -> dict[str, object]:
    rows = load_manifest(processed_dir)
    sample_dir = processed_dir / "samples"
    out_dir = plots_dir / "gifs"
    out_dir.mkdir(parents=True, exist_ok=True)

    by_label: dict[str, list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        by_label[row["label"]].append(row)

    selected_labels = [label for label in CLASS_ORDER if by_label.get(label)]
    if len(selected_labels) < 2:
        return {"created": False, "reason": "Need at least two classes for comparison GIF"}

    frames_available = min(len(by_label[label]) for label in selected_labels)
    frame_count = min(frame_count, frames_available)
    if frame_count <= 0:
        return {"created": False, "reason": "No spectrogram samples available"}

    selected_rows = {
        label: evenly_sample_rows(by_label[label], frame_count)
        for label in selected_labels
    }

    spectrograms: dict[str, list[np.ndarray]] = {}
    all_values = []
    for label, label_rows in selected_rows.items():
        specs = []
        for row in label_rows:
            x = np.load(sample_dir / row["path"])["x"].astype(np.float32)
            specs.append(x)
            all_values.append(x.reshape(-1))
        spectrograms[label] = specs

    values = np.concatenate(all_values)
    vmin, vmax = np.percentile(values, [2, 98])
    gif_frames: list[Image.Image] = []

    for frame_idx in range(frame_count):
        fig, axes = plt.subplots(2, 2, figsize=(10, 7), facecolor="black", squeeze=False)
        fig.suptitle(
            f"RFUAV spectrogram time-sweep comparison | frame {frame_idx + 1}/{frame_count}",
            color="white",
            fontsize=15,
        )

        for ax in axes.ravel():
            ax.set_facecolor("black")
            ax.axis("off")

        for ax, label in zip(axes.ravel(), selected_labels):
            row = selected_rows[label][frame_idx]
            ax.imshow(
                spectrograms[label][frame_idx],
                aspect="auto",
                origin="lower",
                cmap="turbo",
                vmin=vmin,
                vmax=vmax,
            )
            ax.set_title(label, color="white", fontsize=12)
            ax.set_xlabel("Time frame", color="white")
            ax.set_ylabel("Frequency bin", color="white")
            ax.tick_params(colors="white", labelsize=8)
            ax.axis("on")

        fig.tight_layout(rect=(0, 0, 1, 0.93))

        buffer = io.BytesIO()
        fig.savefig(buffer, format="png", dpi=120, facecolor=fig.get_facecolor())
        plt.close(fig)
        buffer.seek(0)
        gif_frames.append(Image.open(buffer).convert("RGB"))

    out_path = out_dir / "time_sweep_comparison.gif"
    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 {
        "created": True,
        "path": str(out_path),
        "frames": frame_count,
        "fps": fps,
        "labels": selected_labels,
    }


def find_metric_files(results_dir: Path) -> list[Path]:
    if not results_dir.exists():
        return []
    return sorted(results_dir.rglob("metrics.json"))


def plot_training_curves(metric_path: Path, plots_dir: Path) -> dict[str, object]:
    metrics = json.loads(metric_path.read_text())
    model_name = str(metrics.get("model") or metric_path.parent.name)
    safe_name = model_name.replace("/", "_")
    model_dir = plots_dir / safe_name
    history = metrics.get("history", [])
    model_dir.mkdir(parents=True, exist_ok=True)

    if history:
        epochs = [row["epoch"] for row in history]
        fig, axes = plt.subplots(1, 2, figsize=(11, 4))
        axes[0].plot(epochs, [row["train_loss"] for row in history], label="train_loss")
        axes[0].plot(epochs, [row["val_loss"] for row in history], label="val_loss")
        axes[0].set_title(f"{model_name}: loss")
        axes[0].set_xlabel("Epoch")
        axes[0].legend()

        axes[1].plot(epochs, [row["val_accuracy"] for row in history], label="val_accuracy")
        axes[1].set_title(f"{model_name}: validation accuracy")
        axes[1].set_xlabel("Epoch")
        axes[1].set_ylim(0, 1)
        axes[1].legend()
        fig.tight_layout()
        fig.savefig(model_dir / "training_curves.png", dpi=160)
        plt.close(fig)

    confusion = metric_path.parent / "confusion_matrix.png"
    if confusion.exists():
        shutil.copy2(confusion, model_dir / "confusion_matrix.png")

    shutil.copy2(metric_path, model_dir / "metrics.json")

    return {
        "model": model_name,
        "best_val_accuracy": metrics.get("best_val_accuracy"),
        "test_accuracy": metrics.get("test_accuracy"),
        "metrics_path": str(metric_path),
        "plot_dir": str(model_dir),
    }


def plot_metric_artifacts(results_dir: Path, plots_dir: Path) -> list[dict[str, object]]:
    summaries = [plot_training_curves(path, plots_dir) for path in find_metric_files(results_dir)]

    if summaries:
        out_dir = plots_dir / "comparison"
        out_dir.mkdir(parents=True, exist_ok=True)
        names = [str(item["model"]) for item in summaries]
        test_acc = [float(item.get("test_accuracy") or 0.0) for item in summaries]
        val_acc = [float(item.get("best_val_accuracy") or 0.0) for item in summaries]
        x = np.arange(len(names))
        width = 0.36
        fig, ax = plt.subplots(figsize=(max(7, 2.5 * len(names)), 4.5))
        ax.bar(x - width / 2, val_acc, width, label="best val acc")
        ax.bar(x + width / 2, test_acc, width, label="test acc")
        ax.set_xticks(x, names, rotation=20, ha="right")
        ax.set_ylim(0, 1)
        ax.set_ylabel("Accuracy")
        ax.set_title("Model accuracy comparison")
        ax.legend()
        fig.tight_layout()
        fig.savefig(out_dir / "comparison_accuracy.png", dpi=160)
        plt.close(fig)

    return summaries


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate plot artifacts for RFUAV spectrograms and CNN metrics.")
    parser.add_argument("--processed-dir", default="/data/RFUAV_processed")
    parser.add_argument("--results-dir", default="/data/results")
    parser.add_argument("--checkpoint-dir", default="/data/checkpoints")
    parser.add_argument("--plots-dir", default="/data/plots")
    parser.add_argument("--max-spectrograms-per-class", type=int, default=3)
    parser.add_argument("--max-3d-spectrograms-per-class", type=int, default=1)
    parser.add_argument("--gif-frames", type=int, default=30)
    parser.add_argument("--gif-fps", type=int, default=2)
    parser.add_argument("--mavic3-3d-gif-frames", type=int, default=20)
    parser.add_argument("--mavic3-3d-gif-fps", type=int, default=2)
    parser.add_argument("--mini4-gradcam-gif-frames", type=int, default=30)
    parser.add_argument("--mini4-gradcam-gif-fps", type=int, default=2)
    args = parser.parse_args()

    processed_dir = Path(args.processed_dir)
    results_dir = Path(args.results_dir)
    checkpoint_dir = Path(args.checkpoint_dir)
    plots_dir = Path(args.plots_dir)
    plots_dir.mkdir(parents=True, exist_ok=True)

    summary = {
        "spectrograms": [],
        "spectrograms_resnet18": [],
        "gifs": {},
        "mavic3_3d_gif": {},
        "mini4_gradcam_gif": {},
        "metrics": [],
    }

    if processed_dir.exists() and (processed_dir / "manifest.csv").exists():
        summary["spectrograms"] = plot_spectrogram_previews(
            processed_dir,
            plots_dir,
            args.max_spectrograms_per_class,
            args.max_3d_spectrograms_per_class,
        )
        summary["spectrograms_resnet18"] = plot_resnet18_spectrogram_previews(
            processed_dir,
            plots_dir,
            args.max_spectrograms_per_class,
            args.max_3d_spectrograms_per_class,
        )
        summary["spectrograms_resnet18_3d_grid"] = plot_resnet18_3d_grid(processed_dir, plots_dir)
        summary["gifs"] = plot_time_sweep_comparison_gif(
            processed_dir,
            plots_dir,
            args.gif_frames,
            args.gif_fps,
        )
        summary["mavic3_3d_gif"] = plot_mavic3_3d_spectrogram_gif(
            processed_dir,
            plots_dir,
            args.mavic3_3d_gif_frames,
            args.mavic3_3d_gif_fps,
        )
        summary["mini4_gradcam_gif"] = plot_mini4_gradcam_gif(
            processed_dir,
            results_dir,
            checkpoint_dir,
            plots_dir,
            args.mini4_gradcam_gif_frames,
            args.mini4_gradcam_gif_fps,
        )
    else:
        print(f"Skipping spectrogram plots; missing {processed_dir / 'manifest.csv'}")

    summary["metrics"] = plot_metric_artifacts(results_dir, plots_dir)
    (plots_dir / "summary.json").write_text(json.dumps(summary, indent=2))

    print(f"Generated plots under {plots_dir}")
    print(
        json.dumps(
            {
                "spectrogram_plots": len(summary["spectrograms"]),
                "resnet18_spectrogram_plots": len(summary["spectrograms_resnet18"]),
                "gif": summary["gifs"],
                "mavic3_3d_gif": summary["mavic3_3d_gif"],
                "mini4_gradcam_gif": summary["mini4_gradcam_gif"],
                "metric_runs": len(summary["metrics"]),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()