"""
zero_shot_calibrator.py — Zero-label conformal calibration for llm-guard-kit

Problem
-------
QuickCalibrator needs 20 real human labels per new domain.  Most deployments
cannot provide labels upfront.

Solution
--------
Use P(True) from Haiku as a pseudo-labeler:
  P(True) risk > ptrue_threshold  →  pseudo-label = 1 (wrong chain)
  P(True) risk ≤ ptrue_threshold  →  pseudo-label = 0 (correct chain)

Then fit isotonic calibration on those pseudo-labeled scores and derive a
conformal threshold at a target FPR.

Expected degradation vs real labels: ~5pp alert precision on HP test sets.

Cost: ~$0.0003 × n_pseudo calls to Haiku (total ~$0.015 for n_pseudo=50).
      $0 if the deployment already calls score_with_ptrue() in production.

Usage
-----
    from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
    from llm_guard import AgentGuard

    guard = AgentGuard(api_key="sk-ant-...")
    cal   = ZeroShotCalibrator(guard)
    cal.fit(unlabeled_chains, n_pseudo=50)

    threshold = cal.get_threshold(target_fpr=0.10)
    # Score a new chain and compare to threshold
    result = guard.score_with_ptrue(question, steps, final_answer)
    alert  = result.risk_score >= threshold

Validated
---------
    HP pseudo-label noise rate:  est. ~15-20% (P(True) precision ~0.80-0.85)
    Expected alert precision:    ~0.75-0.85 at FPR ≤ 10% (vs 0.896 real-label)
Purpose:    Uses Haiku P(True) as an automatic pseudo-labeler to fit isotonic calibration without any human labels. Expected degradation vs real labels is ~5pp alert precision.
Exports:    ZeroShotCalibrator
Added:      v0.29.0
Status:     Stable
"""

from __future__ import annotations

import warnings
from typing import Dict, List, Optional


class ZeroShotCalibrator:
    """
    Calibrate a conformal alert threshold on a new domain with ZERO human labels.

    Uses P(True) to generate pseudo-labels on a small unlabeled calibration
    pool, then fits an isotonic recalibrator and derives a FPR-constrained
    threshold.

    Parameters
    ----------
    guard : AgentGuard
        An existing AgentGuard instance with an Anthropic or NIM API key
        (required for P(True) calls).
    ptrue_threshold : float
        P(True) risk above this value → pseudo-label as wrong (1).
        Default 0.55.  Lower values are more aggressive (more pseudo-wrongs).
        Empirically validated: 0.55 gives ~80% pseudo-label precision on HP.

    Notes
    -----
    - The isotonic regression is fitted on (behavioral_score, pseudo_label)
      pairs, not on ptrue_risk directly.  This reuses the same isotonic path
      as QuickCalibrator, keeping the inference path identical at deploy time.
    - get_threshold() uses the conformal empirical quantile on pseudo-labeled
      correct chains only (one-sided coverage on non-alerts).
    - calibration_quality_estimate() measures bimodality of ptrue_risk
      distribution — high bimodality = confident pseudo-labels = less noise.
    """

    def __init__(
        self,
        guard,
        ptrue_threshold: float = 0.55,
    ):
        self._guard            = guard
        self._ptrue_threshold  = ptrue_threshold
        self._is_fitted: bool  = False
        self._pseudo_labels:  List[int]   = []
        self._beh_scores:     List[float] = []
        self._ptrue_risks:    List[float] = []
        self._n_pseudo: int               = 0
        self._threshold_cache: Dict[float, float] = {}

    # ── Public API ─────────────────────────────────────────────────────────────

    def fit(
        self,
        chains: List[Dict],
        n_pseudo: int = 50,
    ) -> "ZeroShotCalibrator":
        """
        Score the first n_pseudo chains with P(True), generate pseudo-labels,
        then fit an isotonic recalibrator on the pseudo-labeled behavioral scores.

        Parameters
        ----------
        chains : list of dict
            Each dict must contain: "question", "steps", "final_answer".
            NO labels required — this is the zero-shot path.
        n_pseudo : int
            Number of chains to pseudo-label. Default 50.  Max budget ~$0.015.
            Fewer chains → noisier calibration; more chains → better precision.

        Returns
        -------
        self  (for method chaining)

        Raises
        ------
        ValueError
            If chains is empty or n_pseudo < 5.
        RuntimeError
            If fewer than 5 chains score successfully (API error or malformed input).
        """
        if not chains:
            raise ValueError("ZeroShotCalibrator.fit(): chains list is empty.")
        if n_pseudo < 5:
            raise ValueError(
                f"ZeroShotCalibrator.fit(): n_pseudo={n_pseudo} is too small. "
                "Minimum 5 pseudo-labeled chains are required."
            )

        pool = chains[:n_pseudo]
        self._n_pseudo = len(pool)

        beh_scores:   List[float] = []
        ptrue_risks:  List[float] = []
        pseudo_labels: List[int]  = []
        n_failed = 0

        for i, chain in enumerate(pool):
            try:
                result = self._guard.score_with_ptrue(
                    question=chain["question"],
                    steps=chain["steps"],
                    final_answer=chain["final_answer"],
                )
                beh_score  = float(result.behavioral_score)
                ptrue_risk = float(
                    result.behavioral_components.get("ptrue_risk", beh_score)
                )
                pseudo_label = 1 if ptrue_risk > self._ptrue_threshold else 0

                beh_scores.append(beh_score)
                ptrue_risks.append(ptrue_risk)
                pseudo_labels.append(pseudo_label)

            except Exception as exc:
                n_failed += 1
                warnings.warn(
                    f"ZeroShotCalibrator.fit(): chain {i} failed (skipped): {exc}",
                    RuntimeWarning,
                    stacklevel=2,
                )

        n_scored = len(beh_scores)
        if n_scored < 5:
            raise RuntimeError(
                f"ZeroShotCalibrator.fit(): only {n_scored} of {len(pool)} chains "
                "scored successfully.  Check API key and chain format."
            )

        if n_failed > 0:
            warnings.warn(
                f"ZeroShotCalibrator.fit(): {n_failed} chains failed; "
                f"calibration fitted on {n_scored}/{len(pool)} chains.",
                RuntimeWarning,
                stacklevel=2,
            )

        self._beh_scores   = beh_scores
        self._ptrue_risks  = ptrue_risks
        self._pseudo_labels = pseudo_labels

        # Fit isotonic regression on behavioral scores vs pseudo-labels
        # (same path as QuickCalibrator / calibrate_isotonic)
        self._guard.calibrate_isotonic(beh_scores, pseudo_labels)
        self._is_fitted = True
        self._threshold_cache.clear()

        n_pseudo_wrong   = sum(pseudo_labels)
        n_pseudo_correct = n_scored - n_pseudo_wrong
        print(
            f"[ZeroShotCalibrator] Fitted on {n_scored} pseudo-labeled chains "
            f"({n_pseudo_correct} pseudo-correct, {n_pseudo_wrong} pseudo-wrong) "
            f"using ptrue_threshold={self._ptrue_threshold}"
        )
        return self

    def get_threshold(self, target_fpr: float = 0.10) -> float:
        """
        Return conformal alert threshold at the given target FPR.

        The threshold is the minimum behavioral score t such that the fraction
        of pseudo-correct chains with score >= t is <= target_fpr.

        In other words: among the pseudo-labeled correct chains, at most
        target_fpr fraction would be false-alerted at this threshold.

        Parameters
        ----------
        target_fpr : float
            Target false positive rate. Default 0.10. Range (0, 1).

        Returns
        -------
        float
            Alert threshold in [0, 1].  Score a new chain; if its calibrated
            behavioral score >= threshold, alert.

        Raises
        ------
        RuntimeError
            If fit() has not been called yet.
        ValueError
            If no pseudo-correct chains exist in the calibration set.
        """
        self._check_fitted()

        if target_fpr in self._threshold_cache:
            return self._threshold_cache[target_fpr]

        import numpy as np

        beh = np.array(self._beh_scores)
        labels = np.array(self._pseudo_labels)

        # Correct chain scores (pseudo-label = 0)
        correct_scores = beh[labels == 0]
        if len(correct_scores) == 0:
            raise ValueError(
                "ZeroShotCalibrator.get_threshold(): no pseudo-correct chains in "
                "calibration set.  Lower ptrue_threshold or use more chains."
            )

        # Conformal quantile: 1 - target_fpr quantile of correct scores
        # Threshold = percentile at (1 - target_fpr) of correct chain scores
        # → at most target_fpr fraction of correct chains score above threshold
        threshold = float(np.quantile(correct_scores, 1.0 - target_fpr))
        threshold = float(np.clip(threshold, 0.0, 1.0))

        self._threshold_cache[target_fpr] = threshold
        return threshold

    def score(
        self,
        question: str,
        steps: list,
        final_answer: str,
    ) -> float:
        """
        Score a single chain and return the calibrated behavioral risk.

        Parameters
        ----------
        question, steps, final_answer : chain components

        Returns
        -------
        float : calibrated behavioral risk in [0, 1].

        Raises
        ------
        RuntimeError
            If fit() has not been called.
        """
        self._check_fitted()
        result = self._guard.score_chain(
            question=question,
            steps=steps,
            final_answer=final_answer,
        )
        raw = float(result.behavioral_score)
        return self._apply_iso(raw)

    def calibration_quality_estimate(self) -> dict:
        """
        Estimate pseudo-label noise rate from P(True) confidence distribution.

        A high-quality pseudo-labeling produces a bimodal ptrue_risk distribution
        (most chains clearly near 0.0 or 1.0).  Unimodal/flat distributions
        indicate noisy labels and degraded calibration.

        Returns
        -------
        dict with keys:
          n_pseudo            : number of pseudo-labeled chains
          n_pseudo_correct    : chains labeled as correct (ptrue_risk ≤ threshold)
          n_pseudo_wrong      : chains labeled as wrong (ptrue_risk > threshold)
          wrong_rate          : fraction pseudo-labeled as wrong
          ptrue_mean          : mean ptrue_risk across all chains
          ptrue_std           : std of ptrue_risk across all chains
          high_confidence_frac: fraction of chains with ptrue_risk < 0.25 or > 0.75
          estimated_noise_rate: heuristic noise estimate = 1 - high_confidence_frac
          quality             : "high" / "medium" / "low"

        Raises
        ------
        RuntimeError
            If fit() has not been called.
        """
        self._check_fitted()
        import numpy as np

        pr = np.array(self._ptrue_risks)
        labels = np.array(self._pseudo_labels)

        n = len(pr)
        n_wrong   = int(labels.sum())
        n_correct = n - n_wrong

        ptrue_mean = float(pr.mean())
        ptrue_std  = float(pr.std())

        # High-confidence: clearly correct (< 0.25) or clearly wrong (> 0.75)
        high_conf = float(np.mean((pr < 0.25) | (pr > 0.75)))
        noise_est = round(1.0 - high_conf, 3)

        if high_conf >= 0.70:
            quality = "high"
        elif high_conf >= 0.45:
            quality = "medium"
        else:
            quality = "low"

        return {
            "n_pseudo":            n,
            "n_pseudo_correct":    n_correct,
            "n_pseudo_wrong":      n_wrong,
            "wrong_rate":          round(n_wrong / max(1, n), 3),
            "ptrue_mean":          round(ptrue_mean, 3),
            "ptrue_std":           round(ptrue_std, 3),
            "high_confidence_frac": round(high_conf, 3),
            "estimated_noise_rate": noise_est,
            "quality":             quality,
        }

    # ── Properties ─────────────────────────────────────────────────────────────

    @property
    def is_fitted(self) -> bool:
        """True if fit() has been called successfully."""
        return self._is_fitted

    @property
    def ptrue_threshold(self) -> float:
        """P(True) risk threshold used for pseudo-labeling."""
        return self._ptrue_threshold

    @property
    def n_pseudo(self) -> int:
        """Number of chains that were successfully pseudo-labeled."""
        return len(self._beh_scores)

    # ── Internal helpers ───────────────────────────────────────────────────────

    def _check_fitted(self) -> None:
        if not self._is_fitted:
            raise RuntimeError(
                "ZeroShotCalibrator has not been fitted yet. "
                "Call fit(unlabeled_chains) first."
            )

    def _apply_iso(self, raw_score: float) -> float:
        iso = getattr(self._guard, "_iso_calibrator", None)
        if iso is None:
            return raw_score
        try:
            import numpy as np
            return float(iso.predict(np.array([raw_score]))[0])
        except Exception:
            return raw_score

    def __repr__(self) -> str:
        return (
            f"ZeroShotCalibrator("
            f"ptrue_threshold={self._ptrue_threshold}, "
            f"n_pseudo={self.n_pseudo}, "
            f"fitted={self._is_fitted}"
            f")"
        )
