#!/usr/bin/env python3
"""
WP1.1 -- build the alpha(E, t) / beta(E, t) look-up table for continuous-energy
charge-cloud sampling.

    python scripts/01_build_charge_lut.py --config Configs/charge_lut_Moench014.yaml

For every energy the charge-transport simulation gives the transverse carrier
distribution in 512 depth bins.  Each bin is fitted with a generalised normal
distribution exactly as ``McGeneration/McGenerator.py::parameterization`` does,
so the shapes stay identical to the ones validated against measurement.

What is new is what happens next.  ``parameterization`` compresses each energy
into 9 numbers by fitting alpha(t) with a 4-parameter polynomial and beta(t)
with a 5-parameter form; those 9 numbers are strongly correlated and
individually unstable, so interpolating them in energy produces unphysical
curves in between.  Here the fitted *curves* are tabulated instead: the drift
time depends only on the sensor, bias and temperature, so the depth grid is the
same at every energy and (energy, drift time) is already a regular grid --
alpha and beta can simply be looked up and interpolated bilinearly, with no
functional form in between.

Output ``data/charge_lut_<sensor>_<hv>V.npz``:
    energies    (nE,)      keV
    z0          (nZ,)      um, conversion depth measured from the entrance
    drift_time  (nZ,)      ns
    alpha, beta (nE, nZ)   per-depth GGD fit
    alpha_err, beta_err    fit errors
    n_carriers  (nE, nZ)   entries behind each fit
"""
import argparse
import os
import sys
import time
from multiprocessing import Pool

import numpy as np
from array import array
from omegaconf import OmegaConf

from ROOT import TF1, TH1F, TH1, gROOT
gROOT.SetBatch(True)
TH1.AddDirectory(False)

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "src"))
from charge_model import drift_time_allpix2  # noqa: E402

_CONF = None


def fit_generalised_normal(coords, index):
    """The per-depth GGD fit of McGenerator.parameterization, unchanged."""
    rms = np.std(coords)
    if not np.isfinite(rms) or rms <= 0:
        return np.nan, np.nan, np.nan, np.nan
    bin_width = rms / 20
    n_bins = int(100 / bin_width)
    histogram = TH1F(f"_h1_{index}", "", n_bins, -rms * 5, rms * 5)
    histogram.FillN(len(coords), array("d", coords), array("d", [1] * len(coords)))
    function = TF1(
        f"f1_{index}",
        "[0]/(2*[1]*TMath::Gamma(1/[0])) * exp(-pow(TMath::Abs(x)/[1], [0]))",
        -5 * rms, 5 * rms)                      # par[0]: beta, par[1]: alpha
    function.SetParLimits(0, 2, 5)
    function.SetParLimits(1, .5, 30)
    function.SetParameters(2., rms * 1.5)
    histogram.Scale(histogram.GetNbinsX() / (10 * rms) / histogram.Integral())
    histogram.Fit(function, "Q")
    beta, alpha = function.GetParameter(0), function.GetParameter(1)
    beta_err, alpha_err = function.GetParError(0), function.GetParError(1)
    del histogram, function
    return alpha, beta, alpha_err, beta_err


def process_energy(energy_keV):
    conf = _CONF
    folder = conf.simulation.folder.format(sensor=conf.sensor.code,
                                           hv=conf.sensor.hv, energy=energy_keV)
    n_bins = int(conf.simulation.z_bins)
    method = conf.simulation.method
    started = time.time()
    alpha = np.full(n_bins, np.nan)
    beta = np.full(n_bins, np.nan)
    alpha_err = np.full(n_bins, np.nan)
    beta_err = np.full(n_bins, np.nan)
    n_carriers = np.zeros(n_bins, np.int64)
    for index in range(n_bins):
        stem = (f"{conf.sensor.code}_{conf.sensor.hv}V_{{axis}}_{energy_keV}keV_"
                f"{index}_of_{n_bins}_{method}.npy")
        try:
            xs = np.load(os.path.join(folder, stem.format(axis="xs")))
            ys = np.load(os.path.join(folder, stem.format(axis="ys")))
        except FileNotFoundError:
            continue
        coords = np.concatenate((xs, ys))
        n_carriers[index] = len(coords)
        (alpha[index], beta[index],
         alpha_err[index], beta_err[index]) = fit_generalised_normal(coords, index)
    print(f"  {energy_keV:2d} keV: {np.isfinite(alpha).sum()}/{n_bins} depth bins fitted "
          f"({time.time()-started:.0f}s)", flush=True)
    return energy_keV, alpha, beta, alpha_err, beta_err, n_carriers


def initialise(conf):
    global _CONF
    _CONF = conf


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    parser.add_argument("--workers", type=int, default=8)
    args = parser.parse_args()
    conf = OmegaConf.load(args.config)

    energies = list(range(int(conf.energies.first), int(conf.energies.last) + 1,
                          int(conf.energies.step)))
    n_bins = int(conf.simulation.z_bins)
    thickness = float(conf.sensor.thickness_um)
    edges = np.linspace(0, thickness, n_bins + 1)
    z0 = 0.5 * (edges[:-1] + edges[1:])
    drift_time = np.array([drift_time_allpix2(
        depth, thickness_um=thickness, hv=float(conf.sensor.hv),
        depletion_voltage=float(conf.sensor.depletion_voltage),
        temperature=float(conf.sensor.temperature)) for depth in z0])

    print(f"[lut] {conf.sensor.code} @{conf.sensor.hv} V, {len(energies)} energies "
          f"{energies[0]}-{energies[-1]} keV, {n_bins} depth bins")
    print(f"[lut] drift time {drift_time.min():.2f} - {drift_time.max():.2f} ns")

    with Pool(processes=args.workers, initializer=initialise, initargs=(conf,)) as pool:
        results = pool.map(process_energy, energies)

    order = {e: i for i, e in enumerate(energies)}
    shape = (len(energies), n_bins)
    alpha, beta = np.full(shape, np.nan), np.full(shape, np.nan)
    alpha_err, beta_err = np.full(shape, np.nan), np.full(shape, np.nan)
    n_carriers = np.zeros(shape, np.int64)
    for energy, a, b, ae, be, nc in results:
        i = order[energy]
        alpha[i], beta[i], alpha_err[i], beta_err[i], n_carriers[i] = a, b, ae, be, nc

    # Depth bins the simulation never produced (e.g. the deepest bins of an
    # energy whose run was cut short) can be filled from the neighbouring
    # energies at the *same* depth bin, i.e. the same drift time.  The cloud size
    # is set mainly by the drift time, so this is far closer than
    # ChargeCloudLUT._fill_gaps, which holds the last fitted depth constant.
    filled = ~np.isfinite(alpha) | ~np.isfinite(beta)
    if conf.get("fill_missing", None) == "energy" and filled.any():
        grid = np.array(energies, float)
        for j in np.flatnonzero(filled.any(axis=0)):
            good = ~filled[:, j]
            if good.sum() < 2:
                continue
            for surface in (alpha, beta):
                surface[filled[:, j], j] = np.interp(grid[filled[:, j]],
                                                     grid[good], surface[good, j])
        for i in np.flatnonzero(filled.any(axis=1)):
            bins = np.flatnonzero(filled[i])
            print(f"[lut] {energies[i]:2d} keV: {len(bins)} depth bins filled from "
                  f"neighbouring energies (bins {bins.min()}-{bins.max()})")

    os.makedirs(os.path.dirname(conf.output), exist_ok=True)
    np.savez(conf.output, energies=np.array(energies, float), z0=z0,
             drift_time=drift_time, alpha=alpha, beta=beta,
             alpha_err=alpha_err, beta_err=beta_err, n_carriers=n_carriers,
             filled_from_energy=filled if conf.get("fill_missing", None) == "energy"
             else np.zeros(shape, bool),
             sensor=conf.sensor.code, hv=float(conf.sensor.hv),
             thickness_um=thickness)
    print(f"[lut] wrote {conf.output}  "
          f"alpha {np.nanmin(alpha):.2f}-{np.nanmax(alpha):.2f} um, "
          f"beta {np.nanmin(beta):.2f}-{np.nanmax(beta):.2f}")


if __name__ == "__main__":
    main()