#!/usr/bin/env bash
# Nexar Collision Prediction — Full Pipeline
#
# Steps:
#   1. Extract features from test clips      (~36 min with batch=8)
#   2. Extract features from train positive  (~20 min with batch=8)
#   3. Extract features from train negative  (~20 min with batch=8)
#   4. Train NexarTemporalHead on Nexar data (~5  min, CPU-friendly)
#   5. Generate submission (zero-shot + trained + ensemble)
#   6. Evaluate locally against solution.csv
#
# Usage:
#   bash training/Nexar/run_nexar.sh              # full pipeline
#   bash training/Nexar/run_nexar.sh --zero-shot  # only steps 1 + 5 (no training)
#   bash training/Nexar/run_nexar.sh --debug      # 20 clips, smoke test
set -euo pipefail

ROOT=PROJECT_ROOT
SFT_CKPT="$ROOT/checkpoints/SFT/sft_v2/best"
POLICY_CKPT="$ROOT/checkpoints/Policy/policy_warmstart_v2/best"

NEXAR_TEST="$ROOT/nexar-collision-prediction/test"
NEXAR_TRAIN_POS="$ROOT/NEXAR_COLLISION/train/positive"
NEXAR_TRAIN_NEG="$ROOT/NEXAR_COLLISION/train/negative"
BASELINE_CSV="$ROOT/NEXAR_COLLISION/sample_submission.csv"
SOLUTION_CSV="$ROOT/NEXAR_COLLISION/solution.csv"
TEST_CSV="$ROOT/nexar-collision-prediction/test.csv"

CACHE_DIR="$ROOT/data/nexar_cache"
OUTPUT_DIR="$ROOT/checkpoints/Nexar"
SUBMISSION_DIR="$ROOT/submissions"

N_WINDOWS=3
WINDOW_DUR=3.0
N_FRAMES=8
BATCH=8
EPOCHS=30
TRAIN_BATCH=128
LR=3e-4

ZERO_SHOT_ONLY=false
EXTRACTOR_FLAGS=""
TRAIN_EXTRACTOR_FLAGS=""
MAX_CLIPS=0

# Parse args
for arg in "$@"; do
    case $arg in
        --zero-shot) ZERO_SHOT_ONLY=true ;;
        --debug)
            EXTRACTOR_FLAGS="--max_clips 20"
            TRAIN_EXTRACTOR_FLAGS="--max_clips 6"
            MAX_CLIPS=20
            BATCH=4
            EPOCHS=5
            TRAIN_BATCH=16
            echo "=== DEBUG MODE (20 clips) ==="
            ;;
    esac
done

mkdir -p "$CACHE_DIR" "$OUTPUT_DIR" "$SUBMISSION_DIR"
cd "$ROOT"

# ── Step 1: Extract features from TEST clips ─────────────────────────────────
echo ""
echo "Step 1: Extracting test features ..."
python -m training.Nexar.nexar_extractor \
    --sft_checkpoint    "$SFT_CKPT"             \
    --policy_checkpoint "$POLICY_CKPT"           \
    --video_dir         "$NEXAR_TEST"            \
    --out_file          "$CACHE_DIR/test.pt"     \
    --n_windows         $N_WINDOWS               \
    --window_dur        $WINDOW_DUR              \
    --n_frames          $N_FRAMES                \
    --batch_size        $BATCH                   \
    $EXTRACTOR_FLAGS

# ── Step 2 (optional): Zero-shot submission — no training ─────────────────────
echo ""
echo "Step 2: Generating zero-shot submission ..."
for AGG in max_last weighted tta; do
    python -m training.Nexar.nexar_submit \
        --mode          zero_shot                                      \
        --test_cache    "$CACHE_DIR/test.pt"                           \
        --zero_shot_agg $AGG                                           \
        --n_windows     $N_WINDOWS                                     \
        --out_csv       "$SUBMISSION_DIR/nexar_zero_shot_${AGG}.csv"   \
        --test_csv      "$TEST_CSV"                                    \
        --evaluate      "$SOLUTION_CSV"
done

if [[ "$ZERO_SHOT_ONLY" == "true" ]]; then
    echo "Zero-shot only mode — done."
    exit 0
fi

# ── Step 3 & 4: Extract features from TRAIN videos (TTE-aligned) ─────────────
echo ""
echo "Step 3-4: Extracting train features (TTE-aligned for positive clips) ..."
python -m training.Nexar.nexar_train_extractor \
    --sft_checkpoint    "$SFT_CKPT"                    \
    --policy_checkpoint "$POLICY_CKPT"                  \
    --train_csv         "nexar-collision-prediction/train.csv" \
    --train_pos_dir     "$NEXAR_TRAIN_POS"             \
    --train_neg_dir     "$NEXAR_TRAIN_NEG"             \
    --out_dir           "$CACHE_DIR"                   \
    --n_windows         $N_WINDOWS                     \
    --window_dur        $WINDOW_DUR                    \
    --n_frames          $N_FRAMES                      \
    --batch_size        $BATCH                         \
    $TRAIN_EXTRACTOR_FLAGS

# ── Step 5: Train NexarTemporalHead ───────────────────────────────────────────
echo ""
echo "Step 5: Training NexarTemporalHead ..."
python -m training.Nexar.nexar_trainer \
    --cache_pos   "$CACHE_DIR/train_positive.pt"  \
    --cache_neg   "$CACHE_DIR/train_negative.pt"  \
    --output_dir  "$OUTPUT_DIR/nexar_temporal_v1" \
    --arch        temporal                         \
    --n_windows   $N_WINDOWS                      \
    --epochs      $EPOCHS                         \
    --batch_size  $TRAIN_BATCH                    \
    --lr          $LR

# ── Step 6: Generate trained + ensemble submissions ───────────────────────────
echo ""
echo "Step 6: Generating trained submission ..."
python -m training.Nexar.nexar_submit \
    --mode       trained                                       \
    --test_cache "$CACHE_DIR/test.pt"                          \
    --model_dir  "$OUTPUT_DIR/nexar_temporal_v1"               \
    --n_windows  $N_WINDOWS                                    \
    --out_csv    "$SUBMISSION_DIR/nexar_trained.csv"           \
    --test_csv   "$TEST_CSV"                                   \
    --evaluate   "$SOLUTION_CSV"

echo ""
echo "Step 6b: Generating ensemble submissions (varying alpha) ..."
for ALPHA in 0.3 0.5 0.7; do
    python -m training.Nexar.nexar_submit \
        --mode           ensemble                                           \
        --test_cache     "$CACHE_DIR/test.pt"                               \
        --model_dir      "$OUTPUT_DIR/nexar_temporal_v1"                    \
        --baseline_csv   "$BASELINE_CSV"                                    \
        --ensemble_alpha $ALPHA                                             \
        --n_windows      $N_WINDOWS                                         \
        --out_csv        "$SUBMISSION_DIR/nexar_ensemble_a${ALPHA/./_}.csv" \
        --test_csv       "$TEST_CSV"                                        \
        --evaluate       "$SOLUTION_CSV"
done

echo ""
echo "✅ Nexar pipeline complete."
echo "   Submissions in: $SUBMISSION_DIR/"
echo ""
echo "   Evaluate any submission:"
echo "   python NEXAR_COLLISION/evaluate_submission.py SUBMISSION.csv NEXAR_COLLISION/solution.csv"