#!/bin/bash
# Action Recognition Round 3:
# Group A: 8 coarse + prev + window=4s + downsample=2 (more frames) — 9 jobs
# Group B: 20 fine + prev + window=4s — 9 jobs
# Group C: 8 coarse + prev + window=4s + smaller model h=64 — 3 best modalities
# Total: 21 jobs

PYTHON=python
BASEDIR=${PULSE_ROOT}
TRAIN_SCRIPT=${BASEDIR}/experiments/tasks/train_pred_cls.py
OUTDIR_A=${BASEDIR}/results/recog3a
OUTDIR_B=${BASEDIR}/results/recog3b
OUTDIR_C=${BASEDIR}/results/recog3c
mkdir -p ${OUTDIR_A}/slurm_logs ${OUTDIR_B}/slurm_logs ${OUTDIR_C}/slurm_logs

MODS=("imu" "emg" "mocap" "emg,imu" "mocap,imu" "mocap,emg,imu" "mocap,emg,eyetrack" "mocap,emg,eyetrack,imu" "mocap,emg,eyetrack,imu,pressure")

# Group A: 8 coarse + prev + downsample=2 (gives ~100 frames for 2s segments at 100Hz)
COMMON_A="--mode recognition --coarse --use_prev_action --epochs 80 --batch_size 32 --lr 1e-3 --weight_decay 1e-4 --hidden_dim 128 --dropout 0.2 --downsample 2 --patience 20 --seed 42 --augment --noise_std 0.1 --time_mask_ratio 0.1 --label_smoothing 0.1 --window_sec 4.0"

echo "=== Group A: 8 coarse + prev + ds=2 ==="
for mods in "${MODS[@]}"; do
    mod_tag=$(echo $mods | tr ',' '-')
    sbatch \
        -J "rec3a_${mod_tag}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=4 \
        --mem=32G \
        -t 2:00:00 \
        -o "${OUTDIR_A}/slurm_logs/${mod_tag}_%j.out" \
        -e "${OUTDIR_A}/slurm_logs/${mod_tag}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $TRAIN_SCRIPT --modalities $mods $COMMON_A --output_dir $OUTDIR_A"
    echo "Submitted A: $mods"
done

# Group B: 20 fine + prev + window=4s + ds=5
COMMON_B="--mode recognition --use_prev_action --epochs 80 --batch_size 32 --lr 1e-3 --weight_decay 1e-4 --hidden_dim 128 --dropout 0.2 --downsample 5 --patience 20 --seed 42 --augment --noise_std 0.1 --time_mask_ratio 0.1 --label_smoothing 0.1 --window_sec 4.0"

echo ""
echo "=== Group B: 20 fine + prev ==="
for mods in "${MODS[@]}"; do
    mod_tag=$(echo $mods | tr ',' '-')
    sbatch \
        -J "rec3b_${mod_tag}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=4 \
        --mem=32G \
        -t 2:00:00 \
        -o "${OUTDIR_B}/slurm_logs/${mod_tag}_%j.out" \
        -e "${OUTDIR_B}/slurm_logs/${mod_tag}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $TRAIN_SCRIPT --modalities $mods $COMMON_B --output_dir $OUTDIR_B"
    echo "Submitted B: $mods"
done

# Group C: 8 coarse + prev + h=64 (less overfit) — top 3 from Group B round 2
COMMON_C="--mode recognition --coarse --use_prev_action --epochs 80 --batch_size 32 --lr 1e-3 --weight_decay 1e-4 --hidden_dim 64 --dropout 0.3 --downsample 5 --patience 20 --seed 42 --augment --noise_std 0.1 --time_mask_ratio 0.1 --label_smoothing 0.1 --window_sec 4.0"
TOP_MODS=("mocap,emg,eyetrack" "mocap,emg,imu" "imu")

echo ""
echo "=== Group C: 8 coarse + prev + h=64 ==="
for mods in "${TOP_MODS[@]}"; do
    mod_tag=$(echo $mods | tr ',' '-')
    sbatch \
        -J "rec3c_${mod_tag}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=4 \
        --mem=32G \
        -t 2:00:00 \
        -o "${OUTDIR_C}/slurm_logs/${mod_tag}_%j.out" \
        -e "${OUTDIR_C}/slurm_logs/${mod_tag}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $TRAIN_SCRIPT --modalities $mods $COMMON_C --output_dir $OUTDIR_C"
    echo "Submitted C: $mods"
done

echo ""
echo "Total: 21 jobs | Recognition Round 3"
echo "A (ds=2): $OUTDIR_A | B (20fine+prev): $OUTDIR_B | C (h=64): $OUTDIR_C"