#!/usr/bin/env bash
# Stage 1: Supervised 3-class policy warm-start for LKAlert.
#
# Stage flow:
#   Step 0: generate policy label manifests  (CPU, ~30s)
#   Step 1: sanity-check manifests
#   Step 2: build belief vector cache        (GPU, one-time, ~2-3 days for full set)
#   Step 3: warm-start training              (GPU, fast with cache, minutes per epoch)
#   Step 4: evaluate best checkpoint
#
# Usage:
#   bash training/Policy/train_policy.sh           # full training
#   bash training/Policy/train_policy.sh --debug   # smoke test (small subset)
set -euo pipefail

ROOT=PROJECT_ROOT
SFT_CHECKPOINT="$ROOT/checkpoints/SFT/sft_v2/best"
LABEL_DIR="$ROOT/data/policy_labels"
CACHE_DIR="$ROOT/data/belief_cache"
OUTPUT_DIR="$ROOT/checkpoints/Policy"
EXPERIMENT="policy_warmstart_v1"

# ── hyperparams ───────────────────────────────────────────────────────────────
# Cache mode: batch_size can be much larger (no VLM gradient storage)
CACHE_BATCH=8        # batch size for cache BUILD (VLM inference, memory-limited)
TRAIN_BATCH=256      # batch size for TRAINING on cached beliefs (tiny PolicyHead only)
GRAD_ACCUM=1
LR=1e-4
NUM_EPOCHS=20        # fast epochs since each step is just PolicyHead
VAL_EVERY=500

DEBUG_FLAGS=""
if [[ "${1:-}" == "--debug" ]]; then
    DEBUG_FLAGS="--debug --debug_samples 128"
    EXPERIMENT="policy_warmstart_debug"
    CACHE_BATCH=4
    TRAIN_BATCH=32
    NUM_EPOCHS=3
    VAL_EVERY=20
    echo "=== DEBUG / SMOKE TEST MODE ==="
fi

cd "$ROOT"

# ── Step 0: generate policy label manifests ───────────────────────────────────
if [[ ! -f "$LABEL_DIR/train.json" ]] || [[ ! -f "$LABEL_DIR/val.json" ]]; then
    echo "Policy labels not found — generating..."
    python -m training.Policy.make_policy_labels \
        --manifest_dir "$ROOT/data/sft_manifests" \
        --out_dir      "$LABEL_DIR" \
        $DEBUG_FLAGS
fi

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

ok = True
for split in ["train", "val"]:
    p = Path("data/policy_labels") / f"{split}.json"
    if not p.exists():
        print(f"  MISSING: {split}.json"); ok = False; continue
    d = json.loads(p.read_text())
    counts = d.get("label_counts", {})
    excl   = d.get("excluded", {})
    print(f"  {split}: {d['total_samples']} samples  labels={counts}  excluded={excl}")
    if split == "train":
        for cls in ["SILENT", "OBSERVE", "ALERT"]:
            if counts.get(cls, 0) == 0:
                print(f"  ERROR: {cls} = 0 in train!"); ok = False
if not ok:
    sys.exit(1)
print("  ✅ Manifests OK")
PYEOF

# ── Step 2: build belief cache ────────────────────────────────────────────────
# This is a one-time step. Once done, re-running skip automatically.
if [[ ! -f "$CACHE_DIR/train.pt" ]] || [[ ! -f "$CACHE_DIR/val.pt" ]]; then
    echo ""
    echo "Building belief vector cache (one-time, uses frozen SFT)..."
    echo "  This encodes all windows through the VLM and saves belief vectors."
    echo "  After this step, training requires only the tiny PolicyHead (fast)."
    SPLITS_TO_CACHE=""
    [[ ! -f "$CACHE_DIR/train.pt" ]] && SPLITS_TO_CACHE="train"
    [[ ! -f "$CACHE_DIR/val.pt"   ]] && SPLITS_TO_CACHE="$SPLITS_TO_CACHE val"
    python -m training.Policy.make_belief_cache \
        --sft_checkpoint "$SFT_CHECKPOINT" \
        --label_dir      "$LABEL_DIR" \
        --out_dir        "$CACHE_DIR" \
        --batch_size     $CACHE_BATCH \
        --splits         $SPLITS_TO_CACHE
else
    echo "Belief cache found — skipping cache build."
fi

# ── Step 3: warm-start training ───────────────────────────────────────────────
echo ""
echo "Starting policy warm-start (cache-accelerated)..."
echo "  SFT checkpoint : $SFT_CHECKPOINT"
echo "  Belief cache   : $CACHE_DIR"
echo "  Output         : $OUTPUT_DIR/$EXPERIMENT"
echo "  train_batch    : $TRAIN_BATCH  (PolicyHead only, no VLM)"
echo "  lr             : $LR  epochs=$NUM_EPOCHS"

python -m training.Policy.warm_start_trainer \
    --sft_checkpoint    "$SFT_CHECKPOINT"  \
    --label_dir         "$LABEL_DIR"       \
    --belief_cache_dir  "$CACHE_DIR"       \
    --output_dir        "$OUTPUT_DIR"      \
    --experiment_name   "$EXPERIMENT"      \
    --num_epochs        $NUM_EPOCHS        \
    --batch_size        $TRAIN_BATCH       \
    --gradient_accumulation_steps $GRAD_ACCUM \
    --learning_rate     $LR                \
    --val_every_n_steps $VAL_EVERY         \
    --use_wandb                            \
    $DEBUG_FLAGS

# ── Step 4: evaluate best checkpoint ──────────────────────────────────────────
echo ""
echo "Evaluating best checkpoint..."
python -m training.Policy.evaluate_policy \
    --sft_checkpoint    "$SFT_CHECKPOINT"                    \
    --policy_checkpoint "$OUTPUT_DIR/$EXPERIMENT/best"       \
    --label_dir         "$LABEL_DIR"                         \
    --split             val                                  \
    --output_json       "$OUTPUT_DIR/$EXPERIMENT/eval_val.json"

echo ""
echo "✅ Stage 1 policy warm-start complete."
echo "   Best checkpoint : $OUTPUT_DIR/$EXPERIMENT/best"
echo "   Eval results    : $OUTPUT_DIR/$EXPERIMENT/eval_val.json"