#!/usr/bin/env python3
"""Compare saved CAD-mapped weak-risk fits across AC28 and ME1a models."""
from __future__ import annotations

import csv
import html
import json
from pathlib import Path

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

PROJECT = Path(__file__).resolve().parents[1]
OUT = PROJECT / "results/ME1a_DT_comparison/model_ranking"
BASE = PROJECT / "results/ME1a_DT_comparison"


def load(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def residual_row(rows: list[dict], run: str, setup: str) -> dict:
    return next(row for row in rows if row["budget"] == 12000 and row["kind"] == "residual_geo"
                and row["eval_run"] == run and row["setup"] == setup)


def records() -> list[dict]:
    ac = load(PROJECT / "results/models/cad_predictions/report.json")
    me = load(PROJECT / "experiments/ME1a_DT/results/models/cad_predictions/report.json")
    hard = load(BASE / "gnn_joint_hard/report.json")
    if ac["run"] != "18h01" or me["run"] != "14h06":
        raise ValueError("Unexpected comparison run")
    rows = []
    for name, a in ac["models"].items():
        m = me["models"][name]
        rows.append({"model": name, "family": "prior baseline", "graph_nodes_requested": "",
                     "ac28_ap": a["average_precision"], "ac28_iou": a["threshold_iou"],
                     "ac28_dice": a["threshold_dice"], "ac28_coverage": a["display_coverage"],
                     "ac28_epochs": a["epochs"],
                     "ac28_training_runs": "AC28 16h + 18h01" if name == "Risk GNN" else "AC28 18h01",
                     "me1a_ap": m["average_precision"], "me1a_iou": m["threshold_iou"],
                     "me1a_dice": m["threshold_dice"], "me1a_coverage": m["display_coverage"],
                     "me1a_epochs": m["epochs"], "me1a_training_runs": "ME1a 14h06",
                     "ac28_checkpoint": a["checkpoint"],
                     "me1a_checkpoint": "experiments/ME1a_DT/" + m["checkpoint"]})
    for title, setup in (("Residual Geo GNN (individual)", "single"),
                         ("Residual Geo GNN (joint)", "joint")):
        a = residual_row(hard["rows"], "AC28_18h01", "single_AC28_18h01" if setup == "single" else "joint")
        m = residual_row(hard["rows"], "ME1a_14h06", "single_ME1a_14h06" if setup == "single" else "joint")
        rows.append({"model": title, "family": "optimized risk-only GNN", "graph_nodes_requested": 12000,
                     "ac28_ap": a["mapped_ap"], "ac28_iou": a["mapped_iou"],
                     "ac28_dice": a["mapped_dice"], "ac28_coverage": 1.0,
                     "ac28_epochs": hard["epochs_per_fit"], "ac28_training_runs": a["train_runs"],
                     "me1a_ap": m["mapped_ap"], "me1a_iou": m["mapped_iou"],
                     "me1a_dice": m["mapped_dice"], "me1a_coverage": 1.0,
                     "me1a_epochs": hard["epochs_per_fit"], "me1a_training_runs": m["train_runs"],
                     "ac28_checkpoint": a["checkpoint"], "me1a_checkpoint": m["checkpoint"]})
    return rows


def figure(rows: list[dict]) -> None:
    fig, axes = plt.subplots(2, 2, figsize=(15, 11), constrained_layout=True)
    for i, (part, label) in enumerate((('ac28', 'AC28 18h01'), ('me1a', 'ME1a 14h06'))):
        for j, (metric, title) in enumerate((('ap', 'Average precision'), ('iou', 'IoU at 0.5'))):
            key = f"{part}_{metric}"
            ordered = sorted(rows, key=lambda row: row[key])
            colors = [('#cf6d28' if row['model'] == 'Risk GNN' else
                       '#126c9a' if row['family'].startswith('optimized') else '#8998aa') for row in ordered]
            ax = axes[i, j]
            y = np.arange(len(ordered))
            ax.barh(y, [row[key] for row in ordered], color=colors)
            ax.set_yticks(y, [row['model'] for row in ordered], fontsize=8)
            ax.set_xlim(0, max(.05, max(row[key] for row in ordered) * 1.14))
            ax.set_xlabel(metric.upper() if metric == 'ap' else 'IoU')
            ax.set_title(f'{label}: {title}')
            ax.grid(axis='x', alpha=.2)
            for k, row in enumerate(ordered):
                ax.text(row[key] + ax.get_xlim()[1] * .008, k, f"{row[key]:.3f}", va='center', fontsize=7)
    fig.suptitle("Weak-label same-run fit; training setups and native output domains differ", fontsize=15)
    fig.savefig(OUT / 'model_comparison.png', dpi=170)
    plt.close(fig)


def write() -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    rows = records()
    with (OUT / 'report.csv').open('w', newline='', encoding='utf-8') as fh:
        writer = csv.DictWriter(fh, fieldnames=list(rows[0]))
        writer.writeheader(); writer.writerows(rows)
    result = {"status": "saved_same_run_weak_label_risk_comparison", "production_valid": False,
              "part_run_count": {"AC28": 2, "ME1a_DT": 1},
              "evaluation": "CAD-mapped 10 mm CSV-seeded weak risk mask on AC28 18h01 and ME1a 14h06 training runs",
              "held_out_baselines_available": False,
              "limitations": ["Unvalidated CAD-to-scan registration and defect extents",
                              "Different training epochs, source domains and CAD mapping support",
                              "Risk GNN was trained on two AC28 runs; other prior AC28 models used 18h01",
                              "Optimized residual GNN predicts risk only; Risk GNN also predicts deformation"],
              "rows": rows}
    (OUT / 'report.json').write_text(json.dumps(result, indent=2), encoding='utf-8')
    figure(rows)
    risk = next(r for r in rows if r['model'] == 'Risk GNN')
    mesh = next(r for r in rows if r['model'] == 'Mesh GNN')
    ac_ap = max(rows, key=lambda r: r['ac28_ap'])
    ac_iou = max(rows, key=lambda r: r['ac28_iou'])
    me_ap = max(rows, key=lambda r: r['me1a_ap'])
    me_iou = max(rows, key=lambda r: r['me1a_iou'])
    table_rows = ''.join(
        f"<tr><th>{html.escape(r['model'])}</th><td>{r['ac28_ap']:.3f}</td><td>{r['ac28_iou']:.3f}</td>"
        f"<td>{r['me1a_ap']:.3f}</td><td>{r['me1a_iou']:.3f}</td><td>{r['ac28_coverage']:.3f}</td>"
        f"<td>{r['me1a_coverage']:.3f}</td><td>{r['ac28_epochs']}</td><td>{r['me1a_epochs']}</td></tr>"
        for r in sorted(rows, key=lambda r: r['me1a_ap'], reverse=True))
    scopes = ''.join(f"<tr><th>{html.escape(r['model'])}</th><td>{html.escape(r['ac28_training_runs'])}</td>"
                     f"<td>{html.escape(r['me1a_training_runs'])}</td></tr>" for r in rows)
    page = f'''<!doctype html><html lang="en"><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"><title>AC28 and ME1a: all risk models compared</title>
<style>body{{font:15px/1.5 system-ui,sans-serif;max-width:1400px;margin:auto;padding:24px;background:#f4f7fb;color:#1c2b3a}}section{{background:white;border:1px solid #dbe3ec;border-radius:10px;padding:18px;margin:17px 0}}.warning{{background:#fff0cc;border-left:5px solid #b98109;padding:14px}}.grid{{display:grid;grid-template-columns:repeat(auto-fit,minmax(270px,1fr));gap:10px}}.grid div{{background:#eaf2f8;border-radius:8px;padding:12px}}table{{border-collapse:collapse;width:100%;min-width:850px}}th,td{{border-bottom:1px solid #dce3eb;padding:7px;text-align:right;white-space:nowrap}}th:first-child,td:first-child{{text-align:left}}.scroll{{overflow:auto}}img{{max-width:100%}}a{{color:#125d91}}</style>
<h1>Is Risk GNN the best risk model?</h1><p class="warning">These are fits to provisional masks on the training runs, with unvalidated scan registration. Output resolutions, CAD mapping coverage, loss functions and epoch counts differ. The table can compare saved diagnostics; it cannot identify the best production predictor. Other model families have no held-out run scores.</p>
<section><div class="grid"><div><b>AC28 highest AP</b><br>{html.escape(ac_ap['model'])}: {ac_ap['ac28_ap']:.3f}</div><div><b>AC28 highest IoU</b><br>{html.escape(ac_iou['model'])}: {ac_iou['ac28_iou']:.3f}</div><div><b>ME1a highest AP</b><br>{html.escape(me_ap['model'])}: {me_ap['me1a_ap']:.3f}</div><div><b>ME1a highest IoU</b><br>{html.escape(me_iou['model'])}: {me_iou['me1a_iou']:.3f}</div></div><p><strong>Risk GNN:</strong> AC28 AP {risk['ac28_ap']:.3f}, IoU {risk['ac28_iou']:.3f}; ME1a AP {risk['me1a_ap']:.3f}, IoU {risk['me1a_iou']:.3f}. Its AC28 AP lead over Mesh GNN is {risk['ac28_ap']-mesh['ac28_ap']:.3f}. It does not lead AC28 IoU or either ME1a metric. The optimized residual GNN predicts risk only; Risk GNN also emits a deformation field.</p></section>
<section><h2>All saved model results on common CAD nodes</h2><div class="scroll"><table><thead><tr><th>Model</th><th>AC28 AP</th><th>AC28 IoU</th><th>ME1a AP</th><th>ME1a IoU</th><th>AC28 coverage</th><th>ME1a coverage</th><th>AC28 epochs</th><th>ME1a epochs</th></tr></thead><tbody>{table_rows}</tbody></table></div><img src="model_comparison.png" alt="AP and IoU of all models on AC28 and ME1a"></section>
<section><h2>Training scope</h2><p>The Risk GNN used both AC28 runs for its AC28 entry. Other original models used AC28 18h01. The residual GNN individual entry uses one run of each part; its joint entry uses both AC28 runs and ME1a. Every displayed score evaluates a run included in that model's training. Only the residual GNN has a separate leave-one-run-out experiment, and its held-out IoU was zero for all three runs. The five conditional diffusion variants generate outputs on different native domains; their diagnostic risk scores are in the main report and are not ranked here.</p><div class="scroll"><table><thead><tr><th>Model</th><th>AC28 training</th><th>ME1a training</th></tr></thead><tbody>{scopes}</tbody></table></div><p><a href="../gnn_joint_hard/index.html">Residual GNN held-out experiment</a> · <a href="../../models/cad_predictions/index.html">AC28 model gallery</a> · <a href="../../../experiments/ME1a_DT/results/models/cad_predictions/index.html">ME1a model gallery</a></p></section>
<p><a href="report.csv">All metrics CSV</a> · <a href="report.json">Machine-readable results</a> · <a href="../index.html">Main comparison report</a></p></html>'''
    (OUT / 'index.html').write_text(page, encoding='utf-8')
    print(OUT / 'index.html')


if __name__ == '__main__': write()