#!/usr/bin/env bash
# ════════════════════════════════════════════════════════════════════════════
#  V3 baseline 上的 ALERT/OBSERVE 漏播修复 A/B 对照
#  使用现有 legacy mean_pool cache: data/belief_cache/{train,val}.pt
#  无需重建 cache, 直接在 v3 trainer 上做 loss-only 修复.
#
#  4 个实验:
#    v3_baseline           老配置, 无修复 (复现 32% 漏播)
#    v3_F1_only            +cost_lambda=0.3
#    v3_F2_only            +ordinal_lambda=0.2
#    v3_F1F2_full          +cost_lambda=0.3 +ordinal_lambda=0.2 (推荐)
#
#  时长: 单实验 ~12-15 min (cache 模式, 15 epoch, early stop ~3 epoch)
#        4 个串行 ~1h
#
#  Usage:
#    bash training/Policy/run_v3_alert_fix.sh                    # 全部 4 组
#    SKIP_BASELINE=1 bash training/Policy/run_v3_alert_fix.sh    # 跳过 baseline
#    ONLY_FULL=1 bash training/Policy/run_v3_alert_fix.sh        # 只跑 F1+F2
#    bash training/Policy/run_v3_alert_fix.sh --debug            # 128 样本快测
# ════════════════════════════════════════════════════════════════════════════
set -euo pipefail

cd "$(dirname "$0")/../.."
source ~/miniconda3/etc/profile.d/conda.sh 2>/dev/null || true
conda activate lkalert 2>/dev/null || true

SFT_CKPT="${SFT_CKPT:-checkpoints/SFT/sft_v2/best}"
LABEL_DIR="${LABEL_DIR:-data/policy_labels}"
CACHE_DIR="${CACHE_DIR:-data/belief_cache}"
OUTPUT_DIR="${OUTPUT_DIR:-checkpoints/Policy}"
NUM_EPOCHS="${NUM_EPOCHS:-15}"
BATCH_SIZE="${BATCH_SIZE:-256}"
LR="${LR:-3e-4}"
PATIENCE="${PATIENCE:-5}"
GPU="${GPU:-0}"
RUNS_DIR="${RUNS_DIR:-runs/v3_alert_fix}"

DEBUG_FLAG=""
for a in "$@"; do [[ "$a" == "--debug" ]] && DEBUG_FLAG="--debug"; done

export CUDA_VISIBLE_DEVICES="${GPU}"
mkdir -p "${RUNS_DIR}"

run_one () {
    local name="$1"; shift
    local extra="$*"
    local ts="$(date +%Y%m%d_%H%M%S)"
    local log="${RUNS_DIR}/${name}_${ts}.log"
    echo
    echo "═════════════════════════════════════════════════════════════════════"
    echo "  EXP: ${name}    extra: ${extra}"
    echo "  log → ${log}"
    echo "═════════════════════════════════════════════════════════════════════"
    python -m training.Policy.warm_start_trainer \
        --sft_checkpoint    "${SFT_CKPT}" \
        --label_dir         "${LABEL_DIR}" \
        --belief_cache_dir  "${CACHE_DIR}" \
        --output_dir        "${OUTPUT_DIR}" \
        --experiment_name   "${name}" \
        --num_epochs        "${NUM_EPOCHS}" \
        --batch_size        "${BATCH_SIZE}" \
        --learning_rate     "${LR}" \
        --focal_alpha       0.1 0.3 0.6 \
        --focal_gamma       2.0 \
        --belief_noise_std  0.01 \
        --label_smoothing   0.1 \
        --use_balanced_sampler \
        --early_stop_patience "${PATIENCE}" \
        --val_every_n_steps 200 \
        ${DEBUG_FLAG} \
        ${extra} 2>&1 | tee "${log}"
}

# ── 4 个实验 ────────────────────────────────────────────────────────────────
if [[ "${ONLY_FULL:-0}" == "1" ]]; then
    run_one "v3_F1F2_full"  --cost_lambda 0.3 --ordinal_lambda 0.2 --ordinal_margin 0.2
else
    [[ "${SKIP_BASELINE:-0}" != "1" ]] && \
        run_one "v3_baseline"
    run_one "v3_F1_only"     --cost_lambda 0.3 --ordinal_lambda 0
    run_one "v3_F2_only"     --cost_lambda 0   --ordinal_lambda 0.2 --ordinal_margin 0.2
    run_one "v3_F1F2_full"   --cost_lambda 0.3 --ordinal_lambda 0.2 --ordinal_margin 0.2
fi

# ── 汇总对比表 ──────────────────────────────────────────────────────────────
echo
echo "═════════════════════════════════════════════════════════════════════"
echo "  汇总: 4 组实验关键指标 (val best ckpt)"
echo "═════════════════════════════════════════════════════════════════════"
python3 - <<PY
import json, pathlib
rows = []
for name in ("v3_baseline", "v3_F1_only", "v3_F2_only", "v3_F1F2_full"):
    p = pathlib.Path("${OUTPUT_DIR}") / name / "best" / "policy_meta.json"
    if not p.exists():
        continue
    d = json.loads(p.read_text())
    rows.append((
        name,
        d.get("policy_score", 0),
        d.get("ego_alert_recall", 0),
        d.get("alert_leak_to_observe", float('nan')),
        d.get("alert_leak_to_silent", float('nan')),
        d.get("safe_neg_silent_rate", 0),
        d.get("overall_acc", 0),
    ))
print(f"{'experiment':<22}{'PScore':>9}{'ego_rec':>9}{'leak→O':>9}{'leak→S':>9}{'sn_silent':>11}{'acc':>8}")
print("─" * 78)
for r in rows:
    leak_o = f"{r[3]:.3f}" if r[3] == r[3] else "  n/a"
    leak_s = f"{r[4]:.3f}" if r[4] == r[4] else "  n/a"
    print(f"{r[0]:<22}{r[1]:>9.4f}{r[2]:>9.3f}{leak_o:>9}{leak_s:>9}{r[5]:>11.3f}{r[6]:>8.3f}")
print()
print("解读: leak→O 是核心指标. v3_baseline 应 ≈0.32, F1+F2 目标 ≤0.10.")
PY