#!/bin/bash
#SBATCH --job-name=mod_ablation
#SBATCH --partition=gpuA800
#SBATCH --gres=gpu:2
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=4:00:00
#SBATCH --output=${PULSE_ROOT}/results/modality_ablation_%j.log

# Modality Ablation Matrix for Scene Recognition (Exp1)
# 7 configs: 3 single + 3 two-modal + 1 three-modal (already done)
# All use Transformer backbone, hidden_dim=128, 5 seeds
# Single modality: early fusion
# Multi modality: late fusion + pretrained strongest branch

set -e
export PYTHONUNBUFFERED=1

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

COMMON="--model transformer --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 backbones (seed=42, from v7/v8)
PT_IMU=${BASEDIR}/results/exp1_v7/transformer_imu_early/model_best.pt
PT_MOCAP=${BASEDIR}/results/exp1_v8/transformer_mocap_early/model_best.pt
PT_EMG=${BASEDIR}/results/exp1_v7/transformer_emg_early/model_best.pt

echo "=== Modality Ablation Matrix ==="
echo "Output: $OUTDIR"

# ============================================================
# GPU 0: Single modality (mocap, emg) + two-modal (mocap+emg)
# ============================================================
(
export CUDA_VISIBLE_DEVICES=0

# --- Phase 0: Single modality × 5 seeds ---
echo "--- GPU0: Single modality mocap ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap seed=$seed"
    $PYTHON $SCRIPT --modalities mocap --fusion early --seed $seed \
        --tag ablation_s${seed} $COMMON 2>&1 | tail -5
done

echo "--- GPU0: Single modality emg ---"
for seed in "${SEEDS[@]}"; do
    echo "  emg seed=$seed"
    $PYTHON $SCRIPT --modalities emg --fusion early --seed $seed \
        --tag ablation_s${seed} $COMMON 2>&1 | tail -5
done

# --- Phase 1: Two-modal mocap+emg / late+pretrained(emg) ---
# modalities=mocap,emg → idx0=mocap, idx1=emg → pretrain emg (idx=1)
echo "--- GPU0: mocap+emg late+pretrained ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap+emg seed=$seed"
    $PYTHON $SCRIPT --modalities mocap,emg --fusion late --seed $seed \
        --pretrained_backbone $PT_EMG --freeze_backbone_idx 1 \
        --tag ablation_pt_s${seed} $COMMON 2>&1 | tail -5
done

echo "--- GPU0 Done ---"
) &
PID0=$!

# ============================================================
# GPU 1: Two-modal (mocap+imu, emg+imu)
# ============================================================
(
export CUDA_VISIBLE_DEVICES=1

# --- mocap+imu / late+pretrained(imu) ---
# modalities=mocap,imu → idx0=mocap, idx1=imu → pretrain imu (idx=1)
echo "--- GPU1: mocap+imu late+pretrained ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap+imu seed=$seed"
    $PYTHON $SCRIPT --modalities mocap,imu --fusion late --seed $seed \
        --pretrained_backbone $PT_IMU --freeze_backbone_idx 1 \
        --tag ablation_pt_s${seed} $COMMON 2>&1 | tail -5
done

# --- emg+imu / late+pretrained(imu) ---
# modalities=emg,imu → idx0=emg, idx1=imu → pretrain imu (idx=1)
echo "--- GPU1: emg+imu late+pretrained ---"
for seed in "${SEEDS[@]}"; do
    echo "  emg+imu seed=$seed"
    $PYTHON $SCRIPT --modalities emg,imu --fusion late --seed $seed \
        --pretrained_backbone $PT_IMU --freeze_backbone_idx 1 \
        --tag ablation_pt_s${seed} $COMMON 2>&1 | tail -5
done

echo "--- GPU1 Done ---"
) &
PID1=$!

wait $PID0 $PID1

# ============================================================
# Collect results
# ============================================================
echo ""
echo "=== Results Summary ==="
$PYTHON -c "
import json, os, numpy as np

base = '$OUTDIR'
configs = [
    ('mocap / early', 'transformer_mocap_early_ablation_s{}'),
    ('emg / early', 'transformer_emg_early_ablation_s{}'),
    ('imu / early', None),  # from v8_multiseed
    ('mocap+emg / late+pt', 'transformer_mocap-emg_late_ablation_pt_s{}'),
    ('mocap+imu / late+pt', 'transformer_mocap-imu_late_ablation_pt_s{}'),
    ('emg+imu / late+pt', 'transformer_emg-imu_late_ablation_pt_s{}'),
    ('mocap+emg+imu / late+pt', None),  # from v9
]

seeds = [42, 123, 456, 789, 2024]
v8_base = '${BASEDIR}/results/exp1_v8_multiseed'
v9_base = '${BASEDIR}/results/exp1_v9'

print(f'{\"Config\":<30} {\"F1 (mean±std)\":<20} {\"Acc (mean±std)\":<20} N')
print('-' * 75)

for label, pattern in configs:
    f1s, accs = [], []
    for s in seeds:
        if label == 'imu / early':
            path = os.path.join(v8_base, f'transformer_imu_early_s{s}', 'results.json')
        elif label == 'mocap+emg+imu / late+pt':
            path = os.path.join(v9_base, f'transformer_imu-mocap-emg_late_pt_s{s}', 'results.json')
        else:
            path = os.path.join(base, pattern.format(s), 'results.json')
        if os.path.exists(path):
            with open(path) as f:
                d = json.load(f)
            f1s.append(d['test_macro_f1'])
            accs.append(d['test_accuracy'])
    if f1s:
        f1 = np.array(f1s)
        acc = np.array(accs)
        print(f'{label:<30} {f1.mean():.3f}±{f1.std():.3f}           {acc.mean():.3f}±{acc.std():.3f}           {len(f1s)}')
    else:
        print(f'{label:<30} (no results)')
"

echo ""
echo "=== All done ==="