#!/usr/bin/env bash
# VLA + CoT smoke test — end-to-end on a 5090 in ~45-90 min.
#
# Steps:
#   1. Pick 30 train clips (15 pos + 15 neg) → GPT-4o CoT labels (~4 min, ~$0.30)
#   2. Pick a disjoint 20 clips for local eval       (list only — no CoT needed)
#   3. LoRA-train Qwen2.5-VL-3B on the 30-clip CoT set (~20-40 min)
#   4. Infer on the 20 eval clips, compute local AP/AUC
#
# Hard-fail on any error so we catch issues early.
set -euo pipefail

cd "$(dirname "$0")/../.."
ROOT="$(pwd)"
export PYTHONUNBUFFERED=1
export TOKENIZERS_PARALLELISM=false
export OPENAI_API_KEY="$(cat ~/Desktop/openai_api_key.txt | tr -d '[:space:]')"

TRAIN_CSV="nexar-collision-prediction/train.csv"
VIDEO_DIR="nexar-collision-prediction/train"
COT_OUT="data/vla_cot/smoke_train_cot.jsonl"
EVAL_CSV="data/vla_cot/smoke_eval.csv"
CKPT_DIR="checkpoints/VLA/qwen_cot_smoke"
INFER_OUT="eval_results/vla_cot_smoke/eval_scores.csv"
LOG_DIR="runs/vla_cot_smoke"
mkdir -p "$LOG_DIR" "$(dirname "$EVAL_CSV")" "$(dirname "$INFER_OUT")"

N_TRAIN=30    # CoT clips → teacher distil
N_EVAL=20     # local eval clips (disjoint)
SEED=0

echo "==== [1/4] Build eval split (disjoint from training) ===="
python - <<PY
import random, pandas as pd
from pathlib import Path
df = pd.read_csv("${TRAIN_CSV}", dtype={"id": str})
df["id"] = df["id"].astype(str).str.zfill(5)
rng = random.Random(${SEED})
pos = df[df["target"]==1]["id"].tolist()
neg = df[df["target"]==0]["id"].tolist()
rng.shuffle(pos); rng.shuffle(neg)
# Train half (handled by build_cot_labels)
train_ids = set(pos[:${N_TRAIN}//2] + neg[:${N_TRAIN}-${N_TRAIN}//2])
# Eval = next ${N_EVAL} disjoint, balanced
eval_pos = [x for x in pos if x not in train_ids][:${N_EVAL}//2]
eval_neg = [x for x in neg if x not in train_ids][:${N_EVAL}-${N_EVAL}//2]
eval_ids = eval_pos + eval_neg
eval_df = df[df["id"].isin(eval_ids)][["id","target"]]
Path("${EVAL_CSV}").parent.mkdir(parents=True, exist_ok=True)
eval_df.to_csv("${EVAL_CSV}", index=False)
with open("data/vla_cot/smoke_skip_ids.txt","w") as f:
    f.write(",".join(sorted(eval_ids)))
print(f"[split] train_pool=%d eval=%d (pos=%d neg=%d)" % (len(train_ids), len(eval_ids), len(eval_pos), len(eval_neg)))
PY
SKIP_IDS="$(cat data/vla_cot/smoke_skip_ids.txt)"

echo
echo "==== [2/4] GPT-4o CoT labels (teacher) ===="
python -m training.VLA.build_cot_labels \
    --train_csv "${TRAIN_CSV}" \
    --video_dir "${VIDEO_DIR}" \
    --out "${COT_OUT}" \
    --n_clips ${N_TRAIN} \
    --n_frames 8 \
    --resize_short 336 \
    --model gpt-4o \
    --detail low \
    --workers 4 \
    --seed ${SEED} \
    --skip_ids "${SKIP_IDS}" \
    2>&1 | tee "${LOG_DIR}/01_cot.log"

echo
echo "==== [3/4] LoRA-train Qwen2.5-VL-3B on CoT ===="
python -m training.VLA.train_vla_cot \
    --cot_jsonl "${COT_OUT}" \
    --video_dir "${VIDEO_DIR}" \
    --out_dir "${CKPT_DIR}" \
    --lora_r 32 --lora_alpha 16 --lora_dropout 0.05 \
    --lr 2e-4 \
    --epochs 3 \
    --batch_size 1 --grad_accum 4 \
    --n_frames 8 --resize_short 336 \
    --seed ${SEED} \
    2>&1 | tee "${LOG_DIR}/02_train.log"

echo
echo "==== [4/4] Inference + local AP ===="
python -m training.VLA.infer_vla_cot \
    --base_model Qwen/Qwen2.5-VL-3B-Instruct \
    --lora_dir "${CKPT_DIR}/best" \
    --video_dir "${VIDEO_DIR}" \
    --ids_csv "${EVAL_CSV}" \
    --out_csv "${INFER_OUT}" \
    --n_frames 8 --resize_short 336 \
    2>&1 | tee "${LOG_DIR}/03_infer.log"

echo
echo "==== DONE ===="
echo "Logs   : ${LOG_DIR}/"
echo "Scores : ${INFER_OUT}"
echo "Ckpt   : ${CKPT_DIR}/best"