#!/bin/bash
# Exp1 small3: Data augmentation + Frozen pretrained IMU + Label smoothing
# Goal: Break the IMU-alone F1=0.771 ceiling with emg+imu fusion
# Phase 0: pretrain IMU with hidden_dim=48 (matches fusion branch)
# Baselines: IMU+aug+ls, emg+imu early+aug+ls
# Group A: 7 fusion + aug + ls (no freeze)
# Group B: 7 fusion + frozen IMU + ls (no aug)  [dep: phase0]
# Group C: 7 fusion + frozen IMU + aug + ls      [dep: phase0]
# Total: 1 + 2 + 7 + 7 + 7 = 24 jobs

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

COMMON="--model transformer --epochs 100 --batch_size 16 --lr 1e-3 --weight_decay 1e-3 --hidden_dim 32 --downsample 5 --patience 15 --seed 42"
FUSIONS=(late attention weighted_late gated_late stacking product moe)

# ============================================================
# Phase 0: Pretrain IMU with hidden_dim=48 (matches fusion branch)
# ============================================================
PHASE0_JOB=$(sbatch --parsable \
    -J "s3_phase0_imu48" \
    -p gpuA800 \
    --gres=gpu:1 \
    -N 1 -n 1 \
    --cpus-per-task=8 \
    --mem=32G \
    -t 1:00:00 \
    -o "${LOGDIR}/phase0_imu48_%j.out" \
    -e "${LOGDIR}/phase0_imu48_%j.err" \
    --export=ALL \
    --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --model transformer --fusion early --modalities imu --hidden_dim 48 --epochs 100 --batch_size 16 --lr 1e-3 --weight_decay 1e-3 --downsample 5 --patience 15 --seed 42 --output_dir ${OUTDIR}/phase0")
echo "Phase 0 (IMU h48): job $PHASE0_JOB"

PRETRAINED="${OUTDIR}/phase0/transformer_imu_early/model_best.pt"

# ============================================================
# Baselines (no dependency)
# ============================================================

# Baseline 1: IMU alone + augment + label_smoothing
sbatch \
    -J "s3_bl_imu_aug" \
    -p gpuA800 \
    --gres=gpu:1 \
    -N 1 -n 1 \
    --cpus-per-task=8 \
    --mem=32G \
    -t 1:00:00 \
    -o "${LOGDIR}/bl_imu_aug_%j.out" \
    -e "${LOGDIR}/bl_imu_aug_%j.err" \
    --export=ALL \
    --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion early --modalities imu $COMMON --augment --label_smoothing 0.1 --tag bl_aug --output_dir $OUTDIR"
echo "Submitted: baseline IMU+aug+ls"

# Baseline 2: emg,imu early + augment + label_smoothing
sbatch \
    -J "s3_bl_ei_aug" \
    -p gpuA800 \
    --gres=gpu:1 \
    -N 1 -n 1 \
    --cpus-per-task=8 \
    --mem=32G \
    -t 1:00:00 \
    -o "${LOGDIR}/bl_ei_aug_%j.out" \
    -e "${LOGDIR}/bl_ei_aug_%j.err" \
    --export=ALL \
    --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion early --modalities emg,imu $COMMON --augment --label_smoothing 0.1 --tag bl_aug --output_dir $OUTDIR"
echo "Submitted: baseline emg+imu early+aug+ls"

# ============================================================
# Group A: emg+imu x 7 fusion + augment + label_smoothing (no freeze)
# ============================================================
for fusion in "${FUSIONS[@]}"; do
    sbatch \
        -J "s3_A_${fusion}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=8 \
        --mem=32G \
        -t 1:00:00 \
        -o "${LOGDIR}/grpA_${fusion}_%j.out" \
        -e "${LOGDIR}/grpA_${fusion}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion $fusion --modalities emg,imu $COMMON --augment --label_smoothing 0.1 --tag grpA --output_dir $OUTDIR"
    echo "Submitted: Group A $fusion"
done

# ============================================================
# Group B: emg+imu x 7 fusion + frozen IMU + label_smoothing (no augment)
# Depends on Phase 0
# ============================================================
for fusion in "${FUSIONS[@]}"; do
    sbatch \
        --dependency=afterok:${PHASE0_JOB} \
        -J "s3_B_${fusion}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=8 \
        --mem=32G \
        -t 1:00:00 \
        -o "${LOGDIR}/grpB_${fusion}_%j.out" \
        -e "${LOGDIR}/grpB_${fusion}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion $fusion --modalities emg,imu $COMMON --label_smoothing 0.1 --pretrained_backbone $PRETRAINED --freeze_backbone_idx 1 --tag grpB --output_dir $OUTDIR"
    echo "Submitted: Group B $fusion (dep: $PHASE0_JOB)"
done

# ============================================================
# Group C: emg+imu x 7 fusion + frozen IMU + augment + label_smoothing
# Depends on Phase 0
# ============================================================
for fusion in "${FUSIONS[@]}"; do
    sbatch \
        --dependency=afterok:${PHASE0_JOB} \
        -J "s3_C_${fusion}" \
        -p gpuA800 \
        --gres=gpu:1 \
        -N 1 -n 1 \
        --cpus-per-task=8 \
        --mem=32G \
        -t 1:00:00 \
        -o "${LOGDIR}/grpC_${fusion}_%j.out" \
        -e "${LOGDIR}/grpC_${fusion}_%j.err" \
        --export=ALL \
        --wrap="export PYTHONUNBUFFERED=1; cd ${PULSE_ROOT}; $PYTHON $SCRIPT --fusion $fusion --modalities emg,imu $COMMON --augment --label_smoothing 0.1 --pretrained_backbone $PRETRAINED --freeze_backbone_idx 1 --tag grpC --output_dir $OUTDIR"
    echo "Submitted: Group C $fusion (dep: $PHASE0_JOB)"
done

echo ""
echo "Total: 1 phase0 + 2 baselines + 7 grpA + 7 grpB + 7 grpC = 24 jobs"
echo "Results: $OUTDIR"
echo "Phase 0 job ID: $PHASE0_JOB (Groups B & C depend on it)"