#!/usr/bin/env bash
# Overnight run (2026-09-18): poly models trained on 3-30 keV, four network sizes.
#
#   setsid nohup bash scripts/overnight_poly3to30keV.sh \
#       > Results/overnight_poly3to30keV/driver.log 2>&1 < /dev/null &
#   DRY_RUN=1 bash scripts/overnight_poly3to30keV.sh     # print the plan only
#
# Every stage logs to Results/overnight_poly3to30keV/logs/<stage>.log and leaves
# state/<stage>.done when it succeeds; a rerun skips finished stages, so the
# script can simply be started again after a failure.  Progress, with times, is
# appended to Results/overnight_poly3to30keV/status.md.  A failure of the frame
# generation stops everything; a failure inside one model skips the rest of that
# model and goes on with the next.  The last stage writes report/report.md.

set -u
cd /home/xie_x1/MLXID/GridBased
PY=/home/xie_x1/software/miniforge3/envs/pytorch_nightly/bin/python
BASE=Results/overnight_poly3to30keV
LOG=$BASE/logs
STATE=$BASE/state
STATUS=$BASE/status.md
INFER=Configs/infer_measurement_poly_Moench014.yaml
GEN=Configs/generate_frames_poly3to30keV_Moench014.yaml
DATA=/mnt/sls_det_storage/moench_data/MLXID/Samples/GridBased/poly3to30keV_Moench014_150V
REFERENCE=Results/260901_poly5to30keV_Moench014_00      # the current 5-30 keV model
MODELS="c128 c96 c64 r3"
declare -A CHANNELS=([c128]="[32,64,128,128]" [c96]="[32,64,96,96]"
                     [c64]="[32,64,64,64]" [r3]="[32,64,128]")
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() < 2.999 or energy.max() > 30.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 3-30 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']) == 300 else 1)"
}

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

# --- 1. charge-cloud table (done interactively before launch) ---------------
run lut $PY scripts/01_build_charge_lut.py --config Configs/charge_lut_Moench014_3to30keV.yaml --workers 14
run lut_validation $PY scripts/01b_validate_charge_lut.py \
    --lut data/charge_lut_Moench014_150V_3to30keV.npz --outdir Results/charge_lut_3to30keV

# --- 2. MC frames, 3-30 keV (~1 h) ------------------------------------------
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. four models, one after the other (~3 h each) ------------------------
# its XRF runs were moved to <run>_poly5to30 when the c96 model became the default
compare_args=(--model "5-30 keV [32,64,128,128] (previous default)" "$REFERENCE" poly5to30)
for model in $MODELS; do
    if ! run train_$model $PY -u scripts/03_train.py \
            --config Configs/train_grid_poly3to30keV_Moench014_$model.yaml; then
        note "$model: training failed, skipping its evaluation"; continue
    fi
    if [ "$DRY_RUN" = 1 ]; then
        experiment="Results/<date>_poly3to30keV_Moench014_${model}_00"
    else
        [ -f "$STATE/$model.experiment" ] || \
            grep -m1 "\[train\] output ->" "$LOG/train_$model.log" | awk '{print $NF}' \
                > "$STATE/$model.experiment"
        experiment=$(cat "$STATE/$model.experiment")
        if ! run check_train_$model check_training "$experiment"; then
            note "$model: training incomplete in $experiment, skipping its evaluation"
            continue
        fi
    fi
    checkpoint=$experiment/Models/best.pth
    run evaluate_$model $PY scripts/04_evaluate.py --checkpoint "$checkpoint"
    run response_$model $PY scripts/07_energy_response.py --checkpoint "$checkpoint"
    for xrf in AgMo_1_factor_10 CuCrCo_1_factor_10 Pedestal_CuCrCo_1; do
        run xrf_${xrf}_$model $PY scripts/05_infer_measurement.py --config $INFER \
            --run $xrf --checkpoint "$checkpoint" --tag $model
    done
    compare_args+=(--model "3-30 keV ${CHANNELS[$model]}" "$experiment" $model)
done

# --- 4. comparison report ----------------------------------------------------
rm -f "$STATE/compare.done"        # always refresh the report
run compare $PY scripts/08_compare_models.py --out $BASE/report "${compare_args[@]}"
note "=== driver finished; report: $BASE/report/report.md ==="