#!/usr/bin/env bash
# SDIL experiment battery. Runs sequentially on ONE GPU (nets are tiny).
# Usage: bash run_battery.sh <wave>   wave in {A,B,all}
set -u
cd "$(dirname "$0")/.."               # -> /home/yurenh2/sdil
PY=/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3
export OMP_NUM_THREADS=2
WAVE="${1:-all}"
OUT=results
LOGS=logs/battery
mkdir -p "$OUT" "$LOGS"
SEEDS="0 1 2"
EP=20
PROG="$LOGS/progress.txt"
echo "=== battery start wave=$WAVE $(date) ===" >> "$PROG"

run () {   # run <tag> <args...>
  local tag="$1"; shift
  if [ -f "$OUT/$tag.json" ]; then echo "skip $tag (exists)" >> "$PROG"; return; fi
  echo ">>> $tag  $(date +%H:%M:%S)" >> "$PROG"
  $PY experiments/run.py --tag "$tag" --outdir "$OUT" "$@" > "$LOGS/$tag.log" 2>&1 \
     && grep -h "DONE" "$LOGS/$tag.log" >> "$PROG" \
     || echo "FAIL $tag" >> "$PROG"
}

# ---------------- Wave A: core claims, seed 0 (fast health + headline) --------
if [ "$WAVE" = "A" ] || [ "$WAVE" = "all" ]; then
  for m in bp dfa sdil; do
    run "A_main_${m}_s0" --mode $m --dataset mnist --depth 3 --width 256 --epochs $EP --seed 0
  done
  # residualization under nuisance (the Harnett-specific claim), seed 0
  for rho in 0 2 8; do
    for ur in 0 1; do
      run "A_nuis_rho${rho}_res${ur}_s0" --mode sdil --nuis_rho $rho --use_residual $ur \
          --depth 3 --width 256 --epochs $EP --seed 0
    done
  done
  # trained vs fixed apical pathway, seed 0
  run "A_fixedA_s0" --mode sdil --learn_A 0 --depth 3 --width 256 --epochs $EP --seed 0
  for nd in 1 4 16; do
    run "A_trainedA_nd${nd}_s0" --mode sdil --learn_A 1 --pert_ndirs $nd --depth 3 --width 256 --epochs $EP --seed 0
  done
fi

# ---------------- Wave B: full grid, all seeds ------------------------------
if [ "$WAVE" = "B" ] || [ "$WAVE" = "all" ]; then
  for s in $SEEDS; do
    # E1 main comparison + fashion-mnist
    for m in bp dfa sdil; do
      run "main_${m}_mnist_s${s}"  --mode $m --dataset mnist  --depth 3 --width 256 --epochs $EP --seed $s
      run "main_${m}_fmnist_s${s}" --mode $m --dataset fmnist --depth 3 --width 256 --epochs $EP --seed $s
    done
    # E2 residualization x nuisance
    for rho in 0 1 2 4 8; do
      for ur in 0 1; do
        run "nuis_rho${rho}_res${ur}_s${s}" --mode sdil --nuis_rho $rho --use_residual $ur \
            --depth 3 --width 256 --epochs $EP --seed $s
      done
    done
    # E3 trained vs fixed A, perturbation directions
    run "abl_fixedA_s${s}" --mode sdil --learn_A 0 --depth 3 --width 256 --epochs $EP --seed $s
    for nd in 1 4 16; do
      run "abl_trainedA_nd${nd}_s${s}" --mode sdil --learn_A 1 --pert_ndirs $nd --depth 3 --width 256 --epochs $EP --seed $s
    done
    # E4 depth scaling
    for d in 1 2 3 5 7; do
      for m in bp dfa sdil; do
        run "depth_${m}_d${d}_s${s}" --mode $m --depth $d --width 256 --epochs $EP --seed $s
      done
    done
    # E5 feedback type + online apical control
    for fb in error error_deriv; do
      for ss in 0 5; do
        run "ctrl_fb${fb}_settle${ss}_s${s}" --mode sdil --feedback $fb --settle_steps $ss \
            --kappa 0.3 --depth 3 --width 256 --epochs $EP --seed $s
      done
    done
    # E6 predictor timescale (rho=2 so P matters); neutral vs task-period update
    for neu in 0 1; do
      for ep in 0.0005 0.002 0.01; do
        run "pred_neu${neu}_etaP${ep}_s${s}" --mode sdil --nuis_rho 2 --learn_P 1 \
            --p_neutral $neu --eta_P $ep --depth 3 --width 256 --epochs $EP --seed $s
      done
    done
  done
fi

echo "=== battery done wave=$WAVE $(date) ===" >> "$PROG"