#!/usr/bin/env bash
# Overnight run (2026-09-27/28): poly 2-20 keV model for MOENCH040 @150 V, then
# the 12 keV MaxIV knife edge at 0-3 filters and the 15 keV SOLEIL Siemens star.
#
#   setsid nohup bash scripts/overnight_poly2to20keV_Moench040.sh \
#       > Results/overnight_poly2to20keV_Moench040/driver.log 2>&1 < /dev/null &
#   DRY_RUN=1 bash scripts/overnight_poly2to20keV_Moench040.sh     # print the plan only
#
# Part A (this file): charge LUT -> MC frames -> training -> MC evaluation.
# Part B (overnight_poly2to20keV_Moench040_partB.sh): measured runs + notebooks.
# Part B is started at the end of part A once it has been marked ready
# (state/partB.ready), so it can be finished and tested while A trains.
#
# Every stage logs to $BASE/logs/<stage>.log and leaves state/<stage>.done when
# it succeeds; a rerun skips finished stages.  Progress with times goes to
# $BASE/status.md.

set -u
cd /home/xie_x1/MLXID/GridBased
PY=/home/xie_x1/software/miniforge3/envs/pytorch_nightly/bin/python
BASE=Results/overnight_poly2to20keV_Moench040
LOG=$BASE/logs
STATE=$BASE/state
STATUS=$BASE/status.md
LUT_CONF=Configs/charge_lut_Moench040_2to20keV.yaml
LUT=data/charge_lut_Moench040_150V_2to20keV.npz
GEN=Configs/generate_frames_poly2to20keV_Moench040.yaml
TRAIN=Configs/train_grid_poly2to20keV_Moench040_c96.yaml
DATA=/mnt/sls_det_storage/moench_data/MLXID/Samples/GridBased/poly2to20keV_Moench040_150V
EPOCHS=200
DRY_RUN=${DRY_RUN:-0}
mkdir -p "$LOG" "$STATE"

note() {
    [ "$DRY_RUN" = 1 ] || echo "- $(date '+%m-%d %H:%M:%S') $*" >> "$STATUS"
    echo "$(date '+%m-%d %H:%M:%S') $*"
}

# run NAME COMMAND...   -- skip if done, log, time, mark
run() {
    local name=$1; shift
    if [ -f "$STATE/$name.done" ]; then
        note "$name: already done, skipped"; return 0
    fi
    if [ "$DRY_RUN" = 1 ]; then
        echo "[dry] $name: $*"; return 0
    fi
    note "$name: started"
    local started=$(date +%s)
    "$@" > "$LOG/$name.log" 2>&1
    local status=$?
    local minutes=$(( ($(date +%s) - started) / 60 ))
    if [ $status -eq 0 ]; then
        touch "$STATE/$name.done"
        note "$name: done (${minutes} min)"
    else
        note "$name: **FAILED** (exit $status after ${minutes} min), see logs/$name.log"
    fi
    return $status
}

check_frames() {
    $PY - "$DATA" <<'EOF'
import glob, sys
import h5py
folder = sys.argv[1]
expected = {"train": (12, 100_000), "val": (2, 40_000), "test": (2, 40_000),
            "eval_scan": (5, 8_000)}
bad = []
for split, (n_files, n_frames) in expected.items():
    files = sorted(glob.glob(f"{folder}/{split}/*.h5"))
    if len(files) != n_files:
        bad.append(f"{split}: {len(files)} files, expected {n_files}")
    for path in files:
        with h5py.File(path) as f:
            frames, energy = f["frames"].shape[0], f["photons"][:, 2]
        if frames != n_frames:
            bad.append(f"{path}: {frames} frames, expected {n_frames}")
        if energy.min() < 1.999 or energy.max() > 20.001:
            bad.append(f"{path}: energy {energy.min():.3f}-{energy.max():.3f} keV")
print("\n".join(bad) if bad else "all shards complete, energies within 2-20 keV")
sys.exit(1 if bad else 0)
EOF
}

check_training() {   # experiment dir
    [ -f "$1/Models/best.pth" ] && $PY -c "
import json, sys
h = json.load(open('$1/history.json'))
print(len(h['val_loss']), 'epochs, best val loss', min(h['val_loss']))
sys.exit(0 if len(h['val_loss']) == $EPOCHS else 1)"
}

note "=== driver started (pid $$) ==="

# --- 1. charge-cloud table ---------------------------------------------------
if ! run lut $PY scripts/01_build_charge_lut.py --config $LUT_CONF --workers 19; then
    note "charge LUT failed: stopping"; exit 1
fi
# diagnostic only: its comparison with the 9-parameter fits needs MOENCH040
# ggdPar files that exist for a few energies only, so a failure does not stop
# the chain
run lut_validation $PY scripts/01b_validate_charge_lut.py --lut $LUT \
    --sensor Moench040 --outdir Results/charge_lut_Moench040_2to20keV

# --- 2. MC frames, 2-20 keV (~35 min) ----------------------------------------
if ! run generate $PY scripts/02_generate_frames.py --config $GEN; then
    note "frame generation failed: stopping"; exit 1
fi
if ! run check_frames check_frames; then
    note "frame check failed: stopping"; exit 1
fi

# --- 3. training, [32, 64, 96, 96], 200 epochs (~1 h 40 min) -----------------
if ! run train $PY -u scripts/03_train.py --config $TRAIN; then
    note "training failed: stopping"; exit 1
fi
if [ "$DRY_RUN" = 1 ]; then
    experiment="Results/<date>_poly2to20keV_Moench040_c96_00"
else
    [ -f "$STATE/experiment" ] || \
        grep -m1 "\[train\] output ->" "$LOG/train.log" | awk '{print $NF}' > "$STATE/experiment"
    experiment=$(cat "$STATE/experiment")
    if ! run check_train check_training "$experiment"; then
        note "training incomplete in $experiment: stopping"; exit 1
    fi
fi
checkpoint=$experiment/Models/best.pth

# --- 4. MC evaluation --------------------------------------------------------
run evaluate $PY scripts/04_evaluate.py --checkpoint "$checkpoint"
run response $PY scripts/07_energy_response.py --checkpoint "$checkpoint"
note "=== part A finished: $experiment ==="

# --- 5. part B: measured runs and notebooks ----------------------------------
PARTB=scripts/overnight_poly2to20keV_Moench040_partB.sh
if [ "$DRY_RUN" = 1 ]; then
    echo "[dry] part B: bash $PARTB $experiment"; exit 0
fi
waited=0
while [ ! -f "$STATE/partB.ready" ] && [ $waited -lt 7200 ]; do
    [ $waited -eq 0 ] && note "waiting for state/partB.ready"
    sleep 60; waited=$((waited + 60))
done
if [ ! -f "$STATE/partB.ready" ]; then
    note "part B was never marked ready: stopping after part A"; exit 1
fi
bash "$PARTB" "$experiment"