#!/bin/bash
# Round 4: Anti-overfit — smaller model + higher dropout + lower lr + stronger augment
# Focus on top 6 modalities (skip eyetrack-only combos which are toxic)
# Also add a prev_action-only baseline (for ablation)
# Total: 7 jobs

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

# Smaller model, stronger regularization
COMMON="--coarse --use_prev_action --epochs 100 --batch_size 32 --lr 3e-4 --weight_decay 5e-4 --hidden_dim 64 --downsample 5 --patience 25 --seed 42 --augment --noise_std 0.2 --time_mask_ratio 0.15 --label_smoothing 0.15 --output_dir $OUTDIR --window_sec 15.0"

# Top modalities only (no eyetrack-only combos)
MODS=("imu" "emg" "mocap" "emg,imu" "mocap,imu" "mocap,emg,imu")

for mods in "${MODS[@]}"; do
    mod_tag=$(echo $mods | tr ',' '-')
    sbatch \
        -J "pcls4_${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}_%j.out" \
        -e "${LOGDIR}/${mod_tag}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $TRAIN_SCRIPT --modalities $mods $COMMON"
    echo "Submitted: $mods"
done

# Ablation: sensor-only (no prev_action) for best combo emg,imu
COMMON_NOPREV="--coarse --epochs 100 --batch_size 32 --lr 3e-4 --weight_decay 5e-4 --hidden_dim 64 --downsample 5 --patience 25 --seed 42 --augment --noise_std 0.2 --time_mask_ratio 0.15 --label_smoothing 0.15 --output_dir $OUTDIR --window_sec 15.0"
sbatch \
    -J "pcls4_emg-imu_noprev" \
    -p gpuA800 \
    --gres=gpu:1 \
    -N 1 -n 1 \
    --cpus-per-task=4 \
    --mem=32G \
    -t 2:00:00 \
    -o "${LOGDIR}/emg-imu_noprev_%j.out" \
    -e "${LOGDIR}/emg-imu_noprev_%j.err" \
    --export=ALL \
    --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $TRAIN_SCRIPT --modalities emg,imu $COMMON_NOPREV"
echo "Submitted: emg,imu (no prev_action ablation)"

echo ""
echo "Total: 7 jobs | anti-overfit: hidden=64, lr=3e-4, wd=5e-4, dropout, noise=0.2"
echo "Results: $OUTDIR"