#!/bin/bash
# =============================================================================
# Run all deterministic baselines: train, predict, evaluate
# =============================================================================
set -e

DATA_ROOT="data"
EPOCHS=100
BATCH_SIZE=32
LR=0.001

echo "============================================================"
echo "Deterministic Segmentation Baselines for LIDC-IDRI"
echo "============================================================"

# Step 1: Prepare data (if not already done)
if [ ! -d "${DATA_ROOT}/flat_train/images" ]; then
    echo ""
    echo "[Step 1] Preparing flat dataset with majority-vote masks..."
    python baselines/prepare_data.py --data_root ${DATA_ROOT} --skip_nnunet
else
    echo "[Step 1] Flat dataset already exists, skipping..."
fi

# Step 2: Train and predict for each model
MODELS=("unet" "attention_unet" "unetpp" "transunet")

for MODEL in "${MODELS[@]}"; do
    echo ""
    echo "============================================================"
    echo "[Step 2] Training ${MODEL}..."
    echo "============================================================"
    
    python baselines/train.py \
        --model ${MODEL} \
        --data_dir ${DATA_ROOT}/flat_train \
        --epochs ${EPOCHS} \
        --batch_size ${BATCH_SIZE} \
        --lr ${LR} \
        --checkpoint_dir checkpoints
    
    echo ""
    echo "[Step 3] Predicting with ${MODEL}..."
    python baselines/predict.py \
        --model ${MODEL} \
        --checkpoint checkpoints/${MODEL}_best.pth \
        --test_dir ${DATA_ROOT}/flat_test \
        --output_dir results/${MODEL} \
        --num_samples 16
    
    echo ""
    echo "[Step 4] Evaluating ${MODEL}..."
    python evaluate.py \
        --samples_dir results/${MODEL} \
        --gt_dir ${DATA_ROOT}/testing \
        --results_file results/${MODEL}_eval.csv
    
    echo ""
    echo ">>> ${MODEL} complete!"
    echo ""
done

# Step 5: Print comparison summary
echo ""
echo "============================================================"
echo "FINAL RESULTS COMPARISON"
echo "============================================================"
python -c "
import pandas as pd
import os

models = ['unet', 'attention_unet', 'unetpp', 'transunet']
results = []

for m in models:
    f = f'results/{m}_eval.csv'
    if os.path.exists(f):
        df = pd.read_csv(f)
        avg = df[df['image_id'] == 'AVERAGE'].iloc[0]
        results.append({
            'Method': m,
            'GED ↓': f\"{avg['ged_iou_paper']:.4f}\",
            'CI ↑': f\"{avg['CI_Score_Paper']:.4f}\",
            'Dmax ↑': f\"{avg['D_max_Paper']:.4f}\",
            'Dice': f\"{avg['avg_dice']:.4f}\",
            'IoU': f\"{avg['avg_iou']:.4f}\",
        })

if results:
    df = pd.DataFrame(results)
    print(df.to_string(index=False))
else:
    print('No results found!')
"

echo ""
echo "All baselines complete!"