#!/bin/bash
# Create folder structure for ALL rows across Tables 1, 3, 4, 5, 7 and
# freeze the current experiments/ code into each one. After this you can
# cd into any <table>/<row>/ and run ./run.sh to submit 5 SLURM seeds.
#
# Re-running this script is safe: it will re-freeze the code (overwrite the
# snapshot), but won't clobber any existing seeds/ outputs.
set -euo pipefail

BASEDIR=${BASEDIR:-${PULSE_ROOT}}
EXP=${BASEDIR}/experiments
SETUP="${EXP}/setup_row.sh"

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

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

row () {
    # $1=table  $2=row  $3=desc  $4=cli
    bash "${SETUP}" --table "$1" --row "$2" --desc "$3" --cli "$4 ${COMMON}"
}

# ============================================================
# Table 1: Main comparison at T_fut=2s
# ============================================================
T1=table1_main_comparison
cat > "${BASEDIR}/${T1}/README.md" <<'EOF'
# Table 1: Main Comparison (Next-Action Prediction, T_fut = 2 s)

Each baseline is run on its most favourable modality subset; our model
(DailyActFormer) uses all 5 synchronised modalities. 5 seeds per row;
report mean ± std of Verb fine Top-1/5, Noun Top-1/5, Hand Top-1, Action
Top-1 (= verb ∧ noun ∧ hand). Action Top-1 is the headline metric.

| Row | Method            | Family          | Modalities          |
|-----|-------------------|-----------------|---------------------|
| 01  | DailyActFormer    | cross-modal Trf | imu+emg+eye+mocap+P |
| 02  | DeepConvLSTM      | CNN+LSTM (IMU)  | imu                 |
| 03  | DeepConvLSTM 3mod | CNN+LSTM        | imu+mocap+emg       |
| 04  | RULSTM            | rolling LSTM    | imu+mocap           |
| 05  | FUTR              | long-term Trf   | mocap+imu+emg       |
| 06  | AFFT              | multimodal Trf  | imu+emg+eye+mocap   |
| 07  | HandFormer        | hand-pose Trf   | mocap (fingers)     |
| 08  | ActionLLM (LoRA)  | LLM-based       | imu+emg+eye         |
EOF

mkdir -p "${BASEDIR}/${T1}"
row ${T1} row01_ours_dailyactformer_all5 \
    "Our model, all 5 modalities (headline row)" \
    "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut 2"

row ${T1} row02_deepconvlstm_imu \
    "DeepConvLSTM on IMU only (classic HAR baseline)" \
    "--model deepconvlstm --modalities imu --t_obs 8 --t_fut 2"

row ${T1} row03_deepconvlstm_3mod \
    "DeepConvLSTM on IMU+MoCap+EMG (best 3-modality concat)" \
    "--model deepconvlstm --modalities imu,mocap,emg --t_obs 8 --t_fut 2"

row ${T1} row04_rulstm_imu_mocap \
    "RULSTM, rolling-unrolling LSTM (IMU + MoCap late fusion)" \
    "--model rulstm --modalities imu,mocap --t_obs 8 --t_fut 2"

row ${T1} row05_futr_3mod \
    "FUTR (causal transformer) on MoCap+IMU+EMG" \
    "--model futr --modalities mocap,imu,emg --t_obs 8 --t_fut 2"

row ${T1} row06_afft_4mod \
    "AFFT (anticipative feature fusion transformer) on 4 modalities" \
    "--model afft --modalities imu,emg,eyetrack,mocap --t_obs 8 --t_fut 2"

row ${T1} row07_handformer_mocap \
    "HandFormer (skeleton-only ECCV'24) on MoCap finger joints" \
    "--model handformer --modalities mocap --t_obs 8 --t_fut 2"

row ${T1} row08_actionllm_3mod \
    "ActionLLM (Qwen2.5-0.5B + LoRA) on IMU+EMG+EyeTrack" \
    "--model actionllm --modalities imu,emg,eyetrack --t_obs 8 --t_fut 2"

# ============================================================
# Table 3: Horizon curve (DailyActFormer)
# ============================================================
T3=table3_horizon_curve
mkdir -p "${BASEDIR}/${T3}"
cat > "${BASEDIR}/${T3}/README.md" <<'EOF'
# Table 3: Prediction Horizon Curve (DailyActFormer, all 5 modalities)

Same model, varying T_fut. Expect monotonic drop in Action Top-1 as
horizon grows; plot line graph in the paper alongside this table.
EOF
HORIZONS=(1 2 5 10 15)
for i in "${!HORIZONS[@]}"; do
    tfut="${HORIZONS[$i]}"
    idx=$(printf "%02d" $((i+1)))
    row ${T3} row${idx}_ours_tfut${tfut}s \
        "Our model at T_fut=${tfut}s" \
        "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut ${tfut}"
done

# ============================================================
# Table 4: Modality ablation on DailyActFormer (T_fut=2s)
# ============================================================
T4=table4_modality_ablation
mkdir -p "${BASEDIR}/${T4}"
cat > "${BASEDIR}/${T4}/README.md" <<'EOF'
# Table 4: Modality Ablation (DailyActFormer, T_fut = 2 s)

Same model, progressively remove modalities. Each row trained from scratch.
EOF
row ${T4} row01_full_5mod    "Full 5-modality (reference)"         "--model dailyactformer --modalities imu,emg,eyetrack,mocap,pressure --t_obs 8 --t_fut 2"
row ${T4} row02_no_pressure  "Drop pressure"                        "--model dailyactformer --modalities imu,emg,eyetrack,mocap          --t_obs 8 --t_fut 2"
row ${T4} row03_no_eyetrack  "Drop eye-tracking"                    "--model dailyactformer --modalities imu,emg,mocap,pressure          --t_obs 8 --t_fut 2"
row ${T4} row04_no_emg       "Drop EMG"                             "--model dailyactformer --modalities imu,eyetrack,mocap,pressure     --t_obs 8 --t_fut 2"
row ${T4} row05_no_imu       "Drop IMU"                             "--model dailyactformer --modalities emg,eyetrack,mocap,pressure     --t_obs 8 --t_fut 2"
row ${T4} row06_no_mocap     "Drop MoCap"                           "--model dailyactformer --modalities imu,emg,eyetrack,pressure       --t_obs 8 --t_fut 2"
row ${T4} row07_imu_emg_only "Only IMU + EMG (physiology-light)"    "--model dailyactformer --modalities imu,emg                         --t_obs 8 --t_fut 2"
row ${T4} row08_mocap_only   "Only MoCap (skeleton-only)"           "--model dailyactformer --modalities mocap                           --t_obs 8 --t_fut 2"

# ============================================================
# Table 5: Component ablation (DailyActFormer switches)
# ============================================================
T5=table5_component_ablation
mkdir -p "${BASEDIR}/${T5}"
cat > "${BASEDIR}/${T5}/README.md" <<'EOF'
# Table 5: Component Ablation (DailyActFormer, T_fut = 2 s)

Each row toggles one architectural/training component of our model.
Component flags are implemented as CLI switches on train_seqpred.py;
see models_seqpred.py for the corresponding model options.
EOF
row ${T5} row01_full \
    "Full model (reference)" \
    "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut 2"
row ${T5} row02_no_composite_head \
    "Drop the auxiliary verb-composite head (lambda=0)" \
    "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut 2 --lambda_verb_composite 0.0"
row ${T5} row03_equal_lambda \
    "Equal-weight all 4 heads (no prior on verb>hand)" \
    "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut 2 --lambda_verb_composite 1.0 --lambda_hand 1.0"
row ${T5} row04_no_class_weight \
    "No inverse-frequency class weighting" \
    "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut 2 --lambda_verb_composite 0.5"
# row04 re-exposes the default; the variable-off is the absence of --use_class_weights
# We patch this manually — strip the flag out of COMMON.
ROW_DIR="${BASEDIR}/${T5}/row04_no_class_weight/run.sh"
if [[ -e "${ROW_DIR}" ]]; then
    sed -i 's/--use_class_weights //g' "${ROW_DIR}"
fi

row ${T5} row05_no_label_smoothing \
    "Label smoothing off" \
    "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut 2 --label_smoothing 0.0"

# ============================================================
# Table 7: Missing-modality robustness (train once, eval 6 ways)
# ============================================================
T7=table7_missing_modality
mkdir -p "${BASEDIR}/${T7}"
cat > "${BASEDIR}/${T7}/README.md" <<'EOF'
# Table 7: Missing-Modality Robustness (T_fut = 2 s)

Train DailyActFormer with random per-modality dropout (p=0.3). At test time,
evaluate under 6 configurations: full / drop one modality each. Only the
training job has its own folder; eval uses the trained checkpoint to fill
multiple rows of the final table.
EOF
row ${T7} row01_train_with_modality_dropout \
    "DailyActFormer trained with --modality_dropout 0.3" \
    "--model dailyactformer --modalities ${ALL5} --t_obs 8 --t_fut 2 --modality_dropout 0.3"
# The 6 test-time configurations (full / no_P / no_E / no_emg / no_imu /
# no_mocap) will be produced by a separate eval script that loads the
# checkpoint from row01 and runs evaluate() with modality subsets. See
# experiments/tasks/eval_missing_modality.py (TBD).

echo ""
echo "[ok] Froze rows under:"
echo "     ${BASEDIR}/{${T1},${T3},${T4},${T5},${T7}}/"