#!/usr/bin/env bash
# v4: 500-clip GPT-4o CoT → LoRA SFT (assistant mode) → eval on 100 disjoint clips.
# Estimated total: ~15 min CoT + ~2h train + ~10 min eval.
set -euo pipefail

cd "$(dirname "$0")/../.."
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/train500_cot.jsonl"
EVAL_CSV="data/vla_cot/eval100.csv"
CKPT_DIR="checkpoints/VLA/qwen_cot_v4"
INFER_OUT="eval_results/vla_cot_v4/eval100_scores.csv"
LOG_DIR="runs/vla_cot_v4"
mkdir -p "$LOG_DIR" "$(dirname "$INFER_OUT")"

N_TRAIN=500
SEED=0

# Skip eval100 IDs during CoT gen
SKIP_IDS="$(python -c "import pandas as pd; print(','.join(pd.read_csv('${EVAL_CSV}', dtype=str)['id'].str.zfill(5).tolist()))")"

echo "==== [1/3] GPT-4o CoT labels (n=${N_TRAIN}, resume-safe) ===="
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 8 \
    --seed ${SEED} \
    --skip_ids "${SKIP_IDS}" \
    2>&1 | tee "${LOG_DIR}/01_cot.log"

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

echo
echo "==== [3/3] Inference on eval100 ===="
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"