#!/bin/bash
#SBATCH --job-name=ablation_fuse
#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/ablation_fusion_%j.log

# Test confidence-weighted and learned-weight fusion on all multi-modal combos
# Compare against existing mean fusion results

set -e
export PYTHONUNBUFFERED=1

PYTHON=python
BASEDIR=${PULSE_ROOT}
SCRIPT=${BASEDIR}/experiments/train_exp1.py
OUTDIR=${BASEDIR}/results/modality_ablation
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)

PT_IMU=${BASEDIR}/results/exp1_v7/transformer_imu_early/model_best.pt
PT_MOCAP=${BASEDIR}/results/exp1_v8/transformer_mocap_early/model_best.pt

echo "=== Ablation: Confidence & Learned Fusion ==="

# ============================================================
# GPU 0: confidence-weighted fusion
# ============================================================
(
export CUDA_VISIBLE_DEVICES=0

# mocap+imu / confidence / pretrained imu (idx=1)
echo "--- GPU0: mocap+imu / confidence ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap+imu confidence seed=$seed"
    $PYTHON $SCRIPT --modalities mocap,imu --fusion late --late_agg confidence \
        --seed $seed --pretrained_backbone $PT_IMU --freeze_backbone_idx 1 \
        --tag ablation_conf_s${seed} $COMMON 2>&1 | tail -3
done

# emg+imu / confidence / pretrained imu (idx=1)
echo "--- GPU0: emg+imu / confidence ---"
for seed in "${SEEDS[@]}"; do
    echo "  emg+imu confidence seed=$seed"
    $PYTHON $SCRIPT --modalities emg,imu --fusion late --late_agg confidence \
        --seed $seed --pretrained_backbone $PT_IMU --freeze_backbone_idx 1 \
        --tag ablation_conf_s${seed} $COMMON 2>&1 | tail -3
done

# mocap+emg / confidence / pretrained mocap (idx=0)
echo "--- GPU0: mocap+emg / confidence ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap+emg confidence seed=$seed"
    $PYTHON $SCRIPT --modalities mocap,emg --fusion late --late_agg confidence \
        --seed $seed --pretrained_backbone $PT_MOCAP --freeze_backbone_idx 0 \
        --tag ablation_conf_s${seed} $COMMON 2>&1 | tail -3
done

# mocap+emg+imu / confidence / pretrained imu (idx=2, modalities=mocap,emg,imu)
echo "--- GPU0: mocap+emg+imu / confidence ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap+emg+imu confidence seed=$seed"
    $PYTHON $SCRIPT --modalities imu,mocap,emg --fusion late --late_agg confidence \
        --seed $seed --pretrained_backbone $PT_IMU --freeze_backbone_idx 0 \
        --tag ablation_conf_s${seed} $COMMON 2>&1 | tail -3
done

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

# ============================================================
# GPU 1: learned-weight fusion
# ============================================================
(
export CUDA_VISIBLE_DEVICES=1

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

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

# mocap+emg / learned / pretrained mocap (idx=0)
echo "--- GPU1: mocap+emg / learned ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap+emg learned seed=$seed"
    $PYTHON $SCRIPT --modalities mocap,emg --fusion late --late_agg learned \
        --seed $seed --pretrained_backbone $PT_MOCAP --freeze_backbone_idx 0 \
        --tag ablation_lrn_s${seed} $COMMON 2>&1 | tail -3
done

# mocap+emg+imu / learned / pretrained imu (idx=0, modalities=imu,mocap,emg)
echo "--- GPU1: mocap+emg+imu / learned ---"
for seed in "${SEEDS[@]}"; do
    echo "  mocap+emg+imu learned seed=$seed"
    $PYTHON $SCRIPT --modalities imu,mocap,emg --fusion late --late_agg learned \
        --seed $seed --pretrained_backbone $PT_IMU --freeze_backbone_idx 0 \
        --tag ablation_lrn_s${seed} $COMMON 2>&1 | tail -3
done

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

wait $PID0 $PID1

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

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

configs = [
    # (label, pattern_template)
    # mean (from previous ablation run)
    ('mocap+imu / mean',      base + '/transformer_mocap-imu_late_ablation_pt_s{}/results.json'),
    ('mocap+imu / confidence', base + '/transformer_mocap-imu_late_ablation_conf_s{}/results.json'),
    ('mocap+imu / learned',   base + '/transformer_mocap-imu_late_ablation_lrn_s{}/results.json'),
    ('emg+imu / mean',        base + '/transformer_emg-imu_late_ablation_pt_s{}/results.json'),
    ('emg+imu / confidence',  base + '/transformer_emg-imu_late_ablation_conf_s{}/results.json'),
    ('emg+imu / learned',     base + '/transformer_emg-imu_late_ablation_lrn_s{}/results.json'),
    ('mocap+emg / mean',      base + '/transformer_mocap-emg_late_ablation_pt_s{}/results.json'),
    ('mocap+emg / confidence', base + '/transformer_mocap-emg_late_ablation_conf_s{}/results.json'),
    ('mocap+emg / learned',   base + '/transformer_mocap-emg_late_ablation_lrn_s{}/results.json'),
    ('3mod / mean',           v9_base + '/transformer_imu-mocap-emg_late_pt_s{}/results.json'),
    ('3mod / confidence',     base + '/transformer_imu-mocap-emg_late_ablation_conf_s{}/results.json'),
    ('3mod / learned',        base + '/transformer_imu-mocap-emg_late_ablation_lrn_s{}/results.json'),
]

print(f'{\"Config\":<30} {\"F1 (mean±std)\":<20} {\"Acc (mean±std)\":<20} N')
print('-' * 75)
for label, pat in configs:
    f1s, accs = [], []
    for s in seeds:
        path = pat.format(s)
        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 ==="