#!/usr/bin/env python3
"""
WP5.2 -- run the grid model over a measured run, straight from the raw files.

    python scripts/05_infer_measurement.py --config Configs/infer_measurement_12keV.yaml \
        --run Edge_3filters_pos0_12keV
    python scripts/05_infer_measurement.py --config ... --run ... --n-frames 200000
    python scripts/05_infer_measurement.py --config ... --run ... --save-photons

raw ADU -> pedestal -> common mode -> 3D-LUT -> GridNet -> position-space merge
-> sub-pixel binned image + energy spectrum + per-frame photon count.

Only the valid region (the frame minus the halo R) is kept: near the ROI border
the network sees zero padding instead of the neighbouring pixels that exist on
the chip, so those cells are not trustworthy.

With `inference.precision: fp16` the network runs under FP16 autocast in
channels-last layout, 2.3x faster for the poly model.  The offsets and energy of
the cells above threshold are then recomputed in FP32 from the same FP16
features (`predict_fp16`), so energies are not rounded to the FP16 grid.  The models are trained
with AMP, and on MC and measured frames FP16 reproduces the FP32 photon list
(efficiency, purity, position and energy resolution unchanged; 99.997 % of the
photons identical, positions within ~20 nm).

Images written per run, each as .npy (analysis) and .png (looking at it):
    binned_image   reconstructed photons at pixel_size / binning_factor
    native_image   the same photons rebinned back to the 25 um pitch
    analog_image   mean calibrated energy per frame -- the classical "native
                   readout" picture
plus comparison.png, the two side by side.

With `inference.self_flattening` the sub-pixel phase of the reconstructed
positions is flattened with the Rosenblatt transform of
`DeepLearning/src/etaInterpolationFromPoints.py`: the (frac_x, frac_y) histogram
of the whole run becomes a look-up table that maps every position onto the one a
uniform illumination would have put it at.  A per-bin flat-field division cannot
do this -- it corrects intensities, while a grid artefact is a displacement --
so the flattened image is produced here alongside the raw one and both are kept.
"""
import argparse
import json
import math
import os
import queue
import sys
import threading
import time
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 measurement import (GpuCalibrator, ImageAccumulator,            # noqa: E402
                         RawFrameReader, estimate_pedestal)
from models import build_model                                       # noqa: E402
from etaInterpolationFromPoints import (NEtaBins,                    # noqa: E402
                                        interpolate_eta_from_points)
from flattening import StreamingFlattener                           # noqa: E402


@torch.no_grad()
def predict_fp16(model, inputs, threshold):
    """FP16 backbone and head; offsets and energy of candidate cells in FP32.

    An FP16 head rounds every energy to the FP16 grid (7.3 meV below 15 keV,
    14.6 meV above), which shows up as a comb in a 50 eV spectrum.  Running the
    whole head in FP32 means converting the full feature map and costs ~15 %,
    so only the cells above the objectness threshold -- the only ones decode can
    report -- get channels 1-3 recomputed in FP32 from the same FP16 features.
    """
    with torch.autocast("cuda", dtype=torch.float16):
        features = model.backbone(inputs.contiguous(memory_format=torch.channels_last))
        prediction = model.head(features)
    prediction = prediction.float()
    logit_threshold = math.log(threshold / (1.0 - threshold))
    batch_of, row_of, column_of = torch.nonzero(prediction[:, 0] > logit_threshold,
                                                as_tuple=True)
    if len(batch_of):
        weight = model.head.weight[1:, :, 0, 0].float()          # (3, channels)
        bias = model.head.bias[1:].float()
        cells = features[batch_of, :, row_of, column_of].float()  # (n, channels)
        values = cells @ weight.T + bias                          # (n, 3)
        for channel in range(values.shape[1]):
            prediction[batch_of, channel + 1, row_of, column_of] = values[:, channel]
    return prediction


def load_model(checkpoint_path):
    checkpoint = torch.load(checkpoint_path, 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"])
    return model, train_conf, checkpoint


def prefetch(iterator, depth=4):
    """Overlap file reading with GPU work."""
    output = queue.Queue(maxsize=depth)

    def worker():
        try:
            for item in iterator:
                output.put(item)
        finally:
            output.put(None)

    threading.Thread(target=worker, daemon=True).start()
    while True:
        item = output.get()
        if item is None:
            return
        yield item


def exposed_occupancy(density, fraction=0.5):
    """Occupancy inside the illuminated area only.

    A knife-edge or slit run leaves most of the ROI in shadow, so an average
    over the whole frame badly understates the flux the model actually sees:
    on Edge_3filters_pos0_12keV it gives 0.16 % against the 0.67 % the
    illuminated pixels really see.
    """
    if density.size == 0 or density.max() <= 0:
        return float("nan"), 0
    reference = np.percentile(density, 99)
    exposed = density > fraction * reference
    if not exposed.any():
        return float("nan"), 0
    return float(density[exposed].mean()), int(exposed.sum())


def subpixel_histogram(x, y, bins=NEtaBins, chunk=50_000_000):
    """2D histogram of the sub-pixel phase of the reconstructed positions.

    In slices: histogram2d stacks its input as float64 and adds int64 bin
    indices, ~40 bytes per position, which ran out of memory on the 1e9-photon
    flat fields.  The histogram is additive, so the result is the same.
    """
    histogram = np.zeros((bins, bins))
    for start in range(0, len(x), chunk):
        xs, ys = x[start:start + chunk], y[start:start + chunk]
        histogram += np.histogram2d(xs - np.floor(xs), ys - np.floor(ys),
                                    bins=bins, range=[[0, 1], [0, 1]])[0]
    return histogram


def subpixel_nonuniformity(histogram):
    """Excess rms/mean over the Poisson floor; 0 means a flat sub-pixel phase."""
    mean = histogram.mean()
    if mean <= 0:
        return float("nan")
    relative = histogram.std() / mean
    return float(np.sqrt(max(relative ** 2 - 1.0 / mean, 0.0)))


def save_subpixel_figure(output_dir, before, after, run_name):
    panels = [("before flattening", before)]
    if after is not None:
        panels.append(("after flattening", after))
    figure, axes = plt.subplots(1, len(panels) + 1,
                                figsize=(5.2 * (len(panels) + 1), 4.2))
    for axis, (title, histogram) in zip(axes, panels):
        image = axis.imshow(histogram.T / histogram.mean(), origin="lower",
                            extent=[0, 1, 0, 1], vmin=0.9, vmax=1.1, cmap="viridis")
        axis.set_title(f"{title}\nexcess rms "
                       f"{subpixel_nonuniformity(histogram)*100:.2f} %", fontsize=10)
        axis.set_xlabel("sub-pixel x")
        axis.set_ylabel("sub-pixel y")
        plt.colorbar(image, ax=axis, fraction=0.046)
    axis = axes[-1]
    for title, histogram in panels:
        profile = histogram.sum(axis=1)
        axis.plot(np.linspace(0, 1, len(profile)), profile / profile.mean(),
                  lw=1, label=title)
    axis.axhline(1.0, color="k", lw=.6)
    axis.set_xlabel("sub-pixel x")
    axis.set_ylabel("relative density")
    axis.legend(fontsize=8)
    axis.grid(alpha=.3)
    figure.suptitle(run_name)
    figure.tight_layout()
    figure.savefig(output_dir / "subpixel_phase.png", dpi=150)
    plt.close(figure)


def save_images(output_dir, binned, native, analog, run_name, binning_factor,
                pixel_size_um, halo):
    """Write the reconstruction and the analog picture as .npy and .png."""
    np.save(output_dir / "binned_image.npy", binned)
    np.save(output_dir / "native_image.npy", native)
    np.save(output_dir / "analog_image.npy", analog)

    bin_um = pixel_size_um / binning_factor
    extent_binned = [0, binned.shape[1] * bin_um, 0, binned.shape[0] * bin_um]
    extent_native = [0, native.shape[1] * pixel_size_um,
                     0, native.shape[0] * pixel_size_um]

    def draw(axis, image, extent, title, unit):
        finite = image[np.isfinite(image)]
        high = np.percentile(finite, 99.5) if finite.size else 1.0
        handle = axis.imshow(image, origin="lower", extent=extent, cmap="gray",
                             vmin=0, vmax=high if high > 0 else 1.0,
                             interpolation="nearest")
        axis.set_title(title, fontsize=10)
        axis.set_xlabel(r"x [$\mu$m]")
        axis.set_ylabel(r"y [$\mu$m]")
        plt.colorbar(handle, ax=axis, label=unit, fraction=0.046)

    panels = (
        (binned, extent_binned, "binned_image",
         f"reconstruction, {bin_um:g} " + r"$\mu$m bins", "photons"),
        (native, extent_native, "native_image",
         f"reconstructed photons, {pixel_size_um:g} " + r"$\mu$m pitch", "photons"),
        (analog, extent_native, "analog_image",
         "native readout (mean energy per frame)", "keV/frame"),
    )
    for image, extent, name, title, unit in panels:
        figure, axis = plt.subplots(figsize=(6.4, 5.4))
        draw(axis, image, extent, f"{run_name}\n{title}", unit)
        figure.tight_layout()
        figure.savefig(output_dir / f"{name}.png", dpi=150)
        plt.close(figure)

    figure, axes = plt.subplots(1, 2, figsize=(12.4, 5.4))
    draw(axes[0], analog, extent_native,
         "native readout (mean energy per frame)", "keV/frame")
    draw(axes[1], binned, extent_binned,
         f"GridBased reconstruction, {bin_um:g} " + r"$\mu$m bins", "photons")
    figure.suptitle(f"{run_name}   (valid region = frame minus halo R={halo})")
    figure.tight_layout()
    figure.savefig(output_dir / "comparison.png", dpi=150)
    plt.close(figure)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    parser.add_argument("--run", required=True)
    parser.add_argument("--n-frames", type=int, default=None)
    parser.add_argument("--checkpoint", default=None)
    parser.add_argument("--tag", default=None)
    parser.add_argument("--save-photons", action="store_true",
                        help="write photons.npy (x, y, E, p) even if the config says no")
    args = parser.parse_args()

    conf = OmegaConf.load(args.config)
    if args.run not in conf.runs:
        sys.exit(f"unknown run {args.run!r}; have {list(conf.runs.keys())}")
    run_conf = conf.runs[args.run]

    # the stored frame is n_x x n_y; detector.crop = [x0, x1, y0, y1] reads and
    # reconstructs only that window of it (the halo R is lost at its border, so
    # the crop should reach R pixels past the region of interest)
    frame_n_x, frame_n_y = int(conf.detector.n_x), int(conf.detector.n_y)
    crop = conf.detector.get("crop", None)
    crop = None if crop is None else [int(v) for v in crop]
    if crop is None:
        n_x, n_y, crop_x0, crop_y0 = frame_n_x, frame_n_y, 0, 0
    else:
        n_x, n_y, crop_x0, crop_y0 = crop[1] - crop[0], crop[3] - crop[2], crop[0], crop[2]
    pixel_size_um = float(conf.detector.pixel_size_um)
    checkpoint_path = args.checkpoint or conf.model.checkpoint
    model, train_conf, checkpoint = load_model(checkpoint_path)
    precision = str(conf.inference.get("precision", "fp32"))
    if precision not in ("fp32", "fp16"):
        sys.exit(f"unknown inference.precision {precision!r}")
    if precision == "fp16":
        model = model.to(memory_format=torch.channels_last)
        torch.backends.cudnn.benchmark = True
    halo = model.rf_radius
    input_scale = float(train_conf.data.input_scale)
    energy_scale = float(train_conf.data.energy_scale)
    print(f"[infer] run {args.run}")
    print(f"[infer] model {checkpoint_path} (epoch {checkpoint['epoch']}), "
          f"R={halo} -> valid region {n_x-2*halo}x{n_y-2*halo} of {n_x}x{n_y}")

    first_file, last_file = run_conf.signal_files
    file_step = int(run_conf.get("signal_file_step", 1))
    signal_files = [run_conf.signal_pattern.format(i)
                    for i in range(int(first_file), int(last_file) + 1, file_step)]
    reader = RawFrameReader(signal_files, frame_n_x, frame_n_y,
                            header_words=int(conf.detector.header_words), crop=crop)
    n_frames = args.n_frames or conf.inference.n_frames or reader.n_frames_available
    n_frames = min(int(n_frames), reader.n_frames_available)

    corner = (int(conf.detector.corner[0]) + crop_x0, int(conf.detector.corner[1]) + crop_y0)
    calibrator = GpuCalibrator(conf.calibration.file, n_x, n_y, corner=corner,
                               common_mode=conf.calibration.common_mode)
    if conf.calibration.pedestal_source == "pedestal_run":
        pedestal_reader = RawFrameReader([run_conf.pedestal], frame_n_x, frame_n_y,
                                         header_words=int(conf.detector.header_words),
                                         crop=crop)
        pedestal = estimate_pedestal(pedestal_reader,
                                     n_frames=int(conf.calibration.pedestal_frames))
    elif conf.calibration.pedestal_source == "per_file":
        pedestal = estimate_pedestal(reader,
                                     n_frames=int(conf.calibration.pedestal_frames))
    else:
        sys.exit(f"unknown pedestal_source {conf.calibration.pedestal_source!r}")
    calibrator.set_pedestal(pedestal)
    print(f"[infer] pedestal from {conf.calibration.pedestal_source}, "
          f"common mode '{conf.calibration.common_mode}', "
          f"{n_frames} frames of {reader.n_frames_available} "
          f"({len(signal_files)} files, step {file_step})")
    if crop is not None:
        print(f"[infer] crop x {crop[0]}-{crop[1]}, y {crop[2]}-{crop[3]} of the "
              f"{frame_n_x}x{frame_n_y} frame -> {n_x}x{n_y}, valid region "
              f"x {crop[0]+halo}-{crop[1]-halo}, y {crop[2]+halo}-{crop[3]-halo}")

    binning_factor = int(conf.inference.binning_factor)
    accumulator = ImageAccumulator(n_x, n_y, binning_factor=binning_factor, halo=halo)
    energy_sum = torch.zeros((n_y, n_x), dtype=torch.float64, device="cuda")
    energy_low, energy_high, energy_n = conf.inference.energy_bins
    energy_edges = np.linspace(float(energy_low), float(energy_high), int(energy_n) + 1)
    energy_histogram = np.zeros(int(energy_n), np.int64)
    photons_per_frame = []
    save_photons = args.save_photons or bool(conf.inference.save_photons)
    photon_dump = [] if save_photons else None
    # self_flattening: true keeps every position and remaps it (the published
    # runs); "streaming" histograms them on the GPU instead (src/flattening.py),
    # for runs whose photon list would not fit in memory
    flattening_mode = conf.inference.get("self_flattening", False)
    streaming = str(flattening_mode).lower() == "streaming"
    self_flattening = bool(flattening_mode) and not streaming
    position_dump = [] if self_flattening else None
    occupancy_edges = conf.inference.get("flattening_occupancy_edges", None)
    flattener = (StreamingFlattener(
        n_x, n_y, binning_factor, halo=halo,
        lut_photons=int(conf.inference.get("flattening_lut_photons", 100_000_000)),
        occupancy_edges=(None if occupancy_edges is None
                         else [float(e) for e in occupancy_edges]))
        if streaming else None)

    output_dir = Path(conf.inference.output_folder) / (
        args.run + (f"_{args.tag}" if args.tag else ""))
    output_dir.mkdir(parents=True, exist_ok=True)

    batch_size = int(conf.inference.batch_size)
    threshold = float(conf.model.threshold)
    merge_distance = float(conf.model.merge_distance)
    started = time.time()
    n_done = 0

    batches = reader.batches(batch_size, n_frames=n_frames)
    with torch.no_grad():
        for file_index, frame_index, adu_frames in prefetch(batches):
            if len(adu_frames) == 0:
                continue
            energy_frames = calibrator.to_energy(torch.from_numpy(adu_frames).cuda())
            energy_sum += energy_frames.sum(dim=0, dtype=torch.float64)
            inputs = (energy_frames / input_scale).unsqueeze(1)
            if precision == "fp16":
                prediction = predict_fp16(model, inputs, threshold)
            else:
                prediction = model(inputs).float()
            points_per_frame = decode(prediction, threshold=threshold,
                                      merge_distance=merge_distance, halo=halo,
                                      predict_energy=bool(train_conf.model.predict_energy),
                                      energy_scale=energy_scale)
            points_per_frame = [p.numpy() for p in points_per_frame]
            accumulator.add(points_per_frame)
            photons_per_frame.extend(len(p) for p in points_per_frame)
            batch_energies = [p[:, 2] for p in points_per_frame if len(p)]
            if batch_energies:
                energy_histogram += np.histogram(np.concatenate(batch_energies),
                                                 bins=energy_edges)[0]
            if photon_dump is not None:
                photon_dump.extend(points_per_frame)
            if position_dump is not None or flattener is not None:
                # one array per batch, not per frame: 14 M small arrays would
                # cost more in list overhead than the positions themselves
                batch_points = [p[:, :2] for p in points_per_frame if len(p)]
                if batch_points:
                    batch_points = np.concatenate(batch_points).astype(np.float32)
                    if flattener is not None:
                        flattener.add(batch_points, n_frames=len(adu_frames))
                    else:
                        position_dump.append(batch_points)

            n_done += len(adu_frames)
            if n_done % (batch_size * 200) < batch_size:
                elapsed = time.time() - started
                rate = n_done / elapsed
                print(f"  {n_done}/{n_frames} frames  {rate:.0f} frame/s  "
                      f"eta {(n_frames-n_done)/rate/60:.1f} min  "
                      f"<photons/frame> {np.mean(photons_per_frame):.2f}", flush=True)

    elapsed = time.time() - started
    photons_per_frame = np.asarray(photons_per_frame)

    binned = accumulator.image_2d
    native = binned.reshape(n_y, binning_factor, n_x, binning_factor).sum(axis=(1, 3))
    analog = (energy_sum / max(n_done, 1)).cpu().numpy()

    phase_before = phase_after = None
    flattening_classes = None
    if self_flattening and position_dump:
        positions = np.concatenate(position_dump)
        position_dump.clear()
        phase_before = subpixel_histogram(positions[:, 0], positions[:, 1])
        print(f"[infer] self-flattening {len(positions)} positions, sub-pixel "
              f"excess rms {subpixel_nonuniformity(phase_before)*100:.2f} %", flush=True)
        corrected = interpolate_eta_from_points(positions)
        phase_after = subpixel_histogram(corrected[:, 0], corrected[:, 1])
        print(f"[infer] after flattening, sub-pixel excess rms "
              f"{subpixel_nonuniformity(phase_after)*100:.2f} %", flush=True)
        flattened_accumulator = ImageAccumulator(n_x, n_y,
                                                 binning_factor=binning_factor,
                                                 halo=halo)
        # in slices: one call on 1e9 positions builds three int64 index arrays
        # of the full length (~28 GB at the 0-filter flat field)
        for start in range(0, len(corrected), 50_000_000):
            flattened_accumulator.add([corrected[start:start + 50_000_000]])
        np.save(output_dir / "binned_image_flattened.npy",
                flattened_accumulator.image_2d)
        np.save(output_dir / "subpixel_phase_before.npy", phase_before)
        np.save(output_dir / "subpixel_phase_after.npy", phase_after)
        save_subpixel_figure(output_dir, phase_before, phase_after, args.run)
        del positions, corrected, flattened_accumulator
    if flattener is not None and flattener.n_photons:
        flattened, phase_before, phase_after = flattener.finalize()
        print(f"[infer] streaming self-flattening of {flattener.n_photons} positions "
              f"(table from the first {flattener.table_photons}): sub-pixel excess rms "
              f"{subpixel_nonuniformity(phase_before)*100:.2f} % -> "
              f"{subpixel_nonuniformity(phase_after)*100:.2f} %", flush=True)
        for c in flattener.classes:
            print(f"[infer]   occupancy {c['occupancy_from']*100:5.2f}-{c['occupancy_to']*100:6.2f} %: "
                  f"{c['pixels']:6d} px, table from {c['table_photons']} photons", flush=True)
        flattening_classes = flattener.classes
        np.save(output_dir / "binned_image_flattened.npy", flattened)
        np.save(output_dir / "subpixel_phase_before.npy", phase_before)
        np.save(output_dir / "subpixel_phase_after.npy", phase_after)
        save_subpixel_figure(output_dir, phase_before, phase_after, args.run)
        del flattened, flattener

    valid = (slice(halo, n_y - halo), slice(halo, n_x - halo))
    density = native[valid] / max(n_done, 1)        # photons per pixel per frame
    occupancy_exposed, n_exposed = exposed_occupancy(density)
    n_valid_pixels = (n_x - 2 * halo) * (n_y - 2 * halo)

    save_images(output_dir, binned, native, analog, args.run, binning_factor,
                pixel_size_um, halo)
    np.save(output_dir / "energy_spectrum.npy",
            np.stack([0.5 * (energy_edges[:-1] + energy_edges[1:]),
                      energy_histogram.astype(float)]))
    np.save(output_dir / "photons_per_frame.npy", photons_per_frame)
    if photon_dump is not None:
        np.save(output_dir / "photons.npy", np.concatenate(photon_dump))

    summary = dict(
        run=args.run, checkpoint=str(checkpoint_path), epoch=int(checkpoint["epoch"]),
        n_frames=int(n_done), n_photons=int(accumulator.n_photons),
        photons_per_frame=float(photons_per_frame.mean()),
        occupancy_valid_region=float(photons_per_frame.mean() / n_valid_pixels),
        occupancy_exposed=occupancy_exposed,
        exposed_pixels=n_exposed, valid_pixels=n_valid_pixels,
        threshold=threshold, merge_distance=merge_distance, halo=halo,
        common_mode=str(conf.calibration.common_mode),
        pedestal_source=str(conf.calibration.pedestal_source),
        binning_factor=binning_factor,
        crop=crop, frame_size=[frame_n_x, frame_n_y],
        signal_files=[int(first_file), int(last_file), file_step],
        self_flattening="streaming" if streaming else self_flattening,
        flattening_classes=flattening_classes,
        precision=precision,
        subpixel_nonuniformity_before=(subpixel_nonuniformity(phase_before)
                                       if phase_before is not None else None),
        subpixel_nonuniformity_after=(subpixel_nonuniformity(phase_after)
                                      if phase_after is not None else None),
        seconds=elapsed, frames_per_second=float(n_done / elapsed),
    )
    with open(output_dir / "summary.json", "w") as handle:
        json.dump(summary, handle, indent=1)

    print(f"[infer] {n_done} frames in {elapsed:.0f}s "
          f"({n_done/elapsed:.0f} frame/s), {accumulator.n_photons} photons, "
          f"{summary['photons_per_frame']:.2f}/frame")
    print(f"[infer] occupancy: {summary['occupancy_valid_region']*100:.3f}% averaged "
          f"over the valid region, {occupancy_exposed*100:.3f}% inside the "
          f"illuminated area ({n_exposed} of {n_valid_pixels} px)")
    print(f"[infer] wrote {output_dir}")


if __name__ == "__main__":
    main()