#!/usr/bin/env bash
# DPO v1: align HazardHead alert timing via Direct Preference Optimization.
#
# Stage flow:
#   Step 0: generate DPO preference pair manifests (fast, CPU)
#   Step 1: sanity check pair manifests
#   Step 2: DPO training (GPU)
#
# Usage:
#   bash training/DPO/train_dpo.sh           # full training
#   bash training/DPO/train_dpo.sh --debug   # smoke test
set -euo pipefail

ROOT=PROJECT_ROOT
MANIFEST_DIR="$ROOT/data/sft_manifests"
PAIR_DIR="$ROOT/data/dpo_pairs"
SFT_CHECKPOINT="$ROOT/checkpoints/SFT/sft_v2/best"
OUTPUT_DIR="$ROOT/checkpoints/DPO"
EXPERIMENT="dpo_v1"

BATCH_SIZE=4
GRAD_ACCUM=2      # effective batch = 8

DEBUG_FLAGS=""
if [[ "${1:-}" == "--debug" ]]; then
    DEBUG_FLAGS="--debug --debug_samples 64"
    EXPERIMENT="dpo_v1_debug"
    echo "=== DEBUG / SMOKE TEST MODE ==="
fi

# ── Step 0: generate DPO pair manifests ─────────────────────────────────────
if [[ ! -f "$PAIR_DIR/nexar_train.json" ]]; then
    echo "DPO pair manifests not found — generating..."
    python -m training.DPO.make_dpo_pairs \
        --manifest_dir "$MANIFEST_DIR" \
        --out_dir      "$PAIR_DIR"
fi

# ── Step 1: quick sanity check ───────────────────────────────────────────────
echo "Sanity-checking DPO pair manifests..."
python - <<'PYEOF'
import json, sys
from pathlib import Path

pair_dir = Path("data/dpo_pairs")
ok = True
for name in ["nexar_train.json", "dada_train.json", "nexar_val.json", "dada_val.json"]:
    p = pair_dir / name
    if not p.exists():
        print(f"  MISSING: {name}")
        ok = False
        continue
    d = json.loads(p.read_text())
    pairs = d.get("pairs", [])
    timing  = sum(1 for x in pairs if x["pair_type"] == "timing")
    category= sum(1 for x in pairs if x["pair_type"] == "category")
    print(f"  {name}: {len(pairs)} pairs  ({timing} timing, {category} category)")

if not ok:
    print("ERROR: missing manifests. Run make_dpo_pairs.py first.")
    sys.exit(1)
print("  ✅ Manifests OK")
PYEOF

# ── Step 2: DPO training ─────────────────────────────────────────────────────
echo "Starting DPO training..."
echo "  SFT checkpoint : $SFT_CHECKPOINT"
echo "  Pair dir       : $PAIR_DIR"
echo "  Output         : $OUTPUT_DIR/$EXPERIMENT"
echo "  batch_size     : $BATCH_SIZE  (grad_accum=$GRAD_ACCUM, eff_batch=$((BATCH_SIZE*GRAD_ACCUM)))"

python -m training.DPO.trainer \
    --sft_checkpoint  "$SFT_CHECKPOINT"  \
    --pair_dir        "$PAIR_DIR"        \
    --output_dir      "$OUTPUT_DIR"      \
    --experiment_name "$EXPERIMENT"      \
    --num_epochs      5                  \
    --batch_size      $BATCH_SIZE        \
    --gradient_accumulation_steps $GRAD_ACCUM \
    --learning_rate   5e-5               \
    --beta            0.1                \
    --lambda_reg      0.5                \
    --max_grad_norm   1.0                \
    --val_every_n_steps 500              \
    --use_wandb                          \
    $DEBUG_FLAGS