#!/bin/bash
# Scene Recognition (Exp1 v3) - Train 14 vols / Test 4 vols (no val)
# v23,v24 moved from val to train; v3 stays in test
# Part 1: 9 modality combos × 3 backbones = 27 jobs (early fusion)
# Part 2: 7 fusion methods × transformer × (3-core + all-5) = 14 jobs
# Total: 41 jobs

PYTHON=python
BASEDIR=${PULSE_ROOT}
SCRIPT=${BASEDIR}/experiments/train_exp1.py
OUTDIR=${BASEDIR}/results/exp1_v3
LOGDIR=${OUTDIR}/slurm_logs
mkdir -p $LOGDIR

COMMON="--epochs 100 --batch_size 16 --lr 1e-3 --weight_decay 1e-4 --hidden_dim 128 --downsample 5 --patience 15 --seed 42 --output_dir $OUTDIR"

MODS=("mocap" "emg" "eyetrack" "imu" "pressure" "mocap,emg,eyetrack" "mocap,emg,eyetrack,imu" "mocap,emg,eyetrack,pressure" "mocap,emg,eyetrack,imu,pressure")
MODELS=("cnn" "lstm" "transformer")

# Part 1: Modality ablation × 3 backbones
echo "=== Part 1: Modality Ablation (27 jobs) ==="
for mods in "${MODS[@]}"; do
    mod_tag=$(echo $mods | tr ',' '-')
    for model in "${MODELS[@]}"; do
        sbatch \
            -J "e1v3_${model}_${mod_tag}" \
            -p gpuA800 \
            --gres=gpu:1 \
            -N 1 -n 1 \
            --cpus-per-task=4 \
            --mem=32G \
            -t 2:00:00 \
            -o "${LOGDIR}/${model}_${mod_tag}_early_%j.out" \
            -e "${LOGDIR}/${model}_${mod_tag}_early_%j.err" \
            --export=ALL \
            --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --model $model --modalities $mods --fusion early $COMMON"
        echo "  $model / $mods / early"
    done
done

# Part 2: Fusion methods × transformer
FUSIONS=("late" "attention" "weighted_late" "gated_late" "stacking" "product" "moe")
FUSION_MODS=("mocap,emg,eyetrack" "mocap,emg,eyetrack,imu,pressure")

echo ""
echo "=== Part 2: Fusion Ablation (14 jobs) ==="
for fmods in "${FUSION_MODS[@]}"; do
    fmod_tag=$(echo $fmods | tr ',' '-')
    for fusion in "${FUSIONS[@]}"; do
        sbatch \
            -J "e1v3_tf_${fusion}" \
            -p gpuA800 \
            --gres=gpu:1 \
            -N 1 -n 1 \
            --cpus-per-task=4 \
            --mem=32G \
            -t 2:00:00 \
            -o "${LOGDIR}/transformer_${fmod_tag}_${fusion}_%j.out" \
            -e "${LOGDIR}/transformer_${fmod_tag}_${fusion}_%j.err" \
            --export=ALL \
            --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --model transformer --modalities $fmods --fusion $fusion $COMMON"
        echo "  transformer / $fmods / $fusion"
    done
done

echo ""
echo "Total: 41 jobs | Scene Recognition v3 | Train=14vols, Test=4vols"
echo "Results: $OUTDIR"