#!/bin/bash
# SLURM launcher for T10 Triplet Next-Action Prediction experiments.
#
# Produces all five tables from the paper plan:
#   Table 1: main comparison (T_fut=2s)     — 1 model × 5 seeds
#   Table 3: horizon curve                  — 5 horizons × 5 seeds (same model)
#   Table 4: modality ablation              — 6 configs × 5 seeds (ours only)
#   Table 5: component ablation             — 5 variants × 5 seeds (ours only)
#   Table 7: missing-modality robustness    — trained once w/ modality dropout,
#                                             evaluated under 6 test-time drops
#
# ~140 jobs in total. Uses `gpuHygonZ100` (2 idle nodes); change PARTITION to
# `gpuA800` if larger slots are available.
#
# Usage:
#     bash experiments/run_seqpred_all.sh
#     bash experiments/run_seqpred_all.sh --dry    # print what would submit
#
# Outputs: results/seqpred/<config>_<seedX>/{config.json, results.json,
#                                             model_best.pt}
# Aggregate into tables with experiments/analysis/aggregate_seqpred.py (TBD).

set -euo pipefail

DRY=${1:-}
PYTHON=${PYTHON:-python3}
BASEDIR=${BASEDIR:-${PULSE_ROOT}}
TRAIN=${BASEDIR}/experiments/tasks/train_seqpred.py
OUTDIR=${BASEDIR}/results/seqpred
LOGDIR=${OUTDIR}/slurm_logs
mkdir -p "${LOGDIR}"

PARTITION=${PARTITION:-gpuHygonZ100}
GPU_GRES=${GPU_GRES:-gpu:1}
CPUS=${CPUS:-4}
MEM=${MEM:-48G}
TIME=${TIME:-6:00:00}

BASE_ARGS="--epochs 40 --batch_size 32 --lr 3e-4 --weight_decay 1e-4 \
           --dropout 0.2 --patience 12 --label_smoothing 0.05 \
           --use_class_weights --num_workers 2"

ALL_MODS="imu,emg,eyetrack,mocap,pressure"

submit() {
    local JOB_NAME=$1
    local OUT_SUB=$2
    shift 2
    local CMD="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; \
               ${PYTHON} ${TRAIN} $* --output_dir ${OUTDIR}/${OUT_SUB}"
    if [[ "${DRY}" == "--dry" ]]; then
        echo "--- ${JOB_NAME} ---"
        echo "  out: ${OUTDIR}/${OUT_SUB}"
        echo "  $*"
        return
    fi
    sbatch \
        -J "sp_${JOB_NAME}" \
        -p "${PARTITION}" \
        --gres="${GPU_GRES}" \
        -N 1 -n 1 \
        --cpus-per-task=${CPUS} \
        --mem=${MEM} \
        -t "${TIME}" \
        -o "${LOGDIR}/${JOB_NAME}_%j.out" \
        -e "${LOGDIR}/${JOB_NAME}_%j.err" \
        --export=ALL \
        --wrap="${CMD}"
    echo "submitted: ${JOB_NAME} -> ${OUT_SUB}"
}

SEEDS=(42 123 456 789 1024)

# ---------------------------------------------------------------------
# Table 1: main comparison at T_fut=2s
#   Baselines (B1..B8) run on their preferred modality subsets;
#   DailyActFormer runs on ALL 5 modalities.
# ---------------------------------------------------------------------
echo "=== Table 1: main comparison ==="

for seed in "${SEEDS[@]}"; do
    # --- our model, full 5-modality ---
    submit "t1_ours_all5_s${seed}" "t1_ours_all5/seed${seed}" \
        --model dailyactformer --modalities ${ALL_MODS} \
        --t_obs 8 --t_fut 2 --seed ${seed} ${BASE_ARGS}

    # --- DeepConvLSTM (IMU only) ---
    submit "t1_dcl_imu_s${seed}" "t1_dcl_imu/seed${seed}" \
        --model deepconvlstm --modalities imu \
        --t_obs 8 --t_fut 2 --seed ${seed} ${BASE_ARGS}

    # --- DeepConvLSTM (IMU+MoCap+EMG, best 3-modality for baselines) ---
    submit "t1_dcl_3mod_s${seed}" "t1_dcl_3mod/seed${seed}" \
        --model deepconvlstm --modalities imu,mocap,emg \
        --t_obs 8 --t_fut 2 --seed ${seed} ${BASE_ARGS}
done

# ---------------------------------------------------------------------
# Table 3: horizon curve (our model only, 5 horizons × 5 seeds = 25 jobs)
# ---------------------------------------------------------------------
echo ""
echo "=== Table 3: horizon curve ==="
for tfut in 1 2 5 10 15; do
    for seed in "${SEEDS[@]}"; do
        submit "t3_ours_tfut${tfut}_s${seed}" \
            "t3_ours_tfut${tfut}/seed${seed}" \
            --model dailyactformer --modalities ${ALL_MODS} \
            --t_obs 8 --t_fut ${tfut} --seed ${seed} ${BASE_ARGS}
    done
done

# ---------------------------------------------------------------------
# Table 4: modality ablation on our model (remove one modality at a time)
# ---------------------------------------------------------------------
echo ""
echo "=== Table 4: modality ablation ==="
declare -A ABLATIONS
ABLATIONS["noPressure"]="imu,emg,eyetrack,mocap"
ABLATIONS["noEyeTrack"]="imu,emg,mocap,pressure"
ABLATIONS["noEMG"]="imu,eyetrack,mocap,pressure"
ABLATIONS["noIMU"]="emg,eyetrack,mocap,pressure"
ABLATIONS["noMoCap"]="imu,emg,eyetrack,pressure"
ABLATIONS["onlyIMU_EMG"]="imu,emg"
ABLATIONS["onlyMoCap"]="mocap"
ABLATIONS["onlyEMG"]="emg"
for tag in "${!ABLATIONS[@]}"; do
    mods="${ABLATIONS[$tag]}"
    for seed in "${SEEDS[@]}"; do
        submit "t4_${tag}_s${seed}" "t4_${tag}/seed${seed}" \
            --model dailyactformer --modalities ${mods} \
            --t_obs 8 --t_fut 2 --seed ${seed} ${BASE_ARGS}
    done
done

# ---------------------------------------------------------------------
# Table 5: component ablation on our model
# (ablation switches TBD — parameter hooks need to be added to the model
# first. For now submit a placeholder using lambda weights.)
# ---------------------------------------------------------------------
echo ""
echo "=== Table 5: component ablation (placeholders) ==="
# 5a: no aux verb_composite head (set lambda to 0)
for seed in "${SEEDS[@]}"; do
    submit "t5_noComp_s${seed}" "t5_noComp/seed${seed}" \
        --model dailyactformer --modalities ${ALL_MODS} \
        --t_obs 8 --t_fut 2 --seed ${seed} ${BASE_ARGS} \
        --lambda_verb_composite 0.0
done
# 5b: equal-weight heads (remove our lambda prior)
for seed in "${SEEDS[@]}"; do
    submit "t5_equalLambda_s${seed}" "t5_equalLambda/seed${seed}" \
        --model dailyactformer --modalities ${ALL_MODS} \
        --t_obs 8 --t_fut 2 --seed ${seed} ${BASE_ARGS} \
        --lambda_verb_composite 1.0 --lambda_hand 1.0
done

# 5c/5d/5e (modality-stem / fusion / causal-mask toggles) require model
# plumbing — we'll add CLI flags later.

echo ""
echo "All done.  Inspect with:  squeue -u \$USER | head"