#!/bin/bash
# Grid 4 - the method: dependence-weighted entropy regularization.
#
# Prediction, and the thing that makes it falsifiable:
#   agree/parity  (positions dependent)   -> accuracy at high parallelism RISES
#   copy          (positions independent) -> UNCHANGED at 100%, because measured
#                                            dependence is ~0 so no penalty applies
# The un-weighted ablation (dep_weighted=0) should hurt `copy`, which is how we
# show the weighting is doing the work rather than the entropy term alone.
#SBATCH -J g4
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=64g
#SBATCH -t 06:00:00
#SBATCH -o logs/g4-%A_%a.out
#SBATCH --array=0-35%12
source /etc/profile
SCR=/scratch/zt1/project/msml612/user/$USER
module load python/3.10.10
export PYTHONPATH=
PY=$SCR/hf_env/bin/python
cd $SCR/pcmd/src

I=$SLURM_ARRAY_TASK_ID
TASKS=(copy agree parity); GAMMA=(0.5 2.0 8.0); DW=(1 0); SEEDS=(0 1)
T=${TASKS[$((I / 12))]}
G=${GAMMA[$(((I % 12) / 4))]}
D=${DW[$(((I % 4) / 2))]}
S=${SEEDS[$((I % 2))]}

NAME="g4-$T-gam$G-dw$D-s$S"
echo "=== $NAME ==="
$PY train.py --task $T --objective coord --coupling 4 --seed $S \
  --seq_len 128 --n_train 200000 --n_eval 500 --eval_bs 100 \
  --d 384 --layers 6 --heads 6 --bs 64 --steps 20000 --lr 3e-4 \
  --beta 0.0 --gamma $G --dep_weighted $D --elbo_w 1.0 --ramp_frac 0.2 \
  --out $SCR/pcmd/runs/$NAME