#!/usr/bin/env python3
"""
WP4 -- evaluate a grid model on the fixed-occupancy MC scan.

    python scripts/04_evaluate.py --checkpoint Results/<exp>/Models/best.pth
    python scripts/04_evaluate.py --checkpoint ... --data <folder>/eval_scan

Produces, per occupancy point:
  efficiency / purity at the working threshold, position RMS and sigma68,
  energy bias/resolution, sub-pixel non-uniformity (risk R1),
plus the threshold ROC, the pair-resolution curve (efficiency vs distance to the
nearest other photon -- the figure of merit that replaces "efficiency vs
occupancy"), and residual vs conversion depth z.
"""
import argparse
import glob
import json
import os
import re
import sys
from pathlib import Path

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

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "src"))
from datasets import GridFrameDataset, collate_with_photons        # noqa: E402
from decode import decode                                          # noqa: E402
from metrics import (DetectionStats, greedy_match,                 # noqa: E402
                     nearest_neighbour_distance)
from models import build_model                                     # noqa: E402

THRESHOLDS = np.array([.05, .1, .2, .3, .4, .5, .6, .7, .8, .9, .95, .98])
NEIGHBOUR_BIN_EDGES = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 1e9])
ENERGY_BIN_WIDTH = 2.5                         # keV; one bin only for a
                                               # monochromatic set, which is fine


def energy_bin_edges(energies, width=ENERGY_BIN_WIDTH):
    """2.5 keV bins on a fixed grid, spanning the energies of the set."""
    low = np.floor(energies.min() / width) * width
    high = np.ceil(energies.max() / width) * width
    return np.arange(low, high + width / 2, width)


def occupancy_of_filename(path):
    match = re.search(r"occ([0-9p]+)pct", os.path.basename(path))
    return float(match.group(1).replace("p", ".")) / 100 if match else float("nan")


@torch.no_grad()
def run_file(model, conf, path, halo, batch_size=512, workers=8, max_frames=None,
             roc_frames=2000):
    dataset = GridFrameDataset([path], augment=False, return_photons=True,
                               input_scale=float(conf.data.input_scale),
                               energy_scale=float(conf.data.energy_scale),
                               max_frames=max_frames, name=os.path.basename(path))
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False,
                        num_workers=workers, collate_fn=collate_with_photons,
                        pin_memory=True)
    frame_size = dataset.frame_size
    working_threshold = float(conf.eval.threshold)
    lowest_threshold = float(THRESHOLDS.min())

    stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
                           match_radius=float(conf.eval.match_radius),
                           halo=halo, frame_size=frame_size)
    roc_predicted, roc_true = [], []

    for frames, objectness, offset, energy, photons in loader:
        frames = frames.cuda(non_blocking=True)
        prediction = model(frames).float()
        points = decode(prediction, threshold=lowest_threshold,
                        merge_distance=float(conf.eval.merge_distance), halo=halo,
                        predict_energy=bool(conf.model.predict_energy),
                        energy_scale=float(conf.data.energy_scale))
        points = [frame_points.numpy() for frame_points in points]
        truth = [frame_truth.numpy() for frame_truth in photons]
        stats.update([p[p[:, 3] > working_threshold] for p in points], truth)
        if len(roc_predicted) < roc_frames:
            roc_predicted.extend(points)
            roc_true.extend(truth)

    summary = stats.summary()
    differential = stats.differential()

    # ---- threshold sweep on a subset -------------------------------------
    roc = []
    for threshold in THRESHOLDS:
        subset_stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
                                      match_radius=float(conf.eval.match_radius),
                                      halo=halo, frame_size=frame_size)
        subset_stats.update([p[p[:, 3] > threshold] for p in roc_predicted[:roc_frames]],
                            roc_true[:roc_frames])
        subset_summary = subset_stats.summary()
        roc.append((float(threshold), subset_summary["efficiency"],
                    subset_summary["purity"], subset_summary["rms_1d_um"]))

    # ---- pair resolution: efficiency vs truth nearest-neighbour distance --
    n_bins = len(NEIGHBOUR_BIN_EDGES) - 1
    n_true_binned = np.zeros(n_bins)
    n_matched_binned = np.zeros(n_bins)
    for predicted, truth in zip(roc_predicted[:roc_frames], roc_true[:roc_frames]):
        predicted = predicted[predicted[:, 3] > working_threshold]
        if len(truth) == 0:
            continue
        inside = ((truth[:, 0] >= halo) & (truth[:, 0] < frame_size - halo) &
                  (truth[:, 1] >= halo) & (truth[:, 1] < frame_size - halo))
        truth = truth[inside]
        if len(truth) == 0:
            continue
        neighbour_distance = nearest_neighbour_distance(truth[:, :2])
        bin_index = np.digitize(neighbour_distance, NEIGHBOUR_BIN_EDGES) - 1
        for index in range(n_bins):
            n_true_binned[index] += (bin_index == index).sum()
        if len(predicted):
            _, index_true = greedy_match(predicted[:, :2], truth[:, :2],
                                         float(conf.eval.match_radius))
            for index in range(n_bins):
                n_matched_binned[index] += (bin_index[index_true] == index).sum()
    pair_resolution = dict(
        edges=NEIGHBOUR_BIN_EDGES[:-1].tolist(),
        n_true=n_true_binned.tolist(),
        n_match=n_matched_binned.tolist(),
        efficiency=(n_matched_binned / np.maximum(n_true_binned, 1)).tolist())

    return (summary, differential, roc, pair_resolution,
            stats.versus_energy(energy_bin_edges(dataset.photons[:, 2])))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--data", default=None,
                        help="folder with the eval_scan h5 files")
    parser.add_argument("--max-frames", type=int, default=None)
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--outdir", default=None)
    args = parser.parse_args()

    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    conf = OmegaConf.create(checkpoint["conf"])
    model = build_model(conf).cuda().eval()
    model.load_state_dict(checkpoint["model"])
    halo = model.rf_radius
    print(f"[eval] checkpoint epoch {checkpoint['epoch']}, R={halo}, "
          f"{sum(p.numel() for p in model.parameters())} params, "
          f"threshold {conf.eval.threshold}, merge {conf.eval.merge_distance} px")

    data_folder = args.data or os.path.join(conf.data.folder, "eval_scan")
    files = sorted(glob.glob(os.path.join(data_folder, "*.h5")),
                   key=occupancy_of_filename)
    if not files:
        sys.exit(f"no eval files in {data_folder}")
    output_dir = Path(args.outdir or Path(args.checkpoint).parent.parent)
    (output_dir / "Plots").mkdir(parents=True, exist_ok=True)

    results = {}
    for path in files:
        occupancy = occupancy_of_filename(path)
        summary, differential, roc, pair_resolution, per_energy = run_file(
            model, conf, path, halo, workers=args.workers, max_frames=args.max_frames)
        results[f"{occupancy*100:g}%"] = dict(occupancy=occupancy, summary=summary,
                                              roc=roc, pair=pair_resolution,
                                              per_energy=per_energy,
                                              _differential=differential)
        print(f"  occ {occupancy*100:5.2f}%  eff {summary['efficiency']*100:6.2f}%  "
              f"pur {summary['purity']*100:6.2f}%  RMS {summary['rms_1d_um']:.3f} um  "
              f"s68 {summary['sigma68_1d_um']:.3f} um  "
              f"dE {summary.get('energy_rms_keV', float('nan')):.3f} keV  "
              f"subpix-nonunif {summary['subpixel_nonuniformity']*100:.1f}%",
              flush=True)

    populated = [row for row in list(results.values())[len(results)//2]["per_energy"]
                 if row["n_true"] > 100]
    if len(populated) > 1:
        print(f"\n  per energy bin at occupancy "
              f"{list(results.values())[len(results)//2]['occupancy']*100:g}%:")
        print(f"    {'keV':>11} {'n':>9} {'eff%':>7} {'RMS um':>8} {'s68 um':>8} "
              f"{'dE bias':>8} {'dE rms':>8}")
        for row in populated:
            print(f"    {row['energy_low']:4.1f}-{row['energy_high']:4.1f} "
                  f"{row['n_true']:9d} {row['efficiency']*100:7.2f} "
                  f"{row['rms_1d_um']:8.3f} {row['sigma68_1d_um']:8.3f} "
                  f"{row['energy_bias_keV']:8.3f} {row['energy_rms_keV']:8.3f}")
        plot_versus_energy(results, output_dir)

    plot_all(results, output_dir, float(conf.eval.threshold))
    serialisable = {label: {key: value for key, value in result.items()
                            if key != "_differential"}
                    for label, result in results.items()}
    with open(output_dir / "eval_summary.json", "w") as handle:
        json.dump(serialisable, handle, indent=1)
    print(f"[eval] wrote {output_dir/'eval_summary.json'} and {output_dir/'Plots'}")


def plot_versus_energy(results, output_dir):
    """Efficiency, position resolution and energy resolution against photon energy."""
    figure, axes = plt.subplots(1, 3, figsize=(15, 4.4))
    for label, result in results.items():
        rows = [row for row in result["per_energy"] if row["n_true"] > 100]
        if len(rows) < 2:
            continue
        centre = np.array([0.5 * (r["energy_low"] + r["energy_high"]) for r in rows])
        axes[0].plot(centre, [r["efficiency"] * 100 for r in rows], "o-", label=f"occ {label}")
        axes[1].plot(centre, [r["sigma68_1d_um"] for r in rows], "o-", label=f"occ {label}")
        axes[2].plot(centre, [r["energy_rms_keV"] for r in rows], "o-", label=f"occ {label}")
    for axis, ylabel in zip(axes, ("efficiency [%]", r"position $\sigma_{68}$ [$\mu$m]",
                                   r"energy RMS [keV]")):
        axis.set_xlabel("photon energy [keV]")
        axis.set_ylabel(ylabel)
        axis.grid(alpha=.3)
        axis.legend(fontsize=8)
    figure.tight_layout()
    figure.savefig(output_dir / "Plots" / "eval_vs_energy.png", dpi=130)
    plt.close(figure)


def plot_all(results, output_dir, working_threshold):
    occupancy = np.array([r["occupancy"] for r in results.values()]) * 100
    efficiency = np.array([r["summary"]["efficiency"] for r in results.values()]) * 100
    purity = np.array([r["summary"]["purity"] for r in results.values()]) * 100
    rms = np.array([r["summary"]["rms_1d_um"] for r in results.values()])
    sigma68 = np.array([r["summary"]["sigma68_1d_um"] for r in results.values()])

    figure, axes = plt.subplots(2, 3, figsize=(16, 9))

    axes[0, 0].plot(occupancy, efficiency, "o-", label="efficiency")
    axes[0, 0].plot(occupancy, purity, "s-", label="purity")
    axes[0, 0].set_xlabel("occupancy [%]")
    axes[0, 0].set_ylabel("[%]")
    axes[0, 0].set_title(f"working point p>{working_threshold}")
    axes[0, 0].legend()
    axes[0, 0].grid(alpha=.3)

    axes[0, 1].plot(occupancy, rms, "o-", label="RMS")
    axes[0, 1].plot(occupancy, sigma68, "s-", label=r"$\sigma_{68}$")
    axes[0, 1].set_xlabel("occupancy [%]")
    axes[0, 1].set_ylabel(r"residual [$\mu$m]")
    axes[0, 1].legend()
    axes[0, 1].grid(alpha=.3)

    for label, result in results.items():
        roc = np.array(result["roc"])
        axes[0, 2].plot(roc[:, 2] * 100, roc[:, 1] * 100, ".-", label=f"occ {label}")
    axes[0, 2].set_xlabel("purity [%]")
    axes[0, 2].set_ylabel("efficiency [%]")
    axes[0, 2].set_title("threshold sweep")
    axes[0, 2].legend(fontsize=8)
    axes[0, 2].grid(alpha=.3)

    for label, result in results.items():
        residual = result["_differential"]["residual_um"]
        if len(residual):
            axes[1, 0].hist(residual.ravel(), bins=200, range=(-6, 6),
                            histtype="step", density=True, label=f"occ {label}")
    axes[1, 0].set_xlabel(r"residual [$\mu$m]")
    axes[1, 0].set_yscale("log")
    axes[1, 0].legend(fontsize=8)
    axes[1, 0].grid(alpha=.3)

    for label, result in results.items():
        pair = result["pair"]
        edges = np.array(pair["edges"])
        pair_efficiency = np.array(pair["efficiency"]) * 100
        populated = np.array(pair["n_true"]) > 50
        axes[1, 1].plot(edges[populated], pair_efficiency[populated], "o-",
                        label=f"occ {label}")
    axes[1, 1].set_xlabel("distance to nearest other photon [px]")
    axes[1, 1].set_ylabel("efficiency [%]")
    axes[1, 1].set_xlim(0, 6)
    axes[1, 1].set_title("pair resolution")
    axes[1, 1].legend(fontsize=8)
    axes[1, 1].grid(alpha=.3)

    middle = list(results.values())[len(results) // 2]
    subpixel = middle["_differential"]["subpixel"]
    if len(subpixel):
        histogram, _, _ = np.histogram2d(subpixel[:, 0], subpixel[:, 1], bins=25,
                                         range=[[0, 1], [0, 1]])
        image = axes[1, 2].imshow(histogram.T / histogram.mean(), origin="lower",
                                  extent=[0, 1, 0, 1], vmin=0.5, vmax=1.5,
                                  cmap="viridis")
        plt.colorbar(image, ax=axes[1, 2])
    axes[1, 2].set_title(f"sub-pixel uniformity (occ {middle['occupancy']*100:g}%)\n"
                         "flat = no grid artefact (risk R1)")
    axes[1, 2].set_xlabel("sub-pixel x")
    axes[1, 2].set_ylabel("sub-pixel y")

    figure.tight_layout()
    figure.savefig(output_dir / "Plots" / "eval_scan.png", dpi=130)
    plt.close(figure)

    # ---- residual vs depth, vs isolation, and energy ---------------------
    figure, axes = plt.subplots(1, 3, figsize=(16, 4.4))
    depth_edges = np.linspace(0, 650, 14)
    neighbour_edges = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 10])
    for label, result in results.items():
        differential = result["_differential"]
        if not len(differential["z_um"]):
            continue
        depth_bin = np.digitize(differential["z_um"], depth_edges) - 1
        depth_rms = [differential["residual_um"][depth_bin == index].std()
                     if (depth_bin == index).sum() > 50 else np.nan
                     for index in range(len(depth_edges) - 1)]
        axes[0].plot(0.5 * (depth_edges[:-1] + depth_edges[1:]), depth_rms, "o-",
                     label=f"occ {label}")

        neighbour_bin = np.digitize(differential["nn_px"], neighbour_edges) - 1
        neighbour_rms = [differential["residual_um"][neighbour_bin == index].std()
                         if (neighbour_bin == index).sum() > 50 else np.nan
                         for index in range(len(neighbour_edges) - 1)]
        axes[1].plot(neighbour_edges[:-1], neighbour_rms, "o-", label=f"occ {label}")

        if np.isfinite(differential["e_pred"]).any():
            axes[2].hist(differential["e_pred"] - differential["e_true"], bins=150,
                         range=(-4, 4), histtype="step", density=True,
                         label=f"occ {label}")
    axes[0].set_xlabel(r"conversion depth z [$\mu$m]")
    axes[0].set_ylabel(r"RMS [$\mu$m]")
    axes[0].set_title("resolution vs charge-cloud size")
    axes[0].legend(fontsize=8)
    axes[0].grid(alpha=.3)
    axes[1].set_xlabel("distance to nearest other photon [px]")
    axes[1].set_ylabel(r"RMS [$\mu$m]")
    axes[1].legend(fontsize=8)
    axes[1].grid(alpha=.3)
    axes[2].set_xlabel(r"$E_{pred}-E_{true}$ [keV]")
    axes[2].set_yscale("log")
    axes[2].legend(fontsize=8)
    axes[2].grid(alpha=.3)
    figure.tight_layout()
    figure.savefig(output_dir / "Plots" / "eval_differential.png", dpi=130)
    plt.close(figure)


if __name__ == "__main__":
    main()