#!/bin/bash
# Scene Recognition (Exp1 v5) - Only imu, mocap, emg
# Per-modality projection to 50d
# Train 14 vols / Test 4 vols

PYTHON=python
BASEDIR=${PULSE_ROOT}
SCRIPT=${BASEDIR}/experiments/train_exp1.py
OUTDIR=${BASEDIR}/results/exp1_v5
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"
MODELS=("cnn" "lstm" "transformer")

# Part 1: Single modality (3 mods × 3 backbones = 9 jobs)
echo "=== Part 1: Single Modality (9 jobs) ==="
for mods in "imu" "mocap" "emg"; do
    for model in "${MODELS[@]}"; do
        sbatch -J "e1v5_${model}_${mods}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
            --cpus-per-task=4 --mem=32G -t 2:00:00 \
            -o "${LOGDIR}/${model}_${mods}_early_%j.out" \
            -e "${LOGDIR}/${model}_${mods}_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: Multi-modality early fusion (4 combos × 3 backbones = 12 jobs)
echo ""
echo "=== Part 2: Multi-Modality Early Fusion (12 jobs) ==="
for mods in "imu,mocap" "imu,emg" "mocap,emg" "imu,mocap,emg"; do
    mod_tag=$(echo $mods | tr ',' '-')
    for model in "${MODELS[@]}"; do
        sbatch -J "e1v5_${model}_${mod_tag}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
            --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 3: Fusion ablation with imu+mocap+emg × transformer (7 jobs)
FUSIONS=("late" "attention" "weighted_late" "gated_late" "stacking" "product" "moe")
echo ""
echo "=== Part 3: Fusion Ablation - transformer × imu+mocap+emg (7 jobs) ==="
for fusion in "${FUSIONS[@]}"; do
    sbatch -J "e1v5_tf_${fusion}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
        --cpus-per-task=4 --mem=32G -t 2:00:00 \
        -o "${LOGDIR}/transformer_imu-mocap-emg_${fusion}_%j.out" \
        -e "${LOGDIR}/transformer_imu-mocap-emg_${fusion}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --model transformer --modalities imu,mocap,emg --fusion $fusion $COMMON"
    echo "  transformer / imu,mocap,emg / $fusion"
done

echo ""
echo "Total: 28 jobs | 3 modalities: imu(160d→50d), mocap(156d→50d), emg(8d→50d)"
echo "Results: $OUTDIR"