"""Render per-run weak target and single, pooled, held-out GNN predictions."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

from models.gnn_joint_experiment import PROJECT, RUNS, combine, forward, load_samples, make_model
from models.compare_cad_predictions import metrics


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--report", type=Path, required=True)
    parser.add_argument("--kind", default="residual_geo")
    parser.add_argument("--budget", type=int, default=12000)
    args = parser.parse_args()
    report_path = args.report.resolve()
    report = json.loads(report_path.read_text())
    out = report_path.parent
    samples = {s.name: s for s in load_samples(args.budget)}
    fig, axes = plt.subplots(3, 4, figsize=(18, 9), constrained_layout=True)
    saved = {}
    for i, (name, part, root, run) in enumerate(RUNS):
        sample = samples[name]
        with np.load(root / f"data_processed/paired/{run}_cad_ply_targets.npz", allow_pickle=False) as d:
            pos = np.asarray(d["cad_pos_mm"], np.float32)
        axes[i, 0].scatter(pos[:, 0], pos[:, 2], s=1, c="#cbd3da", rasterized=True)
        axes[i, 0].scatter(pos[sample.fine_binary, 0], pos[sample.fine_binary, 2],
                           s=16, c="#d92332", rasterized=True)
        axes[i, 0].set_title(f"{name}: {sample.fine_binary.sum()} weak positives")
        for j, setup in enumerate((f"single_{name}", "joint", f"holdout_{name}"), 1):
            row = next(r for r in report["rows"] if r["kind"] == args.kind
                       and r["budget"] == args.budget and r["setup"] == setup
                       and r["eval_run"] == name)
            path = PROJECT / row["checkpoint"]
            checkpoint = torch.load(path, map_location="cpu", weights_only=True)
            model = make_model(args.kind, sample)
            model.load_state_dict(checkpoint["state_dict"])
            model.eval()
            with torch.no_grad():
                prob = forward(model, combine([sample])).sigmoid().cpu().numpy()
            fine = prob[sample.fine_from_coarse]
            saved[f"{name}_{setup}"] = fine.astype(np.float32)
            axes[i, j].scatter(pos[:, 0], pos[:, 2], s=2.5, c=fine, cmap="magma",
                               vmin=0, vmax=1, rasterized=True)
            axes[i, j].scatter(pos[sample.fine_binary, 0], pos[sample.fine_binary, 2],
                               s=10, facecolors="none", edgecolors="#36d9e7", linewidths=.5,
                               rasterized=True)
            result = metrics(fine, sample.fine_binary)
            axes[i, j].set_title(f"{setup.split('_')[0]}: AP {result['average_precision']:.3f}, "
                                 f"IoU {result['threshold_iou']:.3f}")
        for ax in axes[i]:
            ax.set_aspect("equal"); ax.set_xlabel("CAD X (mm)"); ax.set_ylabel("CAD Z (mm)")
            ax.grid(alpha=.1)
    fig.suptitle("Weak-label diagnostics only: same-run fits and held-out-run predictions", fontsize=16)
    fig.savefig(out / "prediction_maps.png", dpi=170)
    plt.close(fig)
    np.savez_compressed(out / "prediction_maps.npz", **saved)
    print(out / "prediction_maps.png")


if __name__ == "__main__": main()