#!/bin/bash
# Round 5: h=128 (keep capacity) + moderate regularization + multiple seeds
# Best of R3 capacity + some anti-overfit from R4
# Also: 3 seeds for the best config to get confidence intervals

PYTHON=python
BASEDIR=${PULSE_ROOT}
TRAIN_SCRIPT=${BASEDIR}/experiments/tasks/train_pred_cls.py
OUTDIR=${BASEDIR}/results/pred_cls5
LOGDIR=${OUTDIR}/slurm_logs
mkdir -p $LOGDIR

# h=128, lr=5e-4, wd=3e-4, dropout=0.3, moderate augment
COMMON="--coarse --use_prev_action --epochs 80 --batch_size 32 --lr 5e-4 --weight_decay 3e-4 --hidden_dim 128 --dropout 0.3 --downsample 5 --patience 20 --augment --noise_std 0.15 --time_mask_ratio 0.12 --label_smoothing 0.1 --output_dir $OUTDIR --window_sec 15.0"

# Top 6 modality combos
MODS=("imu" "emg" "mocap" "emg,imu" "mocap,imu" "mocap,emg,imu")

for mods in "${MODS[@]}"; do
    mod_tag=$(echo $mods | tr ',' '-')
    sbatch \
        -J "pcls5_${mod_tag}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=4 \
        --mem=32G \
        -t 2:00:00 \
        -o "${LOGDIR}/${mod_tag}_s42_%j.out" \
        -e "${LOGDIR}/${mod_tag}_s42_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $TRAIN_SCRIPT --modalities $mods --seed 42 $COMMON"
    echo "Submitted: $mods seed=42"
done

# 2 extra seeds for emg,imu (best combo) for confidence intervals
for seed in 123 456; do
    sbatch \
        -J "pcls5_emg-imu_s${seed}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=4 \
        --mem=32G \
        -t 2:00:00 \
        -o "${LOGDIR}/emg-imu_s${seed}_%j.out" \
        -e "${LOGDIR}/emg-imu_s${seed}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $TRAIN_SCRIPT --modalities emg,imu --seed $seed $COMMON"
    echo "Submitted: emg,imu seed=$seed"
done

echo ""
echo "Total: 8 jobs | h=128, lr=5e-4, dropout=0.3, wd=3e-4"
echo "Results: $OUTDIR"