#!/usr/bin/env python3
"""
WP4 -- baseline comparison on identical MC frames (gate G2).

    conda activate mlxid_demo          # this env has both torch and aare
    python scripts/06_baseline_comparison.py \
        --checkpoint Results/<exp>/Models/E100.pth --n-frames 4000

Three reconstruction chains are run over the *same* simulated frames and scored
against the same truth with the same matcher:

  grid          GridNet + position-space decode
  cluster+dl    aare VarClusterFinder -> cluster energy picks the multiplicity
                -> the pretrained 1/2/3-photon CNNs on 3x3 / 6x6 / 9x9 crops
  cluster+eta   the same cluster finder, 1-photon clusters only, classical eta
                interpolation with a Rosenblatt LUT fitted on the same events

Cluster cropping, energy windows and coordinate conventions follow
``DataProcess/UsefulFuncs.py::_processFrames``; the CNN input conventions follow
``DeepLearning/src/datasets.py``.  Every chain is scored on the same valid
region (frame minus the grid model's halo), so the numbers are comparable.
"""
import argparse
import importlib.util
import json
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
# GridBased/src and MLXIDDemo/src both define a `models` module.  MLXIDDemo's
# model_zoo does `from models import ...`, so that name is left to the demo and
# the GridBased one is loaded by path under its own name.
sys.path.insert(0, str(HERE.parent / "src"))
sys.path.insert(0, str(HERE.parent.parent / "MLXIDDemo" / "src"))
sys.path.append("/home/xie_x1/software/aareBuild")

import aare                                                        # noqa: E402
from decode import decode                                          # noqa: E402
from eta_interpolation_functions import (bilinear_xy_lookup,       # noqa: E402
                                         build_xy_lut_Rosenblatt)
from metrics import DetectionStats                                 # noqa: E402
from model_zoo import load_pretrained                              # noqa: E402


def _load_by_path(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


build_model = _load_by_path("grid_models", HERE.parent / "src" / "models.py").build_model

CLUSTER_SIZE = {1: 3, 2: 6, 3: 9}
ETA_BINS = 201
ETA_POSITION_WEIGHTS = np.array([-1.0, 0.0, 1.0]) + 0.5


def coordinate_channels(size):
    """The two coordinate planes singlePhotonDataset appends to its input."""
    axis = np.linspace(-0.5, 0.5, size)
    x_grid, y_grid = np.meshgrid(axis, axis, indexing="ij")
    return torch.from_numpy(np.stack([x_grid, y_grid])).float()


def find_clusters(frames, noise_keV, threshold_sigma=5, peripheral_factor=3,
                  n_neighbours=4):
    """Run the variable cluster finder frame by frame.

    Returns a list per frame of dicts with the seed pixel, the charge-weighted
    centre and the cluster energy.
    """
    frame_size = frames.shape[-1]
    noise_map = np.ones((frame_size, frame_size)) * noise_keV

    def new_finder():
        # A finder is built per frame.  Reusing one across frames is what
        # DataProcess/UsefulFuncs._processFrames does and is fine at its
        # 101x101 ROI, but on these 32x32 frames the internal hit buffer
        # overflows after a few frames and hits() starts returning stale
        # entries (and can abort the process).
        finder = aare.VarClusterFinder((frame_size, frame_size), threshold_sigma)
        finder.set_peripheralThresholdFactor(peripheral_factor)
        finder.set_noiseMap(noise_map)
        finder.set_numberOfNeighbours(n_neighbours)
        finder.set_empty_surroundingPixels(False)
        return finder

    per_frame = []
    for frame in frames:
        finder = new_finder()
        working_copy = np.array(frame, dtype=np.float64)   # the finder consumes it
        finder.find_clusters_X(working_copy)
        hits = finder.hits()
        clusters = []
        for index in range(len(hits)):
            size = hits["size"][index]
            energies = hits["enes"][index][:size]
            total = energies.sum()
            if total <= 0:
                continue
            clusters.append(dict(
                energy=float(hits["energy"][index]),
                seed_col=int(hits["col"][index]), seed_row=int(hits["row"][index]),
                centre_col=float((hits["cols"][index][:size] * energies).sum() / total),
                centre_row=float((hits["rows"][index][:size] * energies).sum() / total),
            ))
        per_frame.append(clusters)
    return per_frame


def collect_crops(frames, clusters_per_frame, beam_energy, selection_range):
    """Sort clusters into 1/2/3-photon crops, following UsefulFuncs._processFrames."""
    frame_size = frames.shape[-1]
    crops = {n: [] for n in CLUSTER_SIZE}
    origins = {n: [] for n in CLUSTER_SIZE}     # (frame_index, ref_col, ref_row)
    for frame_index, clusters in enumerate(clusters_per_frame):
        for cluster in clusters:
            energy = cluster["energy"]
            for n_photons, size in CLUSTER_SIZE.items():
                window = n_photons * selection_range
                if not (n_photons * beam_energy - window < energy
                        < n_photons * beam_energy + window):
                    continue
                if n_photons == 1:
                    ref_col = cluster["seed_col"] - size // 2
                    ref_row = cluster["seed_row"] - size // 2
                else:
                    ref_col = int(cluster["centre_col"] - size / 2) + 1
                    ref_row = int(cluster["centre_row"] - size / 2) + 1
                if (ref_col < 0 or ref_col + size > frame_size
                        or ref_row < 0 or ref_row + size > frame_size):
                    continue
                crops[n_photons].append(
                    frames[frame_index, ref_row:ref_row + size, ref_col:ref_col + size])
                origins[n_photons].append((frame_index, ref_col, ref_row))
    return ({n: (np.asarray(v, np.float32) if v else np.zeros((0, CLUSTER_SIZE[n],
                                                               CLUSTER_SIZE[n]), np.float32))
             for n, v in crops.items()},
            {n: (np.asarray(v, np.int64) if v else np.zeros((0, 3), np.int64))
             for n, v in origins.items()})


@torch.no_grad()
def cluster_dl_points(crops, origins, n_frames, device="cuda", batch_size=8192):
    """Run the pretrained 1/2/3-photon CNNs -> per-frame photon lists."""
    points = [[] for _ in range(n_frames)]
    for n_photons, size in CLUSTER_SIZE.items():
        sample = crops[n_photons]
        if len(sample) == 0:
            continue
        model = load_pretrained(n_photons, device)
        centre = size / 2.0
        extra = coordinate_channels(size).to(device) if n_photons == 1 else None
        for start in range(0, len(sample), batch_size):
            batch = torch.from_numpy(sample[start:start + batch_size]).to(device)
            batch = batch.unsqueeze(1)
            if extra is not None:
                batch = torch.cat([batch, extra.expand(batch.shape[0], -1, -1, -1)], dim=1)
            output = model(batch).cpu().numpy().reshape(len(batch), n_photons, 2)
            for row, (frame_index, ref_col, ref_row) in enumerate(
                    origins[n_photons][start:start + batch_size]):
                for photon in range(n_photons):
                    points[frame_index].append((
                        ref_col + centre + output[row, photon, 0],
                        ref_row + centre + output[row, photon, 1],
                        float(n_photons and 12.0), 1.0))
        del model
        torch.cuda.empty_cache()
    return [np.asarray(p, np.float64).reshape(-1, 4) for p in points]


def cluster_eta_points(crops, origins, n_frames):
    """Classical eta interpolation on the 1-photon clusters."""
    sample = crops[1]
    points = [[] for _ in range(n_frames)]
    if len(sample) == 0:
        return [np.zeros((0, 4)) for _ in range(n_frames)]
    totals = sample.sum(axis=(1, 2))
    totals[totals == 0] = 1.0
    eta_x = np.clip((sample.sum(axis=1) * ETA_POSITION_WEIGHTS).sum(axis=-1) / totals, 0, 1)
    eta_y = np.clip((sample.sum(axis=2) * ETA_POSITION_WEIGHTS).sum(axis=-1) / totals, 0, 1)
    histogram, _, _ = np.histogram2d(eta_x, eta_y, bins=ETA_BINS, range=[[0, 1], [0, 1]])
    x_table, y_table = build_xy_lut_Rosenblatt(histogram)
    x_fraction, y_fraction = bilinear_xy_lookup(eta_x, eta_y, x_table, y_table)
    for row, (frame_index, ref_col, ref_row) in enumerate(origins[1]):
        points[frame_index].append((ref_col + 1.0 + x_fraction[row],
                                    ref_row + 1.0 + y_fraction[row], 12.0, 1.0))
    return [np.asarray(p, np.float64).reshape(-1, 4) for p in points]


@torch.no_grad()
def grid_points(model, frames, train_conf, halo, threshold, merge_distance,
                batch_size=512):
    input_scale = float(train_conf.data.input_scale)
    energy_scale = float(train_conf.data.energy_scale)
    points = []
    for start in range(0, len(frames), batch_size):
        batch = torch.from_numpy(frames[start:start + batch_size]).cuda().unsqueeze(1)
        prediction = model(batch / input_scale).float()
        points.extend(p.numpy() for p in decode(
            prediction, threshold=threshold, merge_distance=merge_distance,
            halo=halo, predict_energy=bool(train_conf.model.predict_energy),
            energy_scale=energy_scale))
    return points


def score(points_per_frame, truth_per_frame, halo, frame_size, match_radius,
          pixel_size_um):
    stats = DetectionStats(pixel_size_um=pixel_size_um, match_radius=match_radius,
                           halo=halo, frame_size=frame_size)
    stats.update(points_per_frame, truth_per_frame)
    return stats.summary()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--data", default=None, help="eval_scan folder")
    parser.add_argument("--n-frames", type=int, default=4000)
    parser.add_argument("--beam-energy", type=float, default=12.0)
    parser.add_argument("--selection-range", type=float, default=3.0)
    parser.add_argument("--outdir", default=None)
    args = parser.parse_args()

    import h5py
    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    train_conf = OmegaConf.create(checkpoint["conf"])
    model = build_model(train_conf).cuda().eval()
    model.load_state_dict(checkpoint["model"])
    halo = model.rf_radius
    threshold = float(train_conf.eval.threshold)
    merge_distance = float(train_conf.eval.merge_distance)
    match_radius = float(train_conf.eval.match_radius)
    pixel_size_um = float(train_conf.eval.pixel_size_um)
    noise_keV = 0.13
    print(f"[baseline] grid model {args.checkpoint} (epoch {checkpoint['epoch']}), "
          f"R={halo}, threshold {threshold}, merge {merge_distance} px")

    data_folder = Path(args.data or Path(train_conf.data.folder) / "eval_scan")
    files = sorted(data_folder.glob("*.h5"),
                   key=lambda p: float(p.stem.replace("occ", "").replace("pct", "")
                                       .replace("p", ".")))
    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 = float(path.stem.replace("occ", "").replace("pct", "")
                          .replace("p", ".")) / 100
        with h5py.File(path, "r") as handle:
            frames = handle["frames"][:args.n_frames].astype(np.float32)
            pointer = handle["frame_ptr"][:args.n_frames + 1]
            photons = handle["photons"][:pointer[-1]]
        truth = [photons[pointer[i]:pointer[i + 1]] for i in range(len(frames))]
        frame_size = frames.shape[-1]

        clusters_per_frame = find_clusters(frames.astype(np.float64), noise_keV)
        crops, origins = collect_crops(frames, clusters_per_frame,
                                       args.beam_energy, args.selection_range)
        chains = {
            "grid": grid_points(model, frames, train_conf, halo, threshold,
                                merge_distance),
            "cluster+dl": cluster_dl_points(crops, origins, len(frames)),
            "cluster+eta": cluster_eta_points(crops, origins, len(frames)),
        }
        results[f"{occupancy*100:g}%"] = dict(
            occupancy=occupancy,
            n_clusters={str(n): int(len(crops[n])) for n in CLUSTER_SIZE},
            **{name: score(pts, truth, halo, frame_size, match_radius, pixel_size_um)
               for name, pts in chains.items()})

        print(f"\n  occupancy {occupancy*100:g}%  ({len(frames)} frames, clusters "
              f"1ph={len(crops[1])} 2ph={len(crops[2])} 3ph={len(crops[3])})")
        print(f"    {'chain':12s} {'eff%':>7} {'pur%':>7} {'RMS um':>8} {'s68 um':>8}")
        for name in chains:
            s = results[f"{occupancy*100:g}%"][name]
            print(f"    {name:12s} {s['efficiency']*100:7.2f} {s['purity']*100:7.2f} "
                  f"{s['rms_1d_um']:8.3f} {s['sigma68_1d_um']:8.3f}", flush=True)

    with open(output_dir / "baseline_comparison.json", "w") as handle:
        json.dump(results, handle, indent=1)
    plot(results, output_dir, checkpoint["epoch"])
    print(f"\n[baseline] wrote {output_dir/'baseline_comparison.json'}")


def plot(results, output_dir, epoch):
    occupancy = np.array([r["occupancy"] for r in results.values()]) * 100
    figure, axes = plt.subplots(1, 3, figsize=(15, 4.4))
    styles = {"grid": ("o-", "#10695f"), "cluster+dl": ("s--", "#4b2a6b"),
              "cluster+eta": ("^:", "#a0451c")}
    for name, (style, colour) in styles.items():
        efficiency = np.array([r[name]["efficiency"] for r in results.values()]) * 100
        purity = np.array([r[name]["purity"] for r in results.values()]) * 100
        rms = np.array([r[name]["rms_1d_um"] for r in results.values()])
        axes[0].plot(occupancy, efficiency, style, color=colour, label=name)
        axes[1].plot(occupancy, purity, style, color=colour, label=name)
        axes[2].plot(occupancy, rms, style, color=colour, label=name)
    for axis, label in zip(axes, ("efficiency [%]", "purity [%]",
                                  r"position RMS [$\mu$m]")):
        axis.set_xlabel("occupancy [%]")
        axis.set_ylabel(label)
        axis.set_xscale("log")
        axis.grid(alpha=.3)
        axis.legend()
    figure.suptitle(f"same MC frames, same matcher (grid model at epoch {epoch})")
    figure.tight_layout()
    figure.savefig(output_dir / "Plots" / "baseline_comparison.png", dpi=130)
    plt.close(figure)


if __name__ == "__main__":
    main()