#!/bin/bash
# ============================================================
# Run all 6 published baseline models across 4 experiments
# Submit to SLURM gpuA800 partition
# ============================================================

PYTHON=python3
BASEDIR=${PULSE_ROOT}
OUTBASE=${BASEDIR}/results/published_baselines_v2

SEED=42
ENV_SETUP="export PYTHONUNBUFFERED=1; export LD_LIBRARY_PATH=${PULSE_ROOT} cd ${BASEDIR}"

submit() {
    # $1=job_name $2=time $3=mem $4=command
    local LOGDIR="${OUTBASE}/slurm_logs"
    mkdir -p "$LOGDIR"
    sbatch -J "$1" -p gpuA800 --gres=gpu:1 -N1 -n1 \
        --cpus-per-task=4 --mem="$3" -t "$2" \
        -o "${LOGDIR}/${1}_%j.out" \
        -e "${LOGDIR}/${1}_%j.err" \
        --export=ALL \
        --wrap="${ENV_SETUP}; $4"
    echo "  Submitted: $1"
}

# ============================================================
# Exp1: Scene Recognition - DeepConvLSTM + InceptionTime
# ============================================================
echo "=== Exp1: Scene Recognition ==="
OUTDIR_E1=${OUTBASE}/exp1
EXP1_COMMON="--epochs 100 --batch_size 16 --lr 1e-3 --weight_decay 1e-4 --hidden_dim 32 --downsample 5 --patience 15 --seed $SEED --output_dir $OUTDIR_E1"

for model in deepconvlstm inceptiontime; do
    # Single modality
    for mod in imu mocap emg; do
        submit "e1_${model}_${mod}" "2:00:00" "32G" \
            "$PYTHON experiments/train_exp1.py --model $model --modalities $mod --fusion early $EXP1_COMMON"
    done
    # Multi-modal early + late
    submit "e1_${model}_ime_early" "2:00:00" "32G" \
        "$PYTHON experiments/train_exp1.py --model $model --modalities imu,mocap,emg --fusion early $EXP1_COMMON"
    submit "e1_${model}_ime_late" "2:00:00" "32G" \
        "$PYTHON experiments/train_exp1.py --model $model --modalities imu,mocap,emg --fusion late $EXP1_COMMON"
done
# Total Exp1: 2 models × (3 single + 2 multi) = 10 jobs

# ============================================================
# Exp2: Action Segmentation - MS-TCN++ + DiffAct
# ============================================================
echo ""
echo "=== Exp2: Action Segmentation ==="
OUTDIR_E2=${OUTBASE}/exp2
EXP2_COMMON="--epochs 80 --batch_size 16 --lr 5e-4 --weight_decay 1e-4 --hidden_dim 64 --downsample 2 --patience 15 --seed $SEED --output_dir $OUTDIR_E2"

for model in mstcnpp diffact; do
    for mods in mocap mocap,emg,eyetrack mocap,emg,eyetrack,imu mocap,emg,eyetrack,imu,pressure; do
        mod_tag=${mods//,/-}
        submit "e2_${model}_${mod_tag}" "6:00:00" "64G" \
            "$PYTHON experiments/train_exp2.py --model $model --modalities $mods $EXP2_COMMON"
    done
done
# Total Exp2: 2 models × 4 modality combos = 8 jobs

# ============================================================
# Exp3: Contact Detection - DeepConvLSTM + InceptionTime + UnderPressure
# ============================================================
echo ""
echo "=== Exp3: Contact Detection ==="
OUTDIR_E3=${OUTBASE}/exp3
EXP3_COMMON="--epochs 50 --batch_size 32 --lr 1e-3 --weight_decay 1e-4 --hidden_dim 64 --downsample 2 --patience 10 --seed $SEED --output_dir $OUTDIR_E3"

for model in deepconvlstm inceptiontime underpressure; do
    for mods in mocap emg imu mocap,emg mocap,emg,eyetrack,imu; do
        mod_tag=${mods//,/-}
        submit "e3_${model}_${mod_tag}" "4:00:00" "32G" \
            "$PYTHON experiments/train_exp3.py --model $model --modalities $mods $EXP3_COMMON"
    done
done
# Total Exp3: 3 models × 5 modality combos = 15 jobs

# ============================================================
# Exp4: Cross-Modal Prediction - UnderPressure (4a) + emg2pose (4b)
# ============================================================
echo ""
echo "=== Exp4: Cross-Modal Prediction ==="
OUTDIR_E4=${OUTBASE}/exp4
EXP4_COMMON="--epochs 50 --batch_size 32 --lr 5e-4 --weight_decay 1e-4 --hidden_dim 128 --downsample 2 --patience 10 --seed $SEED --output_dir $OUTDIR_E4"

# 4a: MoCap -> Pressure (UnderPressure)
submit "e4_4a_underpressure" "4:00:00" "32G" \
    "$PYTHON experiments/train_exp4.py --subtask 4a --model underpressure $EXP4_COMMON"

# 4b: EMG -> Hand Pose (emg2pose velocity + direct)
submit "e4_4b_emg2pose" "4:00:00" "32G" \
    "$PYTHON experiments/train_exp4.py --subtask 4b --model emg2pose $EXP4_COMMON"
submit "e4_4b_emg2pose_direct" "4:00:00" "32G" \
    "$PYTHON experiments/train_exp4.py --subtask 4b --model emg2pose_direct $EXP4_COMMON"
# Total Exp4: 3 jobs

echo ""
echo "=== Total: 36 jobs submitted ==="
echo "  Exp1: 10 jobs (DeepConvLSTM + InceptionTime)"
echo "  Exp2:  8 jobs (MS-TCN++ + DiffAct)"
echo "  Exp3: 15 jobs (DeepConvLSTM + InceptionTime + UnderPressure)"
echo "  Exp4:  3 jobs (UnderPressure + emg2pose)"
echo ""
echo "Monitor: squeue -u \$(whoami)"
echo "Results: ${OUTBASE}/"