#!/bin/bash
# Exp1 small model: hidden_dim=32, dropout=0.5, weight_decay=1e-3
# 3 modalities: mocap, emg, imu (exclude pressure & eyetrack)
# Output: results/exp1_small

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

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

# ============================================================
# Part 1: Single modality (early fusion = single backbone)
# ============================================================
for mod in mocap emg imu; do
    job_name="s_${mod}"
    sbatch \
        -J "$job_name" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=8 \
        --mem=32G \
        -t 1:00:00 \
        -o "${LOGDIR}/${job_name}_%j.out" \
        -e "${LOGDIR}/${job_name}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion early --modalities $mod $COMMON"
    echo "Submitted: $job_name"
done

# ============================================================
# Part 2: Multi-modality early fusion (4 combos)
# ============================================================
EARLY_COMBOS=("mocap,emg" "mocap,imu" "emg,imu" "mocap,emg,imu")
for mods in "${EARLY_COMBOS[@]}"; do
    mod_tag=$(echo $mods | tr ',' '-')
    job_name="e_${mod_tag}"
    sbatch \
        -J "$job_name" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=8 \
        --mem=32G \
        -t 1:00:00 \
        -o "${LOGDIR}/${job_name}_%j.out" \
        -e "${LOGDIR}/${job_name}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion early --modalities $mods $COMMON"
    echo "Submitted: $job_name"
done

# ============================================================
# Part 3: Fusion methods x modality sets
# ============================================================
FUSIONS=(late attention weighted_late gated_late stacking product moe)
FUSION_MODS=("mocap,emg,imu" "mocap,imu")

for fusion in "${FUSIONS[@]}"; do
    for mods in "${FUSION_MODS[@]}"; do
        mod_tag=$(echo $mods | tr ',' '-')
        job_name="f_${fusion}_${mod_tag}"
        sbatch \
            -J "$job_name" \
            -p gpuA800 \
            --gres=gpu:1 \
            -N 1 -n 1 \
            --cpus-per-task=8 \
            --mem=32G \
            -t 1:00:00 \
            -o "${LOGDIR}/${job_name}_%j.out" \
            -e "${LOGDIR}/${job_name}_%j.err" \
            --export=ALL \
            --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion $fusion --modalities $mods $COMMON"
        echo "Submitted: $job_name"
    done
done

echo ""
echo "Total: 3 single + 4 early + 14 fusion = 21 jobs submitted!"
echo "Results will be saved to: $OUTDIR"