#!/bin/bash
# Scene Recognition (Exp1 v9) - Improvements over v8
# Changes: (A) augmentation, (B) feat_concat fusion, (C) pretrained branches
# All use transformer, imu+mocap+emg, no projection, 5 seeds

PYTHON=python
BASEDIR=${PULSE_ROOT}
SCRIPT=${BASEDIR}/experiments/train_exp1.py
OUTDIR=${BASEDIR}/results/exp1_v9
LOGDIR=${OUTDIR}/slurm_logs
mkdir -p $LOGDIR

BASE="--model transformer --modalities imu,mocap,emg --epochs 100 --batch_size 16 --lr 1e-3 --weight_decay 1e-4 --hidden_dim 128 --downsample 5 --patience 15 --proj_dim 0 --output_dir $OUTDIR"
SEEDS=(42 123 456 789 2024)

# Pretrained single-modality models (modality order: imu=0, mocap=1, emg=2)
PT_IMU=${PULSE_ROOT}/results/exp1_v7/transformer_imu_early/model_best.pt
PT_MOCAP=${PULSE_ROOT}/results/exp1_v8/transformer_mocap_early/model_best.pt
PT_EMG=${PULSE_ROOT}/results/exp1_v7/transformer_emg_early/model_best.pt

# Group A: late fusion + augmentation (5 seeds)
echo "=== A: late + augment (5 seeds) ==="
for seed in "${SEEDS[@]}"; do
    sbatch -J "v9_late_aug_s${seed}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
        --cpus-per-task=4 --mem=64G -t 2:00:00 \
        -o "${LOGDIR}/late_aug_s${seed}_%j.out" \
        -e "${LOGDIR}/late_aug_s${seed}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --fusion late --augment --seed $seed --tag aug_s${seed} $BASE"
    echo "  late+aug seed=$seed"
done

# Group B: feat_concat fusion (5 seeds)
echo ""
echo "=== B: feat_concat (5 seeds) ==="
for seed in "${SEEDS[@]}"; do
    sbatch -J "v9_fc_s${seed}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
        --cpus-per-task=4 --mem=64G -t 2:00:00 \
        -o "${LOGDIR}/feat_concat_s${seed}_%j.out" \
        -e "${LOGDIR}/feat_concat_s${seed}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --fusion feat_concat --seed $seed --tag s${seed} $BASE"
    echo "  feat_concat seed=$seed"
done

# Group C: feat_concat + augmentation (5 seeds)
echo ""
echo "=== C: feat_concat + augment (5 seeds) ==="
for seed in "${SEEDS[@]}"; do
    sbatch -J "v9_fc_aug_s${seed}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
        --cpus-per-task=4 --mem=64G -t 2:00:00 \
        -o "${LOGDIR}/feat_concat_aug_s${seed}_%j.out" \
        -e "${LOGDIR}/feat_concat_aug_s${seed}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --fusion feat_concat --augment --seed $seed --tag aug_s${seed} $BASE"
    echo "  feat_concat+aug seed=$seed"
done

# Group D: late + pretrained IMU branch (freeze_idx=0) (5 seeds)
echo ""
echo "=== D: late + pretrained IMU (5 seeds) ==="
for seed in "${SEEDS[@]}"; do
    sbatch -J "v9_late_pt_s${seed}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
        --cpus-per-task=4 --mem=64G -t 2:00:00 \
        -o "${LOGDIR}/late_pretrained_s${seed}_%j.out" \
        -e "${LOGDIR}/late_pretrained_s${seed}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --fusion late --pretrained_backbone $PT_IMU --freeze_backbone_idx 0 --seed $seed --tag pt_s${seed} $BASE"
    echo "  late+pretrained seed=$seed"
done

# Group E: late + augment + pretrained IMU (5 seeds)
echo ""
echo "=== E: late + augment + pretrained IMU (5 seeds) ==="
for seed in "${SEEDS[@]}"; do
    sbatch -J "v9_late_aug_pt_s${seed}" -p gpuA800 --gres=gpu:1 -N1 -n1 \
        --cpus-per-task=4 --mem=64G -t 2:00:00 \
        -o "${LOGDIR}/late_aug_pt_s${seed}_%j.out" \
        -e "${LOGDIR}/late_aug_pt_s${seed}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${BASEDIR}; $PYTHON $SCRIPT --fusion late --augment --pretrained_backbone $PT_IMU --freeze_backbone_idx 0 --seed $seed --tag aug_pt_s${seed} $BASE"
    echo "  late+aug+pretrained seed=$seed"
done

echo ""
echo "Total: 25 jobs | 5 groups × 5 seeds"
echo "Results: $OUTDIR"