"""
domain_adaptive_calibrator.py — Warm-start conformal calibration.

Reduces calibration friction from 50+ hard labels to 10 hard labels by
combining P(True) pseudo-labels (ZeroShotCalibrator-style) with a small
hard-label set. Trust weight λ is derived from the pseudo-label noise estimate.

Usage
-----
    from llm_guard.calibration.domain_adaptive_calibrator import DomainAdaptiveCalibrator

    # Collect pseudo-labels via ZeroShotCalibrator (0 human labels)
    from llm_guard.calibration.zero_shot_calibrator import ZeroShotCalibrator
    zsc = ZeroShotCalibrator(guard)
    zsc.fit(unlabelled_chains)
    pseudo_scores, pseudo_preds = zsc.get_pseudo_scores()

    # Collect 10 hard labels from humans
    hard_scores = [result.risk_score for result in scored_chains[:10]]
    hard_labels = [1 if chain_was_wrong else 0 for chain in chains[:10]]

    # Fit the adaptive calibrator
    cal = DomainAdaptiveCalibrator(pseudo_noise_estimate=zsc.noise_estimate())
    threshold = cal.fit(hard_scores, hard_labels, pseudo_scores, pseudo_preds)

    guard2 = AgentGuard(alert_threshold=threshold)
Purpose:    Reduces calibration data requirement from 50+ labels to ~10 hard labels by mixing P(True) pseudo-labels with human annotations via noise-weighted trust. Speeds up cross-domain deployment.
Exports:    DomainAdaptiveCalibrator
Added:      v0.29.0
Status:     Stable
"""
from __future__ import annotations

import logging
from typing import List, Optional

import numpy as np

logger = logging.getLogger(__name__)

_DEFAULT_ALPHA = 0.10
_DEFAULT_LAMBDA_HARD = 5.0


class DomainAdaptiveCalibrator:
    """
    Warm-start conformal threshold calibrator.

    Combines pseudo-labels (from P(True) proxy, higher noise) with a small
    set of hard human labels to compute a conformal alert threshold at a
    target FPR.

    Parameters
    ----------
    pseudo_noise_estimate : float
        Estimated fraction of pseudo-labels that are flipped (0–1).
        Higher noise → pseudo-labels weighted less relative to hard labels.
        Default 0.30 (conservative; ZeroShotCalibrator typically 0.25–0.40).
    alpha : float
        Target FPR for conformal threshold. Default 0.10 (10%).
    n_pseudo_min : int
        Minimum number of pseudo-labels required to include them.
        If fewer are provided, falls back to hard-label-only calibration.
        Default 20.
    """

    def __init__(
        self,
        pseudo_noise_estimate: float = 0.30,
        alpha: float = _DEFAULT_ALPHA,
        n_pseudo_min: int = 20,
    ) -> None:
        if not 0.0 <= pseudo_noise_estimate < 1.0:
            raise ValueError(
                f"pseudo_noise_estimate must be in [0, 1); got {pseudo_noise_estimate}"
            )
        self._noise = pseudo_noise_estimate
        self._alpha = alpha
        self._n_pseudo_min = n_pseudo_min
        # Weight for hard labels relative to pseudo: 1/(noise+ε), capped at 10×
        self._lambda_hard = min(1.0 / max(pseudo_noise_estimate, 0.05), 10.0)
        self._threshold: Optional[float] = None

    def fit(
        self,
        hard_scores: List[float],
        hard_labels: List[int],
        pseudo_scores: Optional[List[float]] = None,
        pseudo_labels: Optional[List[int]] = None,
    ) -> float:
        """
        Compute conformal alert threshold from hard + (optional) pseudo-label pools.

        Parameters
        ----------
        hard_scores : list of float
            Risk scores for hard-labelled chains.
        hard_labels : list of int
            Human labels: 1 = wrong/alert, 0 = correct.
        pseudo_scores : list of float, optional
            Risk scores for pseudo-labelled (unlabelled) chains.
        pseudo_labels : list of int, optional
            Pseudo-labels from ZeroShotCalibrator or P(True) thresholding.

        Returns
        -------
        float : threshold t such that risk ≥ t → ALERT.
        """
        hard_scores = [float(s) for s in hard_scores]
        hard_labels = [int(l) for l in hard_labels]

        use_pseudo = (
            pseudo_scores is not None
            and pseudo_labels is not None
            and len(pseudo_scores) >= self._n_pseudo_min
        )

        if use_pseudo:
            pseudo_scores = [float(s) for s in pseudo_scores]
            pseudo_labels = [int(l) for l in pseudo_labels]

            weights_hard   = [self._lambda_hard] * len(hard_scores)
            weights_pseudo = [1.0] * len(pseudo_scores)

            all_scores  = hard_scores  + pseudo_scores
            all_labels  = hard_labels  + pseudo_labels
            all_weights = weights_hard + weights_pseudo
        else:
            if pseudo_scores is not None:
                logger.warning(
                    "DomainAdaptiveCalibrator: only %d pseudo-labels provided (min=%d). "
                    "Falling back to hard-label-only calibration.",
                    len(pseudo_scores or []),
                    self._n_pseudo_min,
                )
            all_scores  = hard_scores
            all_labels  = hard_labels
            all_weights = [1.0] * len(hard_scores)

        self._threshold = self._conformal_threshold(
            all_scores, all_labels, all_weights, self._alpha
        )
        logger.info(
            "DomainAdaptiveCalibrator: threshold=%.4f "
            "(n_hard=%d, n_pseudo=%d, λ=%.1f, noise=%.2f)",
            self._threshold,
            len(hard_scores),
            len(pseudo_scores) if use_pseudo else 0,
            self._lambda_hard,
            self._noise,
        )
        return self._threshold

    def get_threshold(self) -> float:
        """Return the fitted threshold. Raises RuntimeError if fit() not called."""
        if self._threshold is None:
            raise RuntimeError(
                "DomainAdaptiveCalibrator.fit() must be called before get_threshold()."
            )
        return self._threshold

    @staticmethod
    def _conformal_threshold(
        scores: List[float],
        labels: List[int],
        weights: List[float],
        alpha: float,
    ) -> float:
        """
        Weighted empirical quantile on correct chains: return the minimum score t
        such that the weighted fraction of correct chains scoring ≥ t is ≤ alpha.

        Falls back to the (1-alpha) quantile of ALL scores when fewer than 5
        correct examples exist.
        """
        scores_arr  = np.array(scores)
        labels_arr  = np.array(labels)
        weights_arr = np.array(weights)

        correct_mask = labels_arr == 0

        if correct_mask.sum() < 5:
            logger.warning(
                "DomainAdaptiveCalibrator: fewer than 5 correct chains in pool. "
                "Using global (1-α) quantile as fallback."
            )
            return float(np.quantile(scores_arr, 1.0 - alpha))

        correct_scores  = scores_arr[correct_mask]
        correct_weights = weights_arr[correct_mask]

        # Weighted quantile via sorting
        order          = np.argsort(correct_scores)
        sorted_scores  = correct_scores[order]
        sorted_weights = correct_weights[order]
        cumulative     = np.cumsum(sorted_weights)
        total          = cumulative[-1]
        target_cum     = total * (1.0 - alpha)
        idx            = np.searchsorted(cumulative, target_cum)
        idx            = min(idx, len(sorted_scores) - 1)
        return float(sorted_scores[idx])
