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

# Published Baselines for DailyAct-5M
# ASFormer (Yi et al., BMVC 2021) - Temporal Segmentation & Contact Detection
# TinyHAR (Zhou et al., ISWC 2022 Best Paper) - Scene Recognition

set -e
PYTHON=python
PROJECT=${PULSE_ROOT}
cd $PROJECT

EXP1_OUT=$PROJECT/results/published_baselines/exp1_tinyhar
EXP2_OUT=$PROJECT/results/published_baselines/exp2_asformer
EXP3_OUT=$PROJECT/results/published_baselines/exp3_asformer
mkdir -p $EXP1_OUT $EXP2_OUT $EXP3_OUT

echo "=========================================="
echo "Published Baselines - $(date)"
echo "=========================================="

# ============================================================
# Group 1: TinyHAR for Scene Recognition (Exp 1)
# Run on GPU 0
# ============================================================
(
export CUDA_VISIBLE_DEVICES=0

echo ""
echo "=== [GPU0] Exp1: TinyHAR Scene Recognition ==="

# Single modalities
for MOD in imu mocap emg eyetrack pressure; do
    echo "--- TinyHAR / ${MOD} / early ---"
    $PYTHON experiments/train_exp1.py \
        --model tinyhar --modalities $MOD --fusion early \
        --hidden_dim 32 --epochs 100 --batch_size 16 \
        --lr 1e-3 --weight_decay 1e-3 --downsample 5 \
        --seed 42 --output_dir $EXP1_OUT \
        --tag published 2>&1 | tail -5
done

# Best multi-modal combos
for MOD in "emg,imu" "mocap,emg,imu" "mocap,emg,eyetrack,imu"; do
    echo "--- TinyHAR / ${MOD} / early ---"
    $PYTHON experiments/train_exp1.py \
        --model tinyhar --modalities $MOD --fusion early \
        --hidden_dim 32 --epochs 100 --batch_size 16 \
        --lr 1e-3 --weight_decay 1e-3 --downsample 5 \
        --seed 42 --output_dir $EXP1_OUT \
        --tag published 2>&1 | tail -5
done

# TinyHAR with late fusion (emg + imu)
for FUSION in late weighted_late feat_concat; do
    echo "--- TinyHAR / emg,imu / ${FUSION} ---"
    $PYTHON experiments/train_exp1.py \
        --model tinyhar --modalities emg,imu --fusion $FUSION \
        --hidden_dim 32 --epochs 100 --batch_size 16 \
        --lr 1e-3 --weight_decay 1e-3 --downsample 5 \
        --seed 42 --output_dir $EXP1_OUT \
        --tag published 2>&1 | tail -5
done

echo "[GPU0] TinyHAR experiments complete."
) &
PID_GPU0=$!


# ============================================================
# Group 2: ASFormer for Segmentation (Exp 2) + Contact (Exp 3)
# Run on GPU 1
# ============================================================
(
export CUDA_VISIBLE_DEVICES=1

echo ""
echo "=== [GPU1] Exp2: ASFormer Temporal Segmentation ==="

# Key modality combinations
for MOD in mocap emg "mocap,emg,eyetrack" "mocap,emg,eyetrack,imu" "mocap,emg,eyetrack,imu,pressure"; do
    echo "--- ASFormer / ${MOD} ---"
    $PYTHON experiments/train_exp2.py \
        --model asformer --modalities $MOD \
        --hidden_dim 64 --epochs 80 --batch_size 16 \
        --lr 5e-4 --weight_decay 1e-4 --downsample 2 \
        --seed 42 --output_dir $EXP2_OUT 2>&1 | tail -5
done

echo ""
echo "=== [GPU1] Exp3: ASFormer Contact Detection ==="

# Key modality combinations
for MOD in mocap emg imu "mocap,emg" "mocap,emg,eyetrack" "mocap,emg,eyetrack,imu"; do
    echo "--- ASFormer / ${MOD} ---"
    $PYTHON experiments/train_exp3.py \
        --model asformer --modalities $MOD \
        --hidden_dim 64 --epochs 50 --batch_size 32 \
        --lr 1e-3 --weight_decay 1e-4 --downsample 2 \
        --seed 42 --output_dir $EXP3_OUT 2>&1 | tail -5
done

echo "[GPU1] ASFormer experiments complete."
) &
PID_GPU1=$!

# Wait for both GPU groups
wait $PID_GPU0
wait $PID_GPU1

echo ""
echo "=========================================="
echo "All published baseline experiments complete - $(date)"
echo "=========================================="

# ============================================================
# Collect results summary
# ============================================================
echo ""
echo "=== Results Summary ==="

echo ""
echo "--- Exp1: TinyHAR Scene Recognition ---"
for f in $EXP1_OUT/*/results.json; do
    if [ -f "$f" ]; then
        $PYTHON -c "
import json
with open('$f') as fp:
    r = json.load(fp)
mods = ','.join(r.get('modalities', []))
fus = r.get('fusion', 'early')
f1 = r.get('test_macro_f1', 0)
acc = r.get('test_accuracy', 0)
print(f'  TinyHAR | {mods:<30} | {fus:<12} | F1={f1:.4f} Acc={acc:.4f}')
"
    fi
done

echo ""
echo "--- Exp2: ASFormer Temporal Segmentation ---"
for f in $EXP2_OUT/*/results.json; do
    if [ -f "$f" ]; then
        $PYTHON -c "
import json
with open('$f') as fp:
    r = json.load(fp)
mods = ','.join(r.get('modalities', []))
m = r.get('test_metrics', {})
print(f'  ASFormer | {mods:<35} | Acc={m.get(\"frame_acc\",0):.4f} F1={m.get(\"frame_f1\",0):.4f} Seg@50={m.get(\"seg_f1@50\",0):.4f}')
"
    fi
done

echo ""
echo "--- Exp3: ASFormer Contact Detection ---"
for f in $EXP3_OUT/*/results.json; do
    if [ -f "$f" ]; then
        $PYTHON -c "
import json
with open('$f') as fp:
    r = json.load(fp)
mods = ','.join(r.get('input_modalities', []))
m = r.get('test_metrics', {})
print(f'  ASFormer | {mods:<30} | R_F1={m.get(\"right_f1\",0):.4f} L_F1={m.get(\"left_f1\",0):.4f} Avg_F1={m.get(\"avg_f1\",0):.4f}')
"
    fi
done