#!/usr/bin/env python3
"""
WP1.1 validation -- is the alpha(E,t) / beta(E,t) table good enough to sample
continuous energies from?

Four questions:
  1. how noisy are the tabulated surfaces, and are they monotone in depth?
  2. leave-one-out: drop an energy, interpolate it back from its neighbours,
     and see how far off the interpolation is.  With +-1 keV neighbours this
     tests a 2 keV gap, twice the spacing the table actually uses.
  3. does the table reproduce the 9-parameter fits that McGenerator uses at the
     energies where both exist?
  4. what happens if the 9 fit parameters are interpolated in energy instead
     of the curves -- the thing this table exists to avoid.
"""
import argparse
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "src"))
from charge_model import ChargeCloudLUT   # noqa: E402

GGD_PAR_DIR = Path('/home/xie_x1/MLXID/McGeneration/Notebooks')


def parametric_alpha(pars, t):
    return pars[0] + pars[1] * np.sqrt(t) + pars[2] * t + pars[3] * t ** 2


def parametric_beta(pars, t):
    return pars[4] * (t - pars[5]) ** pars[6] + pars[7] * np.exp(pars[8] * t) + 2


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--lut", default="data/charge_lut_Moench014_150V.npz")
    parser.add_argument("--sensor", default="Moench014")
    parser.add_argument("--hv", type=int, default=150)
    parser.add_argument("--outdir", default="Results/charge_lut")
    args = parser.parse_args()

    lut = ChargeCloudLUT(args.lut)
    energies, time = lut.energies, lut._time
    alpha, beta = lut._alpha, lut._beta
    outdir = Path(args.outdir); outdir.mkdir(parents=True, exist_ok=True)

    # ---- 1. noise and monotonicity ------------------------------------
    print("1. surfaces")
    print(f"   grid {alpha.shape[0]} energies x {alpha.shape[1]} depth bins, "
          f"drift time {time.min():.2f}-{time.max():.2f} ns")
    # bin-to-bin scatter after removing the smooth trend
    from scipy.signal import savgol_filter
    jitter_a, jitter_b = [], []
    for row in range(alpha.shape[0]):
        jitter_a.append(np.std(alpha[row] - savgol_filter(alpha[row], 21, 3)) /
                        np.mean(alpha[row]))
        jitter_b.append(np.std(beta[row] - savgol_filter(beta[row], 21, 3)) /
                        np.mean(beta[row]))
    print(f"   bin-to-bin scatter about a local fit: alpha {np.mean(jitter_a)*100:.2f} %, "
          f"beta {np.mean(jitter_b)*100:.2f} %")
    rising = np.mean([np.all(np.diff(alpha[r]) > -0.02) for r in range(alpha.shape[0])])
    print(f"   alpha monotone in drift time for {rising*100:.0f} % of energies")
    at_limit = ((alpha <= 0.501).mean(), (beta <= 2.001).mean())
    print(f"   fits sitting on the parameter limits: alpha {at_limit[0]*100:.1f} %, "
          f"beta {at_limit[1]*100:.1f} %  (shortest drifts, cloud smaller than a micron)")

    # ---- 2. leave-one-out in energy -----------------------------------
    print("\n2. leave-one-out in energy (interpolating across a 2 keV gap)")
    rows = []
    for i in range(1, len(energies) - 1):
        a_hat = 0.5 * (alpha[i - 1] + alpha[i + 1])
        b_hat = 0.5 * (beta[i - 1] + beta[i + 1])
        rows.append((energies[i],
                     np.max(np.abs(a_hat - alpha[i]) / alpha[i]) * 100,
                     np.mean(np.abs(a_hat - alpha[i]) / alpha[i]) * 100,
                     np.max(np.abs(b_hat - beta[i]) / beta[i]) * 100))
    rows = np.array(rows)
    print(f"   alpha: mean |error| {rows[:,2].mean():.3f} %, worst energy "
          f"{rows[rows[:,1].argmax(),0]:.0f} keV at {rows[:,1].max():.2f} % (max over depth)")
    print(f"   beta : worst {rows[:,3].max():.2f} %")
    print("   per energy (keV, mean%, max%):",
          ", ".join(f"{r[0]:.0f}:{r[2]:.2f}/{r[1]:.2f}" for r in rows[::5]))

    # ---- 3. against the 9-parameter fits ------------------------------
    print("\n3. table vs the 9-parameter fits McGenerator uses")
    valid = (time > 0.3) & (time < time.max() - 0.2)
    comparison = []
    for i, energy in enumerate(energies):
        path = GGD_PAR_DIR / f'ggdPar_{args.sensor}_{args.hv}V_{int(energy)}keV.npy'
        if not path.exists():
            continue
        pars = np.load(path)
        da = (parametric_alpha(pars, time[valid]) - alpha[i][valid]) / alpha[i][valid]
        db = (parametric_beta(pars, time[valid]) - beta[i][valid]) / beta[i][valid]
        comparison.append((energy, np.mean(np.abs(da)) * 100, np.mean(np.abs(db)) * 100))
    comparison = np.array(comparison)
    print(f"   {len(comparison)} energies compared over {valid.sum()} depth bins "
          f"inside the parametric fit's validity range")
    print(f"   mean |alpha difference| {comparison[:,1].mean():.2f} % "
          f"(worst {comparison[:,1].max():.2f} % at {comparison[comparison[:,1].argmax(),0]:.0f} keV)")
    print(f"   mean |beta  difference| {comparison[:,2].mean():.2f} %")

    # ---- 4. interpolating the 9 parameters instead of the curves -------
    print("\n4. interpolating the 9 fit parameters instead (what this table avoids)")
    worst = 0.0
    for target in range(6, 30):
        low = GGD_PAR_DIR / f'ggdPar_{args.sensor}_{args.hv}V_{target-1}keV.npy'
        high = GGD_PAR_DIR / f'ggdPar_{args.sensor}_{args.hv}V_{target+1}keV.npy'
        if not (low.exists() and high.exists()):
            continue
        blended = 0.5 * (np.load(low) + np.load(high))
        i = int(np.argmin(np.abs(energies - target)))
        da = np.abs(parametric_alpha(blended, time[valid]) - alpha[i][valid]) / alpha[i][valid]
        worst = max(worst, da.max() * 100)
        if target in (8, 12, 20, 28):
            print(f"   {target:2d} keV: parameter blend is off by up to {da.max()*100:6.2f} %"
                  f" (curve interpolation: {rows[rows[:,0]==target][0,1]:.2f} %)")
    print(f"   worst case over all energies: {worst:.1f} %")

    # ---- figure --------------------------------------------------------
    fig, ax = plt.subplots(2, 3, figsize=(16, 8))
    for i in range(0, len(energies), 5):
        ax[0, 0].plot(time, alpha[i], lw=1, label=f'{energies[i]:.0f} keV')
        ax[0, 1].plot(time, beta[i], lw=1)
        ax[0, 2].plot(time, lut.cloud_sigma(energies[i], time), lw=1)
    ax[0, 0].set_ylabel(r'$\alpha$ [$\mu$m]'); ax[0, 0].legend(fontsize=8)
    ax[0, 1].set_ylabel(r'$\beta$')
    ax[0, 2].set_ylabel(r'cloud RMS [$\mu$m]')
    for a in ax[0]:
        a.set_xlabel('drift time [ns]'); a.grid(alpha=.3)

    im = ax[1, 0].imshow(alpha, aspect='auto', origin='lower', cmap='viridis',
                         extent=[time.min(), time.max(), energies[0], energies[-1]])
    ax[1, 0].set_title(r'$\alpha(E, t)$ [$\mu$m]'); plt.colorbar(im, ax=ax[1, 0])
    im = ax[1, 1].imshow(beta, aspect='auto', origin='lower', cmap='viridis',
                         extent=[time.min(), time.max(), energies[0], energies[-1]])
    ax[1, 1].set_title(r'$\beta(E, t)$'); plt.colorbar(im, ax=ax[1, 1])
    for a in ax[1, :2]:
        a.set_xlabel('drift time [ns]'); a.set_ylabel('energy [keV]')

    ax[1, 2].plot(rows[:, 0], rows[:, 1], 'o-', label=r'curve interp, max over depth')
    ax[1, 2].plot(rows[:, 0], rows[:, 2], 's-', label=r'curve interp, mean')
    ax[1, 2].set_xlabel('energy left out [keV]'); ax[1, 2].set_ylabel(r'$\alpha$ error [%]')
    ax[1, 2].set_yscale('log'); ax[1, 2].legend(fontsize=8); ax[1, 2].grid(alpha=.3)
    ax[1, 2].set_title('leave-one-out across a 2 keV gap')
    fig.tight_layout()
    fig.savefig(outdir / 'charge_lut_validation.png', dpi=130)
    print(f"\nwrote {outdir / 'charge_lut_validation.png'}")


if __name__ == "__main__":
    main()