#!/bin/bash
# ══════════════════════════════════════════════════════════════════════════════
# LKAlert Overnight Experiment Suite
#
# Part 1: Small improvements (~15 min)
#   1a. Conformal with cost_miss_alert=50          ~2 min
#   1b. verify_binary_ap (v3 checkpoint)           ~2 min
#   1c. Threshold + TTA + Ensemble analysis        ~5 min
#
# Part 2: Temporal Belief Aggregation (~2-3h)
#   2a. temporal_base:      seq=8, balanced, no mono       ~30 min
#   2b. temporal_mono:      seq=8, balanced, mono_λ=0.1    ~30 min
#   2c. temporal_long:      seq=16, balanced, no mono      ~30 min
#   2d. temporal_long_mono: seq=16, balanced, mono_λ=0.1   ~30 min
#
# Part 3: Post-analysis on best temporal model (~10 min)
#   3a. Conformal calibration
#   3b. Threshold analysis
#
# Total: ~3-4 hours
#
# Usage:
#   bash training/Policy/run_overnight.sh 2>&1 | tee logs/overnight_$(date +%Y%m%d_%H%M).log
# ══════════════════════════════════════════════════════════════════════════════

set -euo pipefail

cd "$(dirname "$0")/../.."

# Ensure log directory exists
mkdir -p logs

SFT_CKPT="checkpoints/SFT/sft_v2/best"
LABEL_DIR="data/policy_labels"
CACHE_DIR="data/belief_cache"
V3_CKPT="checkpoints/Policy/policy_warmstart_v3/best"
V5_CKPT="checkpoints/Policy/policy_warmstart_v5_mono/best"
OUTPUT_DIR="checkpoints/Policy"

START_TIME=$(date +%s)

echo "╔══════════════════════════════════════════════════════════╗"
echo "║  LKAlert Overnight Experiment Suite                     ║"
echo "║  Started: $(date '+%Y-%m-%d %H:%M:%S')                        ║"
echo "║  Expected: ~3-4 hours                                   ║"
echo "╚══════════════════════════════════════════════════════════╝"


# ══════════════════════════════════════════════════════════════════════════════
# Part 1: Small improvements
# ══════════════════════════════════════════════════════════════════════════════

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "  PART 1: Small Improvements (~15 min)"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# ── 1a. Conformal with higher cost ──
echo ""
echo "── [1a] Conformal risk (cost_miss=50) on v5_mono ──"
python -m training.Policy.conformal_risk \
    --sft_checkpoint   "$SFT_CKPT" \
    --v4_ckpt          "$V5_CKPT" \
    --label_dir        "$LABEL_DIR" \
    --belief_cache_dir "$CACHE_DIR" \
    --output_dir       eval_results/conformal_v5_cost50 \
    --cost_miss_alert  50.0 \
    --epsilon          0.05

# ── 1b. verify_binary_ap on v3 ──
echo ""
echo "── [1b] Binary AP verification (v3) ──"
python -m training.Policy.verify_binary_ap \
    --sft_checkpoint   "$SFT_CKPT" \
    --policy_checkpoint "$V3_CKPT" \
    --label_dir        "$LABEL_DIR" \
    --belief_cache_dir "$CACHE_DIR" \
    --output_dir       eval_results/binary_ap_verification

# ── 1c. Threshold + Ensemble analysis ──
echo ""
echo "── [1c] Threshold / TTA / Ensemble analysis ──"
python -m training.Policy.threshold_analysis \
    --label_dir        "$LABEL_DIR" \
    --belief_cache_dir "$CACHE_DIR" \
    --v3_ckpt          "$V3_CKPT" \
    --v5_ckpt          "$V5_CKPT" \
    --output_dir       eval_results/threshold_analysis

PART1_TIME=$(date +%s)
echo ""
echo "  Part 1 done in $(( (PART1_TIME - START_TIME) / 60 )) min"


# ══════════════════════════════════════════════════════════════════════════════
# Part 2: Temporal Belief Aggregation
# ══════════════════════════════════════════════════════════════════════════════

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "  PART 2: Temporal Belief Aggregation (~2-3h)"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# ── 2a. temporal_base: seq=8, no mono ──
echo ""
echo "── [2a/4] temporal_base: seq=8, balanced, no mono ──"
python -m training.Policy.temporal_trainer \
    --sft_checkpoint    "$SFT_CKPT" \
    --label_dir         "$LABEL_DIR" \
    --belief_cache_dir  "$CACHE_DIR" \
    --output_dir        "$OUTPUT_DIR" \
    --experiment_name   temporal_base \
    --seq_len           8 \
    --num_epochs        15 \
    --batch_size        256 \
    --learning_rate     2e-4 \
    --focal_alpha       0.75 \
    --focal_gamma       2.0 \
    --mono_lambda       0.0 \
    --val_every_n_steps 200 \
    --early_stop_patience 7 \
    --use_balanced_sampler

# ── 2b. temporal_mono: seq=8, mono ──
echo ""
echo "── [2b/4] temporal_mono: seq=8, balanced, mono_λ=0.1 ──"
python -m training.Policy.temporal_trainer \
    --sft_checkpoint    "$SFT_CKPT" \
    --label_dir         "$LABEL_DIR" \
    --belief_cache_dir  "$CACHE_DIR" \
    --output_dir        "$OUTPUT_DIR" \
    --experiment_name   temporal_mono \
    --seq_len           8 \
    --num_epochs        15 \
    --batch_size        256 \
    --learning_rate     2e-4 \
    --focal_alpha       0.75 \
    --focal_gamma       2.0 \
    --mono_lambda       0.1 \
    --val_every_n_steps 200 \
    --early_stop_patience 7 \
    --use_balanced_sampler

# ── 2c. temporal_long: seq=16 ──
echo ""
echo "── [2c/4] temporal_long: seq=16, balanced, no mono ──"
python -m training.Policy.temporal_trainer \
    --sft_checkpoint    "$SFT_CKPT" \
    --label_dir         "$LABEL_DIR" \
    --belief_cache_dir  "$CACHE_DIR" \
    --output_dir        "$OUTPUT_DIR" \
    --experiment_name   temporal_long \
    --seq_len           16 \
    --num_epochs        15 \
    --batch_size        128 \
    --learning_rate     2e-4 \
    --focal_alpha       0.75 \
    --focal_gamma       2.0 \
    --mono_lambda       0.0 \
    --val_every_n_steps 200 \
    --early_stop_patience 7 \
    --use_balanced_sampler

# ── 2d. temporal_long_mono: seq=16 + mono ──
echo ""
echo "── [2d/4] temporal_long_mono: seq=16, balanced, mono_λ=0.1 ──"
python -m training.Policy.temporal_trainer \
    --sft_checkpoint    "$SFT_CKPT" \
    --label_dir         "$LABEL_DIR" \
    --belief_cache_dir  "$CACHE_DIR" \
    --output_dir        "$OUTPUT_DIR" \
    --experiment_name   temporal_long_mono \
    --seq_len           16 \
    --num_epochs        15 \
    --batch_size        128 \
    --learning_rate     2e-4 \
    --focal_alpha       0.75 \
    --focal_gamma       2.0 \
    --mono_lambda       0.1 \
    --val_every_n_steps 200 \
    --early_stop_patience 7 \
    --use_balanced_sampler

PART2_TIME=$(date +%s)
echo ""
echo "  Part 2 done in $(( (PART2_TIME - PART1_TIME) / 60 )) min"


# ══════════════════════════════════════════════════════════════════════════════
# Part 3: Post-analysis on all temporal models
# ══════════════════════════════════════════════════════════════════════════════

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "  PART 3: Post-analysis (~10 min)"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# Find best temporal model by reading policy_meta.json
echo ""
echo "── Comparing temporal models ──"
python3 -c "
import json, sys
from pathlib import Path

models = ['temporal_base', 'temporal_mono', 'temporal_long', 'temporal_long_mono']
best_name, best_score = None, -1

for name in models:
    meta_path = Path('checkpoints/Policy') / name / 'best' / 'policy_meta.json'
    if meta_path.exists():
        with open(meta_path) as f:
            meta = json.load(f)
        score = meta.get('grid_best_policy_score', meta.get('policy_score', 0))
        ap = meta.get('binary_ap', 0)
        print(f'  {name:25s}  PolicyScore={score:.4f}  AP={ap:.4f}')
        if score > best_score:
            best_score = score
            best_name = name
    else:
        print(f'  {name:25s}  (no checkpoint found)')

if best_name:
    print(f'\n  >>> Best: {best_name} (PolicyScore={best_score:.4f})')
    # Write best name for downstream scripts
    Path('checkpoints/Policy/.best_temporal').write_text(best_name)
else:
    print('  No temporal models found!')
    sys.exit(1)
"

BEST_TEMPORAL=$(cat checkpoints/Policy/.best_temporal 2>/dev/null || echo "temporal_base")
BEST_CKPT="${OUTPUT_DIR}/${BEST_TEMPORAL}/best"

echo ""
echo "── [3a] Conformal on best temporal (${BEST_TEMPORAL}) ──"
python -m training.Policy.conformal_risk \
    --sft_checkpoint   "$SFT_CKPT" \
    --v4_ckpt          "$BEST_CKPT" \
    --label_dir        "$LABEL_DIR" \
    --belief_cache_dir "$CACHE_DIR" \
    --output_dir       "eval_results/temporal_conformal" \
    --cost_miss_alert  50.0 \
    --epsilon          0.05 \
    || echo "  (conformal skipped — model version detection may need update for v6)"

PART3_TIME=$(date +%s)

# ══════════════════════════════════════════════════════════════════════════════
# Summary
# ══════════════════════════════════════════════════════════════════════════════

echo ""
echo "╔══════════════════════════════════════════════════════════╗"
echo "║  ALL EXPERIMENTS COMPLETE                               ║"
echo "║  Finished: $(date '+%Y-%m-%d %H:%M:%S')                       ║"
echo "║  Total time: $(( (PART3_TIME - START_TIME) / 60 )) min                                  ║"
echo "╠══════════════════════════════════════════════════════════╣"
echo "║  Results:                                               ║"
echo "║    eval_results/conformal_v5_cost50/                    ║"
echo "║    eval_results/binary_ap_verification/                 ║"
echo "║    eval_results/threshold_analysis/                     ║"
echo "║    eval_results/temporal_conformal/                     ║"
echo "║                                                         ║"
echo "║  Temporal checkpoints:                                  ║"
echo "║    ${OUTPUT_DIR}/temporal_base/best"
echo "║    ${OUTPUT_DIR}/temporal_mono/best"
echo "║    ${OUTPUT_DIR}/temporal_long/best"
echo "║    ${OUTPUT_DIR}/temporal_long_mono/best"
echo "║                                                         ║"
echo "║  Best temporal: ${BEST_TEMPORAL}"
echo "╚══════════════════════════════════════════════════════════╝"