#!/usr/bin/env python3
"""
Compare poly GridNet models: MC metrics, energy response, XRF lines, efficiency
against the classical 3x3 sum, fake rate and speed, in one report.

    python scripts/08_compare_models.py --out Results/<dir>/report \\
        --model "5-30 keV [32,64,128,128]" Results/260901_poly5to30keV_Moench014_00 - \\
        --model "3-30 keV [32,64,96,96]"   Results/<exp> c96

Each model is (label, experiment dir, XRF tag).  The tag is the --tag its XRF
runs were reconstructed with (`-` for none), so the reconstructions are read
from <inference.output_folder>/<run>[_<tag>].  Expected inputs per model:
eval_summary.json and energy_response.npz in the experiment dir (04, 07), and
AgMo_1_factor_10, CuCrCo_1_factor_10 and Pedestal_CuCrCo_1 reconstructions (05).
Whatever is missing shows up as n/a, with the reason listed at the end.

Writes report.md, report.json and report_response.png.
"""
import argparse
import json
import sys
import time
import traceback
from pathlib import Path

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

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "src"))
from models import build_model                                     # noqa: E402
from xrf_analysis import classical_and_model, fit_lines, line_centroid  # noqa: E402

import importlib                                                   # noqa: E402
sys.path.insert(0, str(HERE))
predict_fp16 = importlib.import_module("05_infer_measurement").predict_fp16

RESPONSE_ENERGIES = (3.5, 4.5, 5.4, 6.9, 8.0, 12.0, 17.4, 22.1, 27.5, 29.4)
XRF_FITS = {
    "AgMo_1_factor_10": [
        dict(alpha=["Mo_Ka"], window=(16.0, 18.9)),
        dict(alpha=["Mo_Kb"], window=(18.75, 20.7)),
        dict(alpha=["Ag_Ka"], window=(20.6, 23.6)),
    ],
    "CuCrCo_1_factor_10": [
        dict(alpha=["Cr_Ka", "Co_Ka", "Cu_Ka"],
             beta={"Cr_Kb": "Cr_Ka", "Co_Kb": "Co_Ka", "Cu_Kb": "Cu_Ka"},
             ratio_fixed={"Co_Kb": 0.135}, fit_stretch=False, window=(4.85, 9.6)),
    ],
}
XRF_LINES = ["Cr_Ka", "Co_Ka", "Cu_Ka", "Mo_Ka", "Mo_Kb", "Ag_Ka"]
DARK_RUN = "Pedestal_CuCrCo_1"


def fmt(value, spec, missing="n/a"):
    if value is None or (isinstance(value, float) and not np.isfinite(value)):
        return missing
    return format(value, spec)


def model_meta(experiment, errors):
    out = {}
    try:
        checkpoint = torch.load(experiment / "Models" / "best.pth", map_location="cpu",
                                weights_only=False)
        conf = OmegaConf.create(checkpoint["conf"])
        model = build_model(conf)
        out.update(channels=list(conf.model.channels), R=model.rf_radius,
                   params=sum(p.numel() for p in model.parameters()),
                   best_epoch=int(checkpoint["epoch"]))
        history = json.load(open(experiment / "history.json"))
        out.update(epochs=len(history["val_loss"]),
                   best_val_loss=float(min(history["val_loss"])))
        out["speed_fp16"] = speed(model)
    except Exception as error:                                     # noqa: BLE001
        errors.append(f"meta: {error!r}")
    return out


@torch.no_grad()
def speed(model, size=400, batch=32, n_iter=20):
    torch.backends.cudnn.benchmark = True
    model = model.cuda().eval().to(memory_format=torch.channels_last)
    x = (torch.randn(batch, 1, size, size, device="cuda") * 0.05).to(
        memory_format=torch.channels_last)
    def run():   # exactly the FP16 path of 05_infer_measurement
        return predict_fp16(model, x, 0.3)

    for _ in range(5):
        run()
    torch.cuda.synchronize()
    started = time.perf_counter()
    for _ in range(n_iter):
        run()
    torch.cuda.synchronize()
    return batch * n_iter / (time.perf_counter() - started)


def mc_metrics(experiment, errors):
    try:
        data = json.load(open(experiment / "eval_summary.json"))
        return {key: value["summary"] for key, value in data.items()}
    except Exception as error:                                     # noqa: BLE001
        errors.append(f"eval_summary: {error!r}")
        return {}


def response(experiment, errors):
    out = {}
    try:
        data = np.load(experiment / "energy_response.npz")
        centre = data["table_centre"]
        e_true, e_pred = data["e_true"], data["e_pred"]
        for energy in RESPONSE_ENERGIES:
            if not centre[0] - 0.13 <= energy <= centre[-1] + 0.13:
                continue
            i = int(np.argmin(np.abs(centre - energy)))
            near = np.abs(e_true - centre[i]) < 0.125
            out[energy] = dict(
                shift_meV=float((data["table_core_mu"][i] - centre[i]) * 1000),
                core_sigma=float(data["table_core_sigma"][i]),
                tail_percent=float((np.abs(e_pred[near] - e_true[near]) > 1).mean() * 100),
                efficiency=float(data["table_efficiency"][i]))
        out["_curve"] = dict(centre=centre.tolist(),
                             shift=(data["table_core_mu"] - centre).tolist(),
                             sigma=data["table_core_sigma"].tolist(),
                             efficiency=data["table_efficiency"].tolist())
    except Exception as error:                                     # noqa: BLE001
        errors.append(f"energy_response: {error!r}")
    return out


def xrf_lines(folder_of, errors):
    out = {}
    for run, groups in XRF_FITS.items():
        try:
            centres, counts = np.load(folder_of(run) / "energy_spectrum.npy")
            for group in groups:
                fit = fit_lines(centres, counts, group["alpha"], group.get("beta"),
                                window=group["window"],
                                beta_ratio_fixed=group.get("ratio_fixed"),
                                fit_stretch=group.get("fit_stretch", True))
                for name in group["alpha"]:
                    line = fit["lines"][name]
                    out[name] = dict(shift_meV=line["shift"] * 1000, sigma=line["sigma"],
                                     fwhm=line["fwhm"], chi2_ndf=fit["chi2_ndf"])
        except Exception as error:                                 # noqa: BLE001
            errors.append(f"xrf {run}: {error!r}")
    return out


def classical_lines(config, errors):
    """Line fits of the classical 3x3 spectra (model independent, cached)."""
    out = {}
    edges = np.arange(0.0, 40.025, 0.05)
    centres = 0.5 * (edges[1:] + edges[:-1])
    for run, groups in XRF_FITS.items():
        try:
            reference = classical_and_model(config, run, n_frames=10_000)
            counts = np.histogram(reference["classical"][:, 3], edges)[0]
            for group in groups:
                fit = fit_lines(centres, counts, group["alpha"], group.get("beta"),
                                window=group["window"],
                                beta_ratio_fixed=group.get("ratio_fixed"),
                                fit_stretch=group.get("fit_stretch", True), tails=False)
                for name in group["alpha"]:
                    line = fit["lines"][name]
                    out[name] = dict(shift_meV=line["shift"] * 1000, sigma=line["sigma"],
                                     fwhm=line["fwhm"])
        except Exception as error:                                 # noqa: BLE001
            errors.append(f"classical {run}: {error!r}")
    return out


def efficiency(config, checkpoint, tag, errors):
    out = {}
    for run in XRF_FITS:
        try:
            reference = classical_and_model(config, run, n_frames=10_000,
                                            checkpoint=checkpoint, tag=tag)
            classical, grid = reference["classical"], reference["grid"]
            found = reference["classical_match"] >= 0
            isolated = reference["isolation_classical"] > 4.0
            energy = classical[:, 3]
            deposited = reference["balance"][:, 1].sum()
            out[run] = dict(
                isolated_above5=float(found[isolated & (energy >= 5)].mean()),
                isolated_3to5=float(found[isolated & (energy >= 3) & (energy < 5)].mean()),
                count_ratio=float(len(grid) / len(classical)),
                energy_balance=float(grid[:, 3].sum() / deposited))
        except Exception as error:                                 # noqa: BLE001
            errors.append(f"efficiency {run}: {error!r}")
    return out


def run_summaries(folder_of, errors):
    out = {}
    for run in list(XRF_FITS) + [DARK_RUN]:
        try:
            summary = json.load(open(folder_of(run) / "summary.json"))
            out[run] = dict(n_photons=summary["n_photons"], n_frames=summary["n_frames"],
                            subpixel=summary.get("subpixel_nonuniformity_before"),
                            fps=summary.get("frames_per_second"),
                            precision=summary.get("precision", "fp32"))
        except Exception as error:                                 # noqa: BLE001
            errors.append(f"summary {run}: {error!r}")
    return out


def write_report(models, classical, out_dir):
    labels = list(models)
    lines = [f"# Poly model comparison  ({time.strftime('%Y-%m-%d %H:%M')})", ""]

    lines += ["## Models", "",
              "| model | channels | R | params | best epoch / epochs | FP16 frame/s (400×400, model only) |",
              "|---|---|---|---|---|---|"]
    for label in labels:
        m = models[label]["meta"]
        lines.append(f"| {label} | {m.get('channels', 'n/a')} | {m.get('R', 'n/a')} | "
                     f"{m.get('params', 'n/a')} | {m.get('best_epoch', 'n/a')} / "
                     f"{m.get('epochs', 'n/a')} | {fmt(m.get('speed_fp16'), '.0f')} |")

    lines += ["", "## MC, fixed-occupancy scan (threshold 0.3)", "",
              "| model | occupancy | efficiency | purity | RMS [µm] | σ68 [µm] | E σ68 [keV] | sub-pixel |",
              "|---|---|---|---|---|---|---|---|"]
    for label in labels:
        for occupancy, s in models[label]["mc"].items():
            if occupancy not in ("1%", "5%", "10%"):
                continue
            lines.append(f"| {label} | {occupancy} | {s['efficiency']*100:.2f} % | "
                         f"{s['purity']*100:.2f} % | {s['rms_1d_um']:.3f} | "
                         f"{s['sigma68_1d_um']:.3f} | {s.get('energy_sigma68_keV', float('nan')):.3f} | "
                         f"{s['subpixel_nonuniformity']*100:.1f} % |")

    lines += ["", "## MC energy response (occupancy ≤ 2.5 %)", "",
              "core σ [keV] / line shift g(E) − E [meV] / tails |ΔE| > 1 keV [%] / efficiency [%]", "",
              "| E [keV] | " + " | ".join(labels) + " |", "|---|" + "---|" * len(labels)]
    for energy in RESPONSE_ENERGIES:
        cells = []
        for label in labels:
            r = models[label]["response"].get(energy)
            cells.append("—" if r is None else
                         f"{r['core_sigma']:.3f} / {r['shift_meV']:+.0f} / "
                         f"{r['tail_percent']:.1f} / {r['efficiency']*100:.1f}")
        lines.append(f"| {energy} | " + " | ".join(cells) + " |")

    lines += ["", "## XRF lines (AgMo, CuCrCo slit factor 10)", "",
              "fitted − tabulated [meV] / core σ [keV] / FWHM [keV]", "",
              "| line | " + " | ".join(labels + ["classical 3×3"]) + " |",
              "|---|" + "---|" * (len(labels) + 1)]
    for name in XRF_LINES:
        cells = []
        for source in [models[label]["xrf"] for label in labels] + [classical]:
            r = source.get(name)
            cells.append("n/a" if r is None else
                         f"{r['shift_meV']:+.0f} / {r['sigma']:.3f} / {r['fwhm']:.2f}")
        lines.append(f"| {name} ({line_centroid(name):.2f}) | " + " | ".join(cells) + " |")

    lines += ["", "## Efficiency against the classical finder, fake rate, sub-pixel", "",
              "| model | run | isolated ≥ 5 keV found | isolated 3–5 keV found | grid / classical | energy balance | sub-pixel excess rms |",
              "|---|---|---|---|---|---|---|"]
    for label in labels:
        for run in XRF_FITS:
            e = models[label]["efficiency"].get(run, {})
            s = models[label]["runs"].get(run, {})
            lines.append(f"| {label} | {run} | {fmt(e.get('isolated_above5', np.nan)*100, '.2f')} % | "
                         f"{fmt(e.get('isolated_3to5', np.nan)*100, '.1f')} % | "
                         f"{fmt(e.get('count_ratio'), '.3f')} | {fmt(e.get('energy_balance'), '.4f')} | "
                         f"{fmt((s.get('subpixel') or np.nan)*100, '.1f')} % |")
    lines += ["", "| model | dark run photons | frames |", "|---|---|---|"]
    for label in labels:
        s = models[label]["runs"].get(DARK_RUN, {})
        lines.append(f"| {label} | {s.get('n_photons', 'n/a')} | {s.get('n_frames', 'n/a')} |")

    problems = [(label, e) for label in labels for e in models[label]["errors"]]
    if problems:
        lines += ["", "## Missing pieces", ""]
        lines += [f"- {label}: {error}" for label, error in problems]
    lines += ["", "Plots: `report_response.png`."]
    (out_dir / "report.md").write_text("\n".join(lines) + "\n")


def plot_response(models, out_dir):
    figure, axes = plt.subplots(1, 3, figsize=(17, 4.4))
    for label, model in models.items():
        curve = model["response"].get("_curve")
        if not curve:
            continue
        centre = np.asarray(curve["centre"])
        axes[0].plot(centre, np.asarray(curve["shift"]) * 1000, lw=1.2, label=label)
        axes[1].plot(centre, curve["sigma"], lw=1.2, label=label)
        axes[2].plot(centre, np.asarray(curve["efficiency"]) * 100, lw=1.2, label=label)
    axes[0].set_ylim(-200, 200)
    for axis, ylabel, title in zip(axes, ("g(E) − E [meV]", "core σ [keV]", "efficiency [%]"),
                                   ("line shift", "energy resolution", "detection efficiency")):
        axis.set_xlabel("E_true [keV]"); axis.set_ylabel(ylabel); axis.set_title(title)
        axis.grid(alpha=.3); axis.legend(fontsize=7)
    figure.tight_layout()
    figure.savefig(out_dir / "report_response.png", dpi=130)
    plt.close(figure)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", nargs=3, action="append", required=True,
                        metavar=("LABEL", "EXPERIMENT", "TAG"))
    parser.add_argument("--config", default=str(HERE.parent / "Configs" /
                                                 "infer_measurement_poly_Moench014.yaml"))
    parser.add_argument("--out", required=True)
    args = parser.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    output_folder = Path(OmegaConf.load(args.config).inference.output_folder)
    models = {}
    for label, experiment, tag in args.model:
        experiment = Path(experiment)
        tag = None if tag in ("-", "") else tag
        errors = []
        print(f"[compare] {label}: {experiment} (tag {tag})", flush=True)

        def folder_of(run, tag=tag):
            return output_folder / (run + (f"_{tag}" if tag else ""))

        models[label] = dict(
            meta=model_meta(experiment, errors), mc=mc_metrics(experiment, errors),
            response=response(experiment, errors), xrf=xrf_lines(folder_of, errors),
            efficiency=efficiency(args.config, str(experiment / "Models" / "best.pth"),
                                  tag, errors),
            runs=run_summaries(folder_of, errors), errors=errors)
        for error in errors:
            print(f"   missing: {error}", flush=True)
    classical_errors = []
    classical = classical_lines(args.config, classical_errors)
    if classical_errors:
        models[next(iter(models))]["errors"] += classical_errors

    write_report(models, classical, out_dir)
    plot_response(models, out_dir)

    def clean(value):
        if isinstance(value, dict):
            return {str(k): clean(v) for k, v in value.items()}
        if isinstance(value, (list, tuple)):
            return [clean(v) for v in value]
        if isinstance(value, (np.floating, np.integer)):
            return value.item()
        return value
    with open(out_dir / "report.json", "w") as handle:
        json.dump(clean(dict(models=models, classical=classical)), handle, indent=1)
    print(f"[compare] wrote {out_dir/'report.md'}")


if __name__ == "__main__":
    try:
        main()
    except Exception:                                              # noqa: BLE001
        traceback.print_exc()
        sys.exit(1)