#!/usr/bin/env bash
# Harnett-specific residualization ablation with per-neuron affine coupling.
#
# Usage:
#   bash experiments/nuisance_sweep.sh <gpu> "<rhos>" "<seeds>" <epochs> <prefix>
#
# Example:
#   bash experiments/nuisance_sweep.sh 5 "0 0.5 1 2" "0 1 2" 20 diag_nuis
set -eu

cd "$(dirname "$0")/.."
PY=/home/yurenh2/miniconda3/envs/ep_pascal/bin/python3
GPU="${1:?GPU index required}"
RHOS="${2:-0 0.5 1 2}"
SEEDS="${3:-0}"
EPOCHS="${4:-20}"
PREFIX="${5:-diag_nuis}"

export CUDA_VISIBLE_DEVICES="$GPU"
export OMP_NUM_THREADS=2
mkdir -p results logs/nuisance

run_one() {
  local rho="$1"
  local residual="$2"
  local seed="$3"
  local rho_tag="${rho//./p}"
  local tag="${PREFIX}_rho${rho_tag}_res${residual}_s${seed}"
  local result="results/${tag}.json"
  local log="logs/nuisance/${tag}.log"
  if [[ -f "$result" ]]; then
    echo "skip $tag (result exists)"
    return
  fi
  echo ">>> $tag $(date --iso-8601=seconds) gpu=$GPU"
  "$PY" experiments/run.py \
    --mode sdil --dataset mnist --depth 3 --width 256 --epochs "$EPOCHS" \
    --seed "$seed" --nuis_rho "$rho" --use_residual "$residual" \
    --predictor_mode diagonal --p_warmup_steps 200 --p_warmup_eta 0.05 \
    --pert_ndirs 8 --log_every 100 --tag "$tag" --outdir results \
    > "$log" 2>&1
  grep -h DONE "$log"
}

for seed in $SEEDS; do
  for rho in $RHOS; do
    for residual in 0 1; do
      run_one "$rho" "$residual" "$seed"
    done
  done
done