#!/usr/bin/env bash
set -euo pipefail

ROOT=/home/yurenh2/sdil
AUTHOR=/scratch/yurenh2/energy-based-learning
PYTHON=/scratch/yurenh2/venvs/burstccn/bin/python
PROFILE="$ROOT/results/physical_bias/p0_state_dependence.json"
OUT="$ROOT/results/ep_bias/centered_r2"
mkdir -p "$OUT"
cd "$AUTHOR"

run_cell() {
  local gpu=$1
  local seed=$2
  local mode=$3
  local tag=$4
  CUBLAS_WORKSPACE_CONFIG=:4096:8 CUDA_VISIBLE_DEVICES="$gpu" "$PYTHON" \
    "$ROOT/experiments/rain_ep_bias_train.py" \
    --author-root "$AUTHOR" --device cuda \
    --adapter dillavou --network-protocol comparative32 \
    --beta-policy centered --beta-value 0.25 \
    --mode "$mode" --bias-ratio 1 --dillavou-drift-ratio 0 \
    --dillavou-profile-json "$PROFILE" --predictor-rate 1 \
    --dillavou-calibration-steps 1 --neutral-cadence 10 \
    --epochs 1 --schedule-epochs 100 \
    --train-limit 10000 --test-limit 2000 --batch-size 128 \
    --training-iterations 15 --inference-iterations 60 \
    --evaluation-split train_holdout --data-seed 6200 \
    --seed "$seed" --beta-seed 7100 --deterministic \
    --output "$OUT/$tag.json" > "$OUT/$tag.log" 2>&1
}

seeds=(1988 1989 1990 1991)
gpu=0
for seed in "${seeds[@]}"; do
  run_cell "$gpu" "$seed" constant "seed_${seed}_intercept" &
  gpu=$((gpu + 1))
  run_cell "$gpu" "$seed" innovation "seed_${seed}_sdil" &
  gpu=$((gpu + 1))
done
wait

"$PYTHON" - "$OUT" <<'PY'
import json
import math
from pathlib import Path
import statistics
import sys

root = Path(sys.argv[1])
by_seed = {}
for path in sorted(root.glob("seed_*.json")):
    _, seed, method = path.stem.split("_", 2)
    report = json.loads(path.read_text())
    by_seed.setdefault(seed, {})[method] = report["final"]["test_accuracy"]

for seed, values in by_seed.items():
    if set(values) != {"intercept", "sdil"}:
        raise RuntimeError(f"incomplete pair for seed {seed}: {values}")
    values["sdil_minus_intercept"] = values["sdil"] - values["intercept"]

summary = {"by_seed": by_seed}
for key in ("intercept", "sdil", "sdil_minus_intercept"):
    values = [pair[key] for pair in by_seed.values()]
    summary[key] = {
        "mean": statistics.mean(values),
        "sample_std": statistics.stdev(values),
        "values": values,
    }
(root / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
print(json.dumps(summary, indent=2))
PY