#!/usr/bin/env bash
# Frozen validation-only recovery screen after the failed MNIST C1 confirmation.
# Usage: c1_fmnist_validation.sh <gpu> <traffic projection seed> [include no-traffic panel]
set -eu

cd "$(dirname "$0")/.."
GPU="${1:?GPU index required}"
TRAFFIC_SEED="${2:?traffic projection seed required}"
INCLUDE_BASELINE="${3:-0}"
PYTHON="${PYTHON:-/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3}"

run_one() {
  family="$1"
  scale="$2"
  signal="$3"

  use_residual=1
  raw_control=none
  p_neutral=1
  case "$signal" in
    raw)
      use_residual=0
      ;;
    matched)
      use_residual=0
      raw_control=match_innovation_norm
      ;;
    residual)
      ;;
    taskfit)
      p_neutral=0
      ;;
    *)
      echo "unknown signal control: $signal" >&2
      exit 2
      ;;
  esac

  scale_tag="${scale//./p}"
  tag="c1_fmnist_val_v1_${family}_rho${scale_tag}_t${TRAFFIC_SEED}_${signal}_d3_s0"
  out="results/${tag}.json"
  if [ -s "$out" ]; then
    echo "[$tag] exists; skipping"
    return
  fi

  CUDA_VISIBLE_DEVICES="$GPU" "$PYTHON" experiments/run.py \
    --mode sdil --dataset fmnist --device cuda \
    --depth 3 --width 256 --residual 1 --act tanh \
    --epochs 15 --batch_size 128 --eta 0.05 --momentum 0.9 \
    --eta_A 0.02 --eta_P 0.05 \
    --pert_sigma 0.01 --pert_every 4 --pert_ndirs 1 \
    --pert_mode simultaneous \
    --traffic_mode "$family" --traffic_seed "$TRAFFIC_SEED" --nuis_rho "$scale" \
    --use_residual "$use_residual" --raw_scale_control "$raw_control" \
    --learn_A 1 --learn_P 1 --p_neutral "$p_neutral" \
    --p_warmup_steps 200 --p_warmup_eta 0.05 \
    --val_examples 5000 --split_seed 2027 --eval_split validation --eval_every 0 \
    --diagnostics alignment --diagnostics_schedule final --probe_bs 512 \
    --seed 0 --log_every 100000 --tag "$tag"
}

# The no-traffic panel is generated once, on projection seed 1234.
if [ "$INCLUDE_BASELINE" = 1 ]; then
  if [ "$TRAFFIC_SEED" != 1234 ]; then
    echo "the no-traffic panel is frozen to projection seed 1234" >&2
    exit 2
  fi
  for signal in raw matched residual; do
    run_one none 0 "$signal"
  done
fi

# Soma traffic remains at the level used in the MNIST confirmation. Top-down
# traffic gets a two-level validation screen; no test result selects a level.
for signal in raw matched residual taskfit; do
  run_one soma 0.5 "$signal"
done
for scale in 0.2 0.5; do
  for signal in raw matched residual taskfit; do
    run_one topdown "$scale" "$signal"
  done
done