#!/usr/bin/env python3
"""
WP4/5.3 -- MC energy response of a polychromatic grid model.

    python scripts/07_energy_response.py \
        --checkpoint Results/260901_poly5to30keV_Moench014_00/Models/best.pth

The energy head is a regression trained on a flat 5-30 keV spectrum, so it
returns something close to the posterior mean: every prediction is pulled
towards the centre of the training range, most strongly within ~1 keV of the
edges where the prior is one-sided.  A spectral line at E therefore shows up at
g(E) != E.  This script measures g on MC in fine true-energy bins -- the
position of the Gaussian core of E_pred, i.e. what a line fit on measured data
returns -- together with the core width and the detection efficiency, and
writes

    Results/<exp>/energy_response.npz    per-bin table + all matched pairs
    Results/<exp>/Plots/energy_response.png

The correction applied to measured photons is E_corr = g^-1(E_pred) (see
src/xrf_analysis.py: EnergyCorrection).  Inverting the forward curve, rather
than regressing E_true on E_pred, is what puts a line back at its energy: the
reverse regression is the posterior mean again and would shrink a second time.

Only frames with occupancy <= --max-occupancy are used, so that the response
matches the low-flux XRF runs it is applied to (0.7-1.2 %).
"""
import argparse
import glob
import os
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

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

DEFAULT_DATA = [
    "{folder}/eval_scan/occ0p5pct.h5",
    "{folder}/eval_scan/occ1pct.h5",
    "{folder}/eval_scan/occ2pct.h5",
    "{folder}/test/shard_00.h5",
    "{folder}/test/shard_01.h5",
]


def load_frames(path, max_occupancy):
    import h5py
    with h5py.File(path, "r") as handle:
        occupancy = handle["occupancy"][:]
        keep = np.nonzero(occupancy <= max_occupancy)[0]
        frames = handle["frames"][:][keep]
        photons = handle["photons"][:]
        pointer = handle["frame_ptr"][:]
    truth = [photons[pointer[i]:pointer[i + 1]] for i in keep]
    return frames, truth, occupancy[keep]


@torch.no_grad()
def collect(model, conf, files, halo, max_occupancy, batch_size=1024):
    """Matched (E_true, E_pred) pairs, and every valid truth photon with a flag."""
    input_scale = float(conf.data.input_scale)
    energy_scale = float(conf.data.energy_scale)
    threshold = float(conf.eval.threshold)
    merge_distance = float(conf.eval.merge_distance)
    match_radius = float(conf.eval.match_radius)
    e_true, e_pred, truth_energy, truth_found, fake_energy = [], [], [], [], []
    n_frames = 0
    for path in files:
        frames, truth, occupancy = load_frames(path, max_occupancy)
        size = frames.shape[-1]
        print(f"[response] {os.path.basename(path)}: {len(frames)} frames "
              f"with occupancy <= {max_occupancy*100:.1f} %", flush=True)
        n_frames += len(frames)
        for start in range(0, len(frames), batch_size):
            batch = torch.from_numpy(frames[start:start + batch_size] / input_scale)
            prediction = model(batch.float().unsqueeze(1).cuda()).float()
            points = decode(prediction, threshold=threshold,
                            merge_distance=merge_distance, halo=halo,
                            predict_energy=True, energy_scale=energy_scale)
            for predicted, true in zip(points, truth[start:start + batch_size]):
                predicted = predicted.numpy().astype(np.float64)
                inside = ((true[:, 0] >= halo) & (true[:, 0] < size - halo) &
                          (true[:, 1] >= halo) & (true[:, 1] < size - halo))
                true = true[inside].astype(np.float64)
                found = np.zeros(len(true), bool)
                matched_prediction = np.zeros(len(predicted), bool)
                if len(true) and len(predicted):
                    index_pred, index_true = greedy_match(predicted[:, :2], true[:, :2],
                                                          match_radius)
                    found[index_true] = True
                    matched_prediction[index_pred] = True
                    e_true.append(true[index_true, 2])
                    e_pred.append(predicted[index_pred, 2])
                truth_energy.append(true[:, 2])
                truth_found.append(found)
                fake_energy.append(predicted[~matched_prediction, 2])
    return (np.concatenate(e_true), np.concatenate(e_pred),
            np.concatenate(truth_energy), np.concatenate(truth_found),
            np.concatenate(fake_energy), n_frames)


def response_table(e_true, e_pred, truth_energy, truth_found, edges):
    rows = {key: [] for key in ("centre", "n", "efficiency", "median", "mean",
                                "core_mu", "core_mu_err", "core_sigma",
                                "core_sigma_err", "sigma68", "rms")}
    for low, high in zip(edges[:-1], edges[1:]):
        selected = (e_true >= low) & (e_true < high)
        in_bin = (truth_energy >= low) & (truth_energy < high)
        residual_reference = e_true[selected]
        predicted = e_pred[selected]
        core = fit_gaussian_core(predicted)
        rows["centre"].append(0.5 * (low + high))
        rows["n"].append(int(selected.sum()))
        rows["efficiency"].append(float(truth_found[in_bin].mean()) if in_bin.any()
                                  else np.nan)
        rows["median"].append(float(np.median(predicted)) if len(predicted) else np.nan)
        rows["mean"].append(float(predicted.mean()) if len(predicted) else np.nan)
        rows["core_mu"].append(core["mu"])
        rows["core_mu_err"].append(core["mu_err"])
        rows["core_sigma"].append(core["sigma"])
        rows["core_sigma_err"].append(core["sigma_err"])
        residual = predicted - residual_reference
        rows["sigma68"].append(float(np.percentile(np.abs(residual - np.median(residual)),
                                                   68.27)) if len(residual) else np.nan)
        rows["rms"].append(float(residual.std()) if len(residual) else np.nan)
    return {key: np.asarray(value) for key, value in rows.items()}


def plot(table, e_true, e_pred, output_path, title):
    centre = table["centre"]
    low, high = float(np.floor(e_true.min())), float(np.ceil(e_true.max()))
    figure, axes = plt.subplots(2, 2, figsize=(12, 8.5))
    axis = axes[0, 0]
    histogram, x_edges, y_edges = np.histogram2d(
        e_true, e_pred, bins=[int(5 * (high - low)), int(5 * (high - low + 5))],
        range=[[low, high], [low - 2.5, high + 2.5]])
    axis.imshow(np.log10(histogram.T + 1), origin="lower", aspect="auto",
                extent=[x_edges[0], x_edges[-1], y_edges[0], y_edges[-1]],
                cmap="viridis")
    axis.plot([low, high], [low, high], "w--", lw=.8)
    axis.set_xlabel("E_true [keV]")
    axis.set_ylabel("E_pred [keV]")
    axis.set_title("matched photons, log10 counts")

    axis = axes[0, 1]
    axis.errorbar(centre, table["core_mu"] - centre, table["core_mu_err"], fmt=".",
                  ms=3, label="Gaussian core position")
    axis.plot(centre, table["median"] - centre, lw=1, label="median")
    axis.plot(centre, table["mean"] - centre, lw=1, label="mean")
    axis.axhline(0, color="k", lw=.6)
    axis.set_xlabel("E_true [keV]")
    axis.set_ylabel("E_pred - E_true [keV]")
    axis.set_title("g(E) - E: pull towards the centre")
    axis.legend(fontsize=8)
    axis.grid(alpha=.3)

    axis = axes[1, 0]
    axis.plot(centre, table["core_sigma"], ".", ms=4, label=r"core $\sigma$")
    axis.plot(centre, table["sigma68"], lw=1, label=r"$\sigma_{68}$ of residual")
    axis.plot(centre, table["rms"], lw=1, label="rms of residual")
    axis.set_xlabel("E_true [keV]")
    axis.set_ylabel("keV")
    axis.set_ylim(0, 1.6)
    axis.set_title("energy resolution")
    axis.legend(fontsize=8)
    axis.grid(alpha=.3)

    axis = axes[1, 1]
    axis.plot(centre, table["efficiency"] * 100, ".", ms=4)
    axis.set_xlabel("E_true [keV]")
    axis.set_ylabel("efficiency [%]")
    axis.set_title("detection efficiency")
    axis.grid(alpha=.3)
    figure.suptitle(title)
    figure.tight_layout()
    figure.savefig(output_path, dpi=130)
    plt.close(figure)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--data", nargs="*", default=None)
    parser.add_argument("--max-occupancy", type=float, default=0.025)
    parser.add_argument("--bin-width", type=float, default=0.25)
    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
    files = args.data or [pattern.format(folder=conf.data.folder)
                          for pattern in DEFAULT_DATA]
    files = [f for pattern in files for f in sorted(glob.glob(pattern))]

    e_true, e_pred, truth_energy, truth_found, fake_energy, n_frames = collect(
        model, conf, files, halo, args.max_occupancy)
    # the true-energy range of the MC set: 5-30 keV or 3-30 keV
    edges = np.arange(np.floor(truth_energy.min()), np.ceil(truth_energy.max()) + 1e-9,
                      args.bin_width)
    table = response_table(e_true, e_pred, truth_energy, truth_found, edges)

    experiment_dir = Path(args.checkpoint).resolve().parent.parent
    output = experiment_dir / "energy_response.npz"
    np.savez_compressed(
        output, edges=edges, **{f"table_{k}": v for k, v in table.items()},
        e_true=e_true.astype(np.float32), e_pred=e_pred.astype(np.float32),
        truth_energy=truth_energy.astype(np.float32), truth_found=truth_found,
        fake_energy=fake_energy.astype(np.float32),
        checkpoint=str(Path(args.checkpoint).resolve()), epoch=int(checkpoint["epoch"]),
        max_occupancy=args.max_occupancy, n_frames=n_frames,
        files=np.array([str(f) for f in files]))
    (experiment_dir / "Plots").mkdir(exist_ok=True)
    plot(table, e_true, e_pred, experiment_dir / "Plots" / "energy_response.png",
         f"{experiment_dir.name} (epoch {checkpoint['epoch']}), "
         f"occupancy <= {args.max_occupancy*100:.1f} %, {len(e_true)} matched photons")
    print(f"[response] {n_frames} frames, {len(truth_energy)} truth photons, "
          f"{len(e_true)} matched, {len(fake_energy)} unmatched predictions")
    print(f"[response] efficiency {truth_found.mean()*100:.2f} %")
    for energy in [e for e in (3.5, 4.5, 5.5, 7.0, 8.0, 12.0, 17.5, 22.0, 26.0, 29.5)
                   if edges[0] < e < edges[-1]]:
        i = np.searchsorted(edges, energy) - 1
        print(f"  E={table['centre'][i]:6.3f}  g-E={table['core_mu'][i]-table['centre'][i]:+.3f}"
              f"  core sigma={table['core_sigma'][i]:.3f}  eff={table['efficiency'][i]*100:.2f} %")
    print(f"[response] wrote {output}")


if __name__ == "__main__":
    main()