#!/usr/bin/env bash
# MViT-v2-s Fine-tuning Pipeline for Nexar Collision Prediction
#
# Replicates the 1st-place approach (0.898 on private LB):
#   - MViT-v2-s pretrained on Kinetics-400
#   - Binary classification head
#   - Data-centric filtering (remove short-warning positives)
#   - Full fine-tuning with differential LR
#
# Usage:
#   bash training/Nexar/run_mvit.sh              # full training
#   bash training/Nexar/run_mvit.sh --debug      # 2 epochs, 16 samples
#   bash training/Nexar/run_mvit.sh --strict     # more aggressive data filtering
set -euo pipefail

ROOT=PROJECT_ROOT
TRAIN_CSV="$ROOT/nexar-collision-prediction/train.csv"
# Use flat train dir (both pos/neg in same dir; dataset reads labels from CSV)
TRAIN_DIR="$ROOT/nexar-collision-prediction/train"
TEST_DIR="$ROOT/nexar-collision-prediction/test"
TEST_CSV="$ROOT/nexar-collision-prediction/test.csv"
BASELINE_CSV="$ROOT/NEXAR_COLLISION/sample_submission.csv"
SOLUTION_CSV="$ROOT/NEXAR_COLLISION/solution.csv"

OUTPUT_BASE="$ROOT/checkpoints/Nexar"
SUBMISSION_DIR="$ROOT/submissions"

# Hyperparams (1st place config)
EPOCHS=20
BATCH=8
LR=5e-5
LR_MIN=1e-7
MIN_WARNING=0.3    # filter positives with warning window < 0.3s
PATIENCE=6
N_FRAMES=16
IMG_SIZE=224

DEBUG=false
STRICT=false

for arg in "$@"; do
    case $arg in
        --debug)
            DEBUG=true
            EPOCHS=2
            BATCH=4
            echo "=== DEBUG MODE ==="
            ;;
        --strict)
            STRICT=true
            MIN_WARNING=1.0
            echo "=== STRICT DATA FILTERING (min_warning=1.0s) ==="
            ;;
    esac
done

mkdir -p "$OUTPUT_BASE" "$SUBMISSION_DIR"
cd "$ROOT"

# ── Stage 1: Fine-tune MViT-v2-s ─────────────────────────────────────────────
EXP_NAME="mvit_v2_s_mw${MIN_WARNING/./_}"
CKPT_DIR="$OUTPUT_BASE/$EXP_NAME"

echo ""
echo "Training MViT-v2-s (min_warning=${MIN_WARNING}s) ..."
python -m training.Nexar.mvit_trainer \
    --train_csv   "$TRAIN_CSV"    \
    --video_dir   "$TRAIN_DIR"    \
    --output_dir  "$CKPT_DIR"     \
    --epochs      $EPOCHS         \
    --batch_size  $BATCH          \
    --lr          $LR             \
    --lr_min      $LR_MIN         \
    --min_warning $MIN_WARNING    \
    --patience    $PATIENCE       \
    --n_frames    $N_FRAMES       \
    --img_size    $IMG_SIZE

# ── Stage 2: Generate submissions ─────────────────────────────────────────────
echo ""
echo "Generating submissions ..."

# MViT only
python -m training.Nexar.mvit_submit \
    --model_dir  "$CKPT_DIR"                              \
    --test_dir   "$TEST_DIR"                              \
    --test_csv   "$TEST_CSV"                              \
    --batch_size 16                                       \
    --out_csv    "$SUBMISSION_DIR/mvit_${EXP_NAME}.csv"   \
    --evaluate   "$SOLUTION_CSV"

# Ensemble with baseline at various alphas
for ALPHA in 0.5 0.6 0.7 0.8; do
    OUT="$SUBMISSION_DIR/mvit_${EXP_NAME}_ensemble_a${ALPHA/./_}.csv"
    python -m training.Nexar.mvit_submit \
        --model_dir      "$CKPT_DIR"         \
        --test_dir       "$TEST_DIR"         \
        --test_csv       "$TEST_CSV"         \
        --batch_size     16                  \
        --baseline_csv   "$BASELINE_CSV"     \
        --ensemble_alpha $ALPHA              \
        --out_csv        "$OUT"              \
        --evaluate       "$SOLUTION_CSV"
done

# ── Optional Stage 3: Strict filtering run ────────────────────────────────────
if [[ "$STRICT" == "true" ]]; then
    echo ""
    echo "Stage 3: Strict data-filtered run (min_warning=1.0s) ..."
    STRICT_CKPT="$OUTPUT_BASE/mvit_v2_s_strict"
    python -m training.Nexar.mvit_trainer \
        --train_csv   "$TRAIN_CSV"    \
        --video_dir   "$TRAIN_DIR"    \
        --output_dir  "$STRICT_CKPT"  \
        --epochs      $EPOCHS         \
        --batch_size  $BATCH          \
        --lr          $LR             \
        --lr_min      $LR_MIN         \
        --min_warning 1.0             \
        --patience    $PATIENCE       \
        --n_frames    $N_FRAMES       \
        --img_size    $IMG_SIZE

    python -m training.Nexar.mvit_submit \
        --model_dir      "$STRICT_CKPT"                       \
        --test_dir       "$TEST_DIR"                          \
        --test_csv       "$TEST_CSV"                          \
        --batch_size     16                                   \
        --baseline_csv   "$BASELINE_CSV"                      \
        --ensemble_alpha 0.7                                  \
        --out_csv        "$SUBMISSION_DIR/mvit_strict_ensemble_0.7.csv" \
        --evaluate       "$SOLUTION_CSV"
fi

echo ""
echo "✅ MViT pipeline complete."
echo ""
echo "Submissions:"
ls -la "$SUBMISSION_DIR"/*.csv 2>/dev/null | tail -20
echo ""
echo "Evaluate any submission:"
echo "  python NEXAR_COLLISION/evaluate_submission.py SUBMISSION.csv NEXAR_COLLISION/solution.csv"