"""
Generate a small single-photon MC cluster sample at 12 keV (Moench040, 150V).

Mirrors McGeneration/generate_Moench040_150V.py exactly: `parameterization()`
fits the generalized-Gaussian charge-cloud parameters alpha(t)/beta(t) from
raw per-depth-bin charge-transport simulation output (see
https://github.com/slsdetectorgroup/ChargeTransportSimulation for how that
raw output itself is produced -- data/mc_raw_simulation/ ships the 256
files needed for this one sensor/voltage/energy point, ~34 MB), then
`generateFromParameters()` draws the MC events. Only the event count is
much smaller here than production.

Output:
    data/mc_params/ggdPar_Moench040_150V_12keV.npy (+ _fit.png) -- fitted
        alpha(t)/beta(t) parameters, regenerated fresh each run
    data/mc_samples/12keV_Moench040_150V_{thread}.npz (+ .h5)
        samples: (N, 5, 5) float32 pixel clusters in keV
        labels:  (N, 4) float32 [x, y, z_um, energy_eV] -- MC ground truth

The beta(t) fit is a 5-parameter nonlinear fit and occasionally converges
such that it's undefined (NaN) for a handful of the shortest drift times;
singleProcess() already discards those events (~0.1-0.2% of the total,
printed per event) the same way it discards out-of-range z0 draws -- not a
bug, just how that fit behaves at its domain edge.

Needs: numpy, scipy, ROOT, h5py. Use the `mlxid_demo` conda env.
"""
import sys
from pathlib import Path
import numpy as np

sys.path.append(str(Path(__file__).resolve().parent.parent / 'src'))
import mc_generator

HERE = Path(__file__).resolve().parent.parent
RAW_SIM_DIR = HERE / 'data/mc_raw_simulation'
GGD_PARAM_DIR = HERE / 'data/mc_params'
OUTPUT_DIR = HERE / 'data/mc_samples'

### 12 keV attenuation length in Moench040 (Si), from https://henke.lbl.gov/optical_constants/atten2.html
ATTENUATION_LENGTH_12KEV_CM = 228.738

N_TOTAL_INCIDENT = 40_000  ### demo scale; production uses 10_000_000
N_THREAD = 4

mcConfig = {
    'element': '12keV',
    'energy': 12_000,  ### eV
    'attenuationLength': ATTENUATION_LENGTH_12KEV_CM,
    'sensorCode': 'Moench040',
    'nTotalIncident': N_TOTAL_INCIDENT,
    'nThread': N_THREAD,
    'Roi': [50, 350, 50, 350],
    'clusterWidth': 5,  ### pixel
    'pixelSize': 25,  ### um
    'noiseEneFrame': np.zeros((400, 400)),  ### no noise map shipped -> noise-free MC
    'shotNoiseFactor': 0.0,
    'calibrationNoise': 0.0,
    'sensorThickness': 650,  ### um, Moench040
    'T': 273 + 20.0 + 14.8,
    'hv': 150,
    'depletionVoltage': 34.7,
    'zBins': 128,  ### depth bins the raw simulation was run in
    'simulationMethod': 'simu2',
    'resultsPath': str(RAW_SIM_DIR),
    'ggdParOutputDir': str(GGD_PARAM_DIR),
    'sampleOutputPath': str(OUTPUT_DIR),
}

if __name__ == '__main__':
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    GGD_PARAM_DIR.mkdir(parents=True, exist_ok=True)
    mc_generator.init(mcConfig)
    print(f'[01] Fitting alpha(t)/beta(t) from {RAW_SIM_DIR} -> {GGD_PARAM_DIR}...')
    mc_generator.parameterization()
    print(f'[01] Generating {N_TOTAL_INCIDENT} single-photon 12keV MC events '
          f'across {N_THREAD} threads -> {OUTPUT_DIR}')
    mc_generator.generateFromParameters()
    print('[01] Done.')