#!/usr/bin/env python3
"""
Conformal Risk Control for LKAlert PolicyHead (v4 Evidential / v5 Hierarchical).

Provides distribution-free statistical guarantees:
  - "With probability ≥ 1-ε, the true class is in the prediction set."
  - Prediction set size adapts: uncertain inputs → larger set → conservative.

Two modes:
  1. Standard conformal:  coverage guarantee on class membership.
  2. Risk control:        guarantee on asymmetric miss cost (missing ALERT is worse).

Supports both model architectures:
  - v4 EvidentialPolicyModel: Dirichlet α → probs = α / S
  - v5 HierarchicalPolicyModel: (alert_logit, danger_logit) → 3-class probs

Usage:
  python -m training.Policy.conformal_risk \
      --sft_checkpoint   checkpoints/SFT/sft_v2/best \
      --v4_ckpt          checkpoints/Policy/policy_warmstart_v5_mono/best \
      --label_dir        data/policy_labels \
      --belief_cache_dir data/belief_cache \
      --output_dir       eval_results/paper_comparison_v5 \
      --epsilon          0.05
"""

from __future__ import annotations

import argparse
import json
import logging
from collections import defaultdict
from pathlib import Path
from typing import List, Optional

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from training.Policy.policy_model_v4 import EvidentialPolicyModel
from training.Policy.policy_model_v5 import HierarchicalPolicyModel
from training.Policy.policy_dataset import PolicyDataset, policy_collate_fn
from training.Policy.temporal_trainer import TemporalPolicyDataset, TemporalPolicyModel, temporal_collate_fn
from training.Policy.trajectory_trainer import TrajectoryPolicyDataset, TrajectoryPolicyModel, trajectory_collate_fn

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Policy.conformal")


def _detect_model_version(ckpt_dir: str) -> str:
    """Detect whether checkpoint is v4/v5/v6/v7."""
    meta_path = Path(ckpt_dir) / "policy_meta.json"
    if meta_path.exists():
        with open(meta_path) as f:
            meta = json.load(f)
        ver = meta.get("version", "")
        if ver == "v7_trajectory":
            return "v7"
        if ver == "v6_temporal":
            return "v6"
        if ver == "v5_hierarchical":
            return "v5"
    # fallback: check state_dict keys
    head_path = Path(ckpt_dir) / "policy_head.pt"
    if head_path.exists():
        sd = torch.load(head_path, map_location="cpu")
        if "danger_estimator.0.weight" in sd:
            return "v7"
        if "gru.weight_ih_l0" in sd:
            return "v6"
        if "alert_head.weight" in sd:
            return "v5"
    return "v4"


def _v5_logits_to_probs(alert_logit: np.ndarray, danger_logit: np.ndarray) -> np.ndarray:
    """
    Convert v5 hierarchical outputs to 3-class probabilities.

    P(ALERT)   = σ(alert_logit)
    P(DANGER)  = σ(danger_logit)   # = P(OBSERVE ∪ ALERT)
    P(SILENT)  = 1 - P(DANGER)
    P(OBSERVE) = P(DANGER) - P(ALERT)  (clipped ≥ 0)

    Returns: [N, 3] probabilities normalized to sum to 1.
    """
    p_alert = 1.0 / (1.0 + np.exp(-alert_logit))
    p_danger = 1.0 / (1.0 + np.exp(-danger_logit))

    p_silent = 1.0 - p_danger
    p_observe = np.clip(p_danger - p_alert, 0.0, None)

    probs = np.stack([p_silent, p_observe, p_alert], axis=-1)  # [N, 3]
    # renormalize (sigmoid outputs are independent, may not sum to 1)
    probs = probs / probs.sum(axis=-1, keepdims=True).clip(1e-8)
    return probs


def calibrate_conformal(
    alphas: np.ndarray,     # [N, K] Dirichlet concentrations
    labels: np.ndarray,     # [N] int
    epsilon: float = 0.05,  # target miscoverage rate
) -> dict:
    """
    Split-conformal calibration.

    Non-conformity score: s_i = 1 - p_{y_i}(x_i)
    where p = α / S (expected probability under Dirichlet).

    Returns threshold q_hat such that:
      P(y ∈ C(x)) ≥ 1 - ε  where  C(x) = {k : p_k(x) ≥ 1 - q_hat}
    """
    S = alphas.sum(axis=-1, keepdims=True)
    probs = alphas / S  # [N, K]

    idx = np.arange(len(labels))
    p_true = probs[idx, labels]
    scores = 1.0 - p_true  # [N]

    n = len(scores)
    q_level = np.ceil((n + 1) * (1 - epsilon)) / n
    q_level = min(q_level, 1.0)
    q_hat = float(np.quantile(scores, q_level))

    return {
        "q_hat": q_hat,
        "epsilon": epsilon,
        "n_calibration": n,
        "q_level": q_level,
        "score_mean": float(scores.mean()),
        "score_std": float(scores.std()),
        "score_p90": float(np.percentile(scores, 90)),
    }


def calibrate_risk_control(
    alphas: np.ndarray,
    labels: np.ndarray,
    epsilon: float = 0.05,
    cost_miss_alert: float = 5.0,
    cost_fa: float = 1.0,
) -> dict:
    """
    Conformal risk control with asymmetric costs (Angelopoulos et al. 2024).

    Risk function: L(C, y) = cost_miss * 1[y=ALERT, ALERT ∉ C] + cost_fa * |C|/K

    Find threshold λ such that E[L(C_λ, y)] ≤ δ, where
    C_λ(x) = {k : p_k(x) ≥ λ}.
    """
    S = alphas.sum(axis=-1, keepdims=True)
    probs = alphas / S
    K = alphas.shape[1]
    n = len(labels)

    lambdas = np.linspace(0.01, 0.99, 200)
    best_lambda = 0.01
    best_risk = float("inf")

    for lam in lambdas:
        sets = probs >= lam  # [N, K] bool
        risks = []
        for i in range(n):
            r = 0.0
            if labels[i] == 2 and not sets[i, 2]:
                r += cost_miss_alert
            r += cost_fa * sets[i].sum() / K
            risks.append(r)
        mean_risk = np.mean(risks)

        # Hoeffding correction for finite-sample guarantee
        correction = np.sqrt(np.log(1.0 / epsilon) / (2 * n))
        if mean_risk + correction <= best_risk:
            best_risk = mean_risk + correction
            best_lambda = lam

    # compute final metrics at best_lambda
    sets = probs >= best_lambda
    coverage = float(np.mean([labels[i] in np.where(sets[i])[0] for i in range(n)]))
    avg_set_size = float(sets.sum(axis=1).mean())
    alert_miss = float(np.mean([
        (not sets[i, 2]) for i in range(n) if labels[i] == 2
    ])) if (labels == 2).any() else 0.0

    return {
        "lambda": best_lambda,
        "epsilon": epsilon,
        "coverage": coverage,
        "avg_set_size": avg_set_size,
        "alert_miss_rate": alert_miss,
        "cost_miss_alert": cost_miss_alert,
        "cost_fa": cost_fa,
        "n_calibration": n,
    }


def predict_conformal(
    alphas: np.ndarray,
    q_hat: float,
) -> dict:
    """
    Apply conformal prediction to produce prediction sets.

    Returns:
      sets: [N, K] bool — prediction set membership
      sizes: [N] int — size of each prediction set
      preds: [N] int — point prediction (argmax within set, or conservative)
    """
    S = alphas.sum(axis=-1, keepdims=True)
    probs = alphas / S
    threshold = 1.0 - q_hat

    sets = probs >= threshold  # [N, K]
    sizes = sets.sum(axis=1)

    preds = probs.argmax(axis=1)
    # if prediction set is empty (very high q_hat), predict the argmax
    empty = sizes == 0
    if empty.any():
        sets[empty] = False
        sets[empty, preds[empty]] = True
        sizes[empty] = 1

    return {
        "sets": sets,
        "sizes": sizes,
        "preds": preds,
    }


def evaluate_conformal(
    alphas: np.ndarray,
    labels: np.ndarray,
    categories: np.ndarray,
    ttas: np.ndarray,
    video_ids: List[str],
    cal_result: dict,
    risk_result: dict,
) -> dict:
    """Full evaluation with conformal metrics."""
    N, K = alphas.shape
    S = alphas.sum(axis=-1, keepdims=True)
    probs = alphas / S
    u = K / S.squeeze(-1)

    # standard conformal
    conf_pred = predict_conformal(alphas, cal_result["q_hat"])
    sets = conf_pred["sets"]
    sizes = conf_pred["sizes"]

    coverage = float(np.mean([labels[i] in np.where(sets[i])[0] for i in range(N)]))
    avg_size = float(sizes.mean())

    # coverage by category
    ego_mask = categories == "ego_positive"
    ne_mask = categories == "non_ego"
    sn_mask = categories == "safe_neg"

    def _cov(mask):
        idx = np.where(mask)[0]
        if len(idx) == 0:
            return 0.0
        return float(np.mean([labels[i] in np.where(sets[i])[0] for i in idx]))

    # alert-specific safety metric: guaranteed miss rate
    alert_mask = labels == 2
    alert_in_set = np.array([sets[i, 2] for i in range(N)])
    guaranteed_miss_rate = 1.0 - float(alert_in_set[alert_mask].mean()) if alert_mask.any() else 0.0

    # conditional set sizes
    size_by_class = {}
    for k in range(K):
        mask = labels == k
        if mask.any():
            size_by_class[f"set_size_class_{k}"] = float(sizes[mask].mean())

    # uncertainty-coverage curve (for plotting)
    u_flat = u.flatten()
    thresholds = np.percentile(u_flat, np.arange(0, 101, 10))
    u_coverage_curve = []
    for thr in thresholds:
        mask = u_flat <= thr
        if mask.any():
            cov = float(np.mean([labels[i] in np.where(sets[i])[0] for i in np.where(mask)[0]]))
            u_coverage_curve.append({"u_threshold": float(thr), "coverage": cov, "frac": float(mask.mean())})

    # risk control metrics
    risk_lambda = risk_result["lambda"]
    risk_sets = probs >= risk_lambda
    risk_alert_miss = float(np.mean([
        (not risk_sets[i, 2]) for i in range(N) if labels[i] == 2
    ])) if alert_mask.any() else 0.0

    return {
        "conformal": {
            "q_hat": cal_result["q_hat"],
            "epsilon": cal_result["epsilon"],
            "empirical_coverage": coverage,
            "avg_set_size": avg_size,
            "coverage_ego": _cov(ego_mask),
            "coverage_non_ego": _cov(ne_mask),
            "coverage_safe_neg": _cov(sn_mask),
            "guaranteed_alert_miss_rate": guaranteed_miss_rate,
            **size_by_class,
        },
        "risk_control": {
            "lambda": risk_lambda,
            "alert_miss_rate": risk_alert_miss,
            "avg_set_size": float(risk_sets.sum(axis=1).mean()),
            "coverage": float(np.mean([labels[i] in np.where(risk_sets[i])[0] for i in range(N)])),
        },
        "uncertainty_stats": {
            "mean_u": float(u_flat.mean()),
            "u_alert": float(u_flat[alert_mask].mean()) if alert_mask.any() else 0.0,
            "u_silent": float(u_flat[labels == 0].mean()) if (labels == 0).any() else 0.0,
            "u_observe": float(u_flat[labels == 1].mean()) if (labels == 1).any() else 0.0,
        },
        "u_coverage_curve": u_coverage_curve,
        "n_samples": int(N),
    }


def main():
    parser = argparse.ArgumentParser("conformal_risk")
    parser.add_argument("--sft_checkpoint", required=True)
    parser.add_argument("--v4_ckpt", required=True, help="Policy checkpoint (v4 or v5)")
    parser.add_argument("--label_dir", default="data/policy_labels")
    parser.add_argument("--belief_cache_dir", default=None)
    parser.add_argument("--output_dir", default="eval_results/paper_comparison")
    parser.add_argument("--epsilon", type=float, default=0.05)
    parser.add_argument("--cost_miss_alert", type=float, default=5.0)
    parser.add_argument("--cost_fa", type=float, default=1.0)
    parser.add_argument("--batch_size", type=int, default=256)
    args = parser.parse_args()

    cache_dir = Path(args.belief_cache_dir) if args.belief_cache_dir else None

    def _cache_path(split):
        if cache_dir is None:
            return None
        p = cache_dir / f"{split}.pt"
        return p if p.exists() else None

    # auto-detect model version
    model_version = _detect_model_version(args.v4_ckpt)
    logger.info(f"Detected model version: {model_version}")

    # load val set — v6/v7 temporal models need sequence datasets
    if model_version in ("v6", "v7"):
        meta_path = Path(args.v4_ckpt) / "policy_meta.json"
        seq_len = 8
        if meta_path.exists():
            with open(meta_path) as f:
                seq_len = json.load(f).get("seq_len", 8)
        if model_version == "v7":
            val_ds = TrajectoryPolicyDataset(
                manifests=[Path(args.label_dir) / "val.json"],
                split="val",
                belief_cache_path=_cache_path("val"),
                seq_len=seq_len,
            )
            val_loader = DataLoader(
                val_ds, batch_size=args.batch_size, shuffle=False,
                num_workers=4, collate_fn=trajectory_collate_fn,
            )
        else:
            val_ds = TemporalPolicyDataset(
                manifests=[Path(args.label_dir) / "val.json"],
                split="val",
                belief_cache_path=_cache_path("val"),
                seq_len=seq_len,
            )
            val_loader = DataLoader(
                val_ds, batch_size=args.batch_size, shuffle=False,
                num_workers=4, collate_fn=temporal_collate_fn,
            )
    else:
        val_ds = PolicyDataset(
            manifests=[Path(args.label_dir) / "val.json"],
            split="val",
            belief_cache_path=_cache_path("val"),
        )
        val_loader = DataLoader(
            val_ds, batch_size=args.batch_size, shuffle=False,
            num_workers=4, collate_fn=policy_collate_fn,
        )

    # infer hidden_dim from the val cache (backbone-agnostic)
    ds_cache = getattr(val_ds, "_cache", None)
    if ds_cache is not None and "beliefs" in ds_cache:
        cache_hidden_dim = int(ds_cache["beliefs"].shape[-1])
    else:
        cache_hidden_dim = 2048  # legacy fallback

    # load model
    if model_version == "v7":
        # detect use_gru from meta or state_dict
        use_gru = True
        if meta_path.exists():
            with open(meta_path) as f:
                use_gru = json.load(f).get("use_gru", True)
        model = TrajectoryPolicyModel(
            hidden_dim=cache_hidden_dim, seq_len=seq_len, use_gru=use_gru
        )
        model.load_policy_checkpoint(args.v4_ckpt)
    elif model_version == "v6":
        model = TemporalPolicyModel(hidden_dim=cache_hidden_dim, seq_len=seq_len)
        model.load_policy_checkpoint(args.v4_ckpt)
    elif model_version == "v5":
        model = HierarchicalPolicyModel(
            sft_checkpoint_dir=args.sft_checkpoint,
            use_bf16=True,
        )
        model.load_policy_checkpoint(args.v4_ckpt)
    else:
        model = EvidentialPolicyModel(
            sft_checkpoint_dir=args.sft_checkpoint,
            use_bf16=True,
        )
        model.load_policy_checkpoint(args.v4_ckpt)
    model.eval()

    # extract outputs
    all_alphas = []   # [N, 3] — Dirichlet α (v4), pseudo-probs (v5), or softmax probs (v6)
    all_labels = []
    all_cats = []
    all_ttas = []
    all_vids = []

    logger.info(f"Extracting predictions from val set ({model_version})...")
    with torch.no_grad():
        for batch in tqdm(val_loader, desc=f"Extract ({model_version})", ncols=80):
            if model_version == "v7":
                logits, _danger_t = model(
                    batch["belief_seqs"], batch["tta_mean_seqs"], batch["tta_var_seqs"]
                )
                probs = torch.softmax(logits, dim=-1).cpu().numpy()
                all_alphas.append(probs)
            elif model_version == "v6":
                logits = model(batch["belief_seqs"], batch["tta_mean_seqs"], batch["tta_var_seqs"])
                probs = torch.softmax(logits, dim=-1).cpu().numpy()
                all_alphas.append(probs)
            elif model_version == "v5":
                if "beliefs" in batch:
                    out = model.forward_cached(batch["beliefs"], batch["tta_means"], batch["tta_vars"])
                else:
                    out = model(batch["images"], batch["metadata"])
                alert_logit = out[0].cpu().numpy()
                danger_logit = out[1].cpu().numpy()
                all_alphas.append(_v5_logits_to_probs(alert_logit, danger_logit))
            else:
                if "beliefs" in batch:
                    out = model.forward_cached(batch["beliefs"], batch["tta_means"], batch["tta_vars"])
                else:
                    out = model(batch["images"], batch["metadata"])
                all_alphas.append(out.cpu().numpy())

            all_labels.extend(batch["action_labels"].tolist())
            all_cats.extend(batch["categories"])
            all_ttas.extend(batch["tta_raws"].tolist())
            all_vids.extend(batch["video_ids"])

    alphas = np.concatenate(all_alphas, axis=0)
    labels = np.array(all_labels)
    cats = np.array(all_cats)
    ttas = np.array(all_ttas)

    # split val into calibration (50%) and test (50%)
    n = len(labels)
    np.random.seed(42)
    perm = np.random.permutation(n)
    n_cal = n // 2
    cal_idx, test_idx = perm[:n_cal], perm[n_cal:]

    logger.info(f"Calibration: {n_cal} samples, Test: {n - n_cal} samples")

    # calibrate on first half
    cal_result = calibrate_conformal(alphas[cal_idx], labels[cal_idx], args.epsilon)
    logger.info(f"Conformal q_hat = {cal_result['q_hat']:.4f} (epsilon={args.epsilon})")

    risk_result = calibrate_risk_control(
        alphas[cal_idx], labels[cal_idx],
        epsilon=args.epsilon,
        cost_miss_alert=args.cost_miss_alert,
        cost_fa=args.cost_fa,
    )
    logger.info(f"Risk control lambda = {risk_result['lambda']:.4f}")

    # evaluate on second half
    test_vids = [all_vids[i] for i in test_idx]
    eval_result = evaluate_conformal(
        alphas[test_idx], labels[test_idx], cats[test_idx], ttas[test_idx],
        test_vids, cal_result, risk_result,
    )

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    with open(out_dir / "conformal_results.json", "w") as f:
        json.dump({
            "model_version": model_version,
            "calibration": cal_result,
            "risk_control": risk_result,
            "evaluation": eval_result,
        }, f, indent=2, default=str)

    logger.info(f"\nResults saved to {out_dir / 'conformal_results.json'}")
    logger.info(f"  Model version: {model_version}")
    logger.info(f"  Coverage: {eval_result['conformal']['empirical_coverage']:.4f}")
    logger.info(f"  Avg set size: {eval_result['conformal']['avg_set_size']:.2f}")
    logger.info(f"  Guaranteed alert miss: {eval_result['conformal']['guaranteed_alert_miss_rate']:.4f}")
    logger.info(f"  Risk control alert miss: {eval_result['risk_control']['alert_miss_rate']:.4f}")
    if model_version in ("v5", "v6", "v7"):
        logger.info("  Note: uncertainty stats (K/S) are not meaningful for %s "
                     "(no Dirichlet output)" % model_version)


if __name__ == "__main__":
    main()