#!/bin/bash
# Submit all 3 new benchmark experiments (A: missing modality, B: grip force
# regression, C: T5 text retrieval) in parallel to the gpuA800 partition.
# Each single-GPU job is sbatched independently.

set -u
PYTHON=python
BASEDIR=${PULSE_ROOT}
OUTROOT=${BASEDIR}/results/exp_new
mkdir -p ${OUTROOT}/slurm_logs

SUBMIT() {
    # args: job_name time_hrs cmd...
    local jname=$1; shift
    local hrs=$1; shift
    sbatch \
        -J "${jname}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=4 \
        --mem=32G \
        -t "${hrs}:00:00" \
        -o "${OUTROOT}/slurm_logs/${jname}_%j.out" \
        -e "${OUTROOT}/slurm_logs/${jname}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $*"
}

# ---------------------------------------------------------------------------
# Experiment A: Missing-modality robustness
# Train late-fusion Transformer with random modality dropout at p=0.3
# 5 seeds, all 5 modalities
# ---------------------------------------------------------------------------
echo "=== Exp A: Missing-modality robustness (5 jobs) ==="
for seed in 42 123 456 789 2024; do
    SUBMIT "expA_missing_seed${seed}" 2 \
        "$PYTHON experiments/tasks/train_exp_missing.py \
            --model transformer --fusion late \
            --modalities mocap,emg,eyetrack,imu,pressure \
            --mod_dropout_p 0.3 \
            --epochs 100 --batch_size 16 --lr 1e-3 --hidden_dim 128 \
            --patience 15 --augment \
            --seed ${seed} \
            --output_dir ${OUTROOT}/expA_missing \
            --tag ''"
    echo "  Submitted: expA_missing_seed${seed}"
done

# Baseline (no dropout) for comparison, same seeds
for seed in 42 123 456; do
    SUBMIT "expA_baseline_seed${seed}" 2 \
        "$PYTHON experiments/tasks/train_exp_missing.py \
            --model transformer --fusion late \
            --modalities mocap,emg,eyetrack,imu,pressure \
            --mod_dropout_p 0.0 \
            --epochs 100 --batch_size 16 --lr 1e-3 --hidden_dim 128 \
            --patience 15 --augment \
            --seed ${seed} \
            --output_dir ${OUTROOT}/expA_baseline \
            --tag ''"
    echo "  Submitted: expA_baseline_seed${seed}"
done

# ---------------------------------------------------------------------------
# Experiment B: Grip force regression (T4')
# 3 backbones x 3 modality configs x 3 seeds
# ---------------------------------------------------------------------------
echo ""
echo "=== Exp B: Grip force regression ==="
BACKBONES=("transformer" "lstm")
MOD_CONFIGS=(
    "emg"
    "mocap"
    "emg,imu"
    "mocap,emg,imu,eyetrack"
)
for bb in "${BACKBONES[@]}"; do
    for mods in "${MOD_CONFIGS[@]}"; do
        for seed in 42 123 456; do
            mod_tag=$(echo $mods | tr ',' '-')
            SUBMIT "expB_grip_${bb}_${mod_tag}_s${seed}" 1 \
                "$PYTHON experiments/tasks/train_exp_grip.py \
                    --backbone ${bb} --modalities ${mods} \
                    --epochs 60 --batch_size 8 --lr 1e-3 \
                    --hidden_dim 128 --patience 12 \
                    --seed ${seed} \
                    --output_dir ${OUTROOT}/expB_grip \
                    --tag ''"
            echo "  Submitted: expB_grip_${bb}_${mod_tag}_s${seed}"
        done
    done
done

# ---------------------------------------------------------------------------
# Experiment C: T5 text retrieval
# 2 modality configs x 3 seeds
# ---------------------------------------------------------------------------
echo ""
echo "=== Exp C: T5 text retrieval ==="
for mods in "mocap,emg,eyetrack,imu" "emg,imu" "mocap"; do
    for seed in 42 123 456; do
        mod_tag=$(echo $mods | tr ',' '-')
        SUBMIT "expC_retrieval_${mod_tag}_s${seed}" 1 \
            "$PYTHON experiments/tasks/train_exp_retrieval.py \
                --modalities ${mods} \
                --epochs 60 --batch_size 64 --lr 5e-4 \
                --hidden_dim 128 --emb_dim 128 \
                --seed ${seed} \
                --output_dir ${OUTROOT}/expC_retrieval \
                --tag ''"
        echo "  Submitted: expC_retrieval_${mod_tag}_s${seed}"
    done
done

echo ""
echo "All jobs submitted. Monitor with: squeue -u \$USER"
echo "Results in: ${OUTROOT}/"