"""
ModelTransferCalibrator — Platt-scaling transfer between judge models.
=======================================================================

When switching judge models (e.g. Haiku → Sonnet-3.5), existing calibration
breaks (AUROC ~0.508 cross-model). This component adapts a source-model
CalibParams to a new model using Platt scaling with ≥2 labeled chains.

Usage
-----
    from llm_guard import ModelTransferCalibrator

    # After switching from Haiku to Sonnet-3.5, collect 3-5 labeled chains
    cal = ModelTransferCalibrator(haiku_calib_params, target_model="sonnet-3-5")
    raw_scores = [guard.score_chain(q, s, a).risk_score for q, s, a in new_chains]
    cal.fit(new_chains, new_labels, raw_scores=raw_scores)

    # Use for inference
    calibrated = cal.predict_one(raw_score)

    # Export for federated merge
    params = cal.export_params()
"""
from __future__ import annotations

from typing import List, Optional

import numpy as np

from llm_guard.federated_calibrator import CalibParams


def _sigmoid(x: float) -> float:
    return 1.0 / (1.0 + np.exp(-x))


def _source_map(raw: float, source_params: CalibParams) -> float:
    """Apply source isotonic mapping via linear interpolation."""
    return float(np.interp(
        raw,
        source_params.x_breaks,
        source_params.y_breaks,
        left=source_params.y_breaks[0],
        right=source_params.y_breaks[-1],
    ))


class ModelTransferCalibrator:
    """
    Adapt calibration from a source judge model to a target model.

    Requires ≥2 labeled chains from the target model.
    Recommend 3–5 chains for stable Platt scaling.
    """

    def __init__(self, source_params: CalibParams, target_model: str) -> None:
        self._source_params = source_params
        self._target_model  = target_model
        self._T: Optional[float] = None
        self._b: Optional[float] = None
        self._n_new: int = 0

    # ------------------------------------------------------------------

    def fit(
        self,
        new_chains: List[dict],
        new_labels: List[int],
        raw_scores: Optional[List[float]] = None,
    ) -> "ModelTransferCalibrator":
        """
        Fit Platt scaling transfer on ≥2 labeled chains.

        Parameters
        ----------
        new_chains : list of chain dicts (used only for count)
        new_labels : list of int (1=correct/low-risk, 0=wrong/high-risk).
                     NOTE: This is OPPOSITE to FederatedCalibrator's convention (1=wrong).
                     Do not mix label conventions between the two classes.
        raw_scores : raw risk scores from the NEW model for each chain.
                     Must be provided by caller (usually from guard.score_chain()).
        """
        if len(new_labels) < 2:
            raise ValueError(
                "ModelTransferCalibrator.fit() requires at least 2 labeled chains. "
                "Recommend 3–5 for stable transfer."
            )
        if raw_scores is not None and len(raw_scores) != len(new_labels):
            raise ValueError(
                f"raw_scores length ({len(raw_scores)}) must match new_labels length ({len(new_labels)})."
            )
        if len(new_chains) != len(new_labels):
            raise ValueError(
                f"new_chains length ({len(new_chains)}) must match new_labels length ({len(new_labels)})."
            )
        import warnings as _warnings
        if len(set(new_labels)) < 2:
            _warnings.warn(
                "ModelTransferCalibrator.fit() received all-identical labels. "
                "Platt scaling will produce a degenerate (near-constant) output. "
                "Provide at least one label of each class for valid transfer calibration.",
                UserWarning,
                stacklevel=2,
            )
        if raw_scores is None:
            # Fall back to reading from chain dicts
            raw_scores = [c.get("risk_score", 0.5) for c in new_chains]

        # Step 1: apply source isotonic mapping
        source_cal = np.array(
            [_source_map(s, self._source_params) for s in raw_scores],
            dtype=np.float64,
        )
        y = np.array(new_labels, dtype=np.float64)

        # Step 2: fit Platt scaling — minimize binary cross-entropy
        from scipy.optimize import minimize

        def neg_log_loss(params):
            T, b = params
            p = np.clip(1.0 / (1.0 + np.exp(-(T * source_cal + b))), 1e-9, 1 - 1e-9)
            return -np.mean(y * np.log(p) + (1 - y) * np.log(1 - p))

        res = minimize(
            neg_log_loss,
            x0=[1.0, 0.0],
            method="L-BFGS-B",
            bounds=[(0.1, 10.0), (-5.0, 5.0)],
        )
        self._T = float(res.x[0])
        self._b = float(res.x[1])
        self._n_new = len(new_chains)
        return self

    # ------------------------------------------------------------------

    def predict_one(self, raw_score: float) -> float:
        """Calibrate a single raw score from the new model."""
        if self._T is None:
            raise RuntimeError(
                "ModelTransferCalibrator.predict_one() called before fit(). "
                "Call fit() with labeled chains from the target model first."
            )
        source_cal = _source_map(raw_score, self._source_params)
        return float(np.clip(_sigmoid(self._T * source_cal + self._b), 0.0, 1.0))

    # ------------------------------------------------------------------

    def export_params(self) -> CalibParams:
        """
        Export as CalibParams compatible with FederatedCalibrator.merge().

        Samples the composed transform at 20 evenly spaced x-values.
        n_samples = number of transfer chains used in fit().
        """
        if self._T is None:
            raise RuntimeError("Call fit() before export_params().")
        x_grid = np.linspace(0.0, 1.0, 20)
        y_grid = np.array([self.predict_one(float(x)) for x in x_grid])
        # Enforce monotonicity (Platt scaling is monotone but float precision may jitter)
        for i in range(1, len(y_grid)):
            if y_grid[i] < y_grid[i - 1]:
                y_grid[i] = y_grid[i - 1]
        return CalibParams(
            x_breaks=x_grid.tolist(),
            y_breaks=y_grid.tolist(),
            n_samples=self._n_new,
            epsilon=None,
            site_id="transfer",
        )
