"""
EnergyAwareCalibrator — Dynamic precision-recall tradeoff via energy budget (Patent 8 / P8).

Ported from Patent 8 (ECL — Energy-Conscious Learning) experiments.
Adjusts alert thresholds dynamically based on an "energy" signal representing
available computational budget or downstream cost tolerance.

Core idea from Patent 8:
  When energy is HIGH (abundant resources), tighten precision by raising
  the alert threshold → fewer false alarms, miss some real failures.
  When energy is LOW (constrained resources), lower the threshold → catch
  more failures (higher recall) at the cost of more false alarms.

This is a direct port of the precision-recall budget principle from Patent 8's
ECL calibration: energy depletion drives threshold adaptation.

Usage::

    cal = EnergyAwareCalibrator()
    cal.fit(risk_scores, labels)   # optional; enables Platt scaling

    # At inference time, provide current energy level (0–1):
    threshold = cal.threshold(energy=0.6)   # adapts based on energy
    is_alert = risk_score >= threshold

    # Or score directly:
    result = cal.calibrate(risk_score, energy=0.6)
    print(result.alert, result.threshold, result.energy_regime)
Purpose:    Adjusts alert thresholds based on an energy signal representing available computational budget. High energy tightens precision; low energy raises recall to catch more failures.
Exports:    EnergyAwareCalibrator
Added:      v0.44.0
Status:     Stable
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Tuple

import numpy as np


@dataclass
class EnergyCalibrationResult:
    """Result of EnergyAwareCalibrator.calibrate()."""
    risk_score: float
    """Input risk score [0, 1]."""
    calibrated_risk: float
    """Platt-scaled risk (= risk_score if not fitted)."""
    threshold: float
    """Adaptive alert threshold for current energy level."""
    alert: bool
    """True when calibrated_risk >= threshold."""
    energy: float
    """Input energy level [0, 1]."""
    energy_regime: str
    """'high' (≥0.7), 'medium' (0.3–0.7), 'low' (<0.3)."""


class EnergyAwareCalibrator:
    """
    Dynamic precision-recall calibrator driven by energy budget signal.

    The energy parameter represents available computational/cost budget:
    - 1.0 = full resources (strict precision mode, high threshold)
    - 0.0 = depleted resources (aggressive recall mode, low threshold)

    Default threshold schedule (linear interpolation):
      energy=1.0 → threshold = high_threshold  (default 0.80)
      energy=0.0 → threshold = low_threshold   (default 0.50)

    Optionally: fit Platt scaling on labeled risk scores to improve
    calibration before applying the energy-adaptive threshold.

    Parameters
    ----------
    high_threshold : float
        Alert threshold when energy is full (1.0). Maximises precision.
        Default: 0.80
    low_threshold : float
        Alert threshold when energy is depleted (0.0). Maximises recall.
        Default: 0.50
    energy_exponent : float
        Exponent for energy→threshold mapping. >1 = slower drop,
        <1 = faster drop. Default 1.0 (linear).
    """

    def __init__(
        self,
        high_threshold: float = 0.80,
        low_threshold: float = 0.50,
        energy_exponent: float = 1.0,
    ) -> None:
        if not (0.0 < low_threshold <= high_threshold <= 1.0):
            raise ValueError(
                "Require 0 < low_threshold <= high_threshold <= 1.0, "
                f"got low={low_threshold}, high={high_threshold}"
            )
        self.high_threshold = high_threshold
        self.low_threshold = low_threshold
        self.energy_exponent = energy_exponent
        self._T: float = 1.0    # Platt temperature
        self._b: float = 0.0    # Platt bias
        self._fitted = False

    @property
    def is_fitted(self) -> bool:
        return self._fitted

    def fit(
        self, scores: List[float], labels: List[int], max_iter: int = 200
    ) -> "EnergyAwareCalibrator":
        """
        Fit Platt scaling (temperature + bias) on labeled risk scores.

        Parameters
        ----------
        scores : List[float]
            Raw risk scores [0, 1] from AgentGuard.score_chain().
        labels : List[int]
            1 = chain is wrong (risky), 0 = chain is correct. Same length as scores.
        max_iter : int
            Maximum scipy optimization iterations. Default 200.

        Returns
        -------
        self
        """
        try:
            from scipy.optimize import minimize as _minimize
        except ImportError:
            # Without scipy, fall back to identity calibration
            self._fitted = True
            return self

        s = np.array(scores, dtype=float)
        y = np.array(labels, dtype=float)

        def _nll(params: np.ndarray) -> float:
            T, b = params[0], params[1]
            logits = T * s + b
            logits = np.clip(logits, -50, 50)
            probs = 1.0 / (1.0 + np.exp(-logits))
            probs = np.clip(probs, 1e-7, 1 - 1e-7)
            return -float(np.mean(y * np.log(probs) + (1 - y) * np.log(1 - probs)))

        res = _minimize(_nll, x0=[1.0, 0.0], method="Nelder-Mead",
                        options={"maxiter": max_iter, "xatol": 1e-5})
        self._T = float(res.x[0])
        self._b = float(res.x[1])
        self._fitted = True
        return self

    def _platt_scale(self, risk: float) -> float:
        """Apply Platt scaling. Returns risk unchanged if not fitted."""
        if not self._fitted:
            return float(risk)
        logit = self._T * risk + self._b
        logit = max(-50.0, min(50.0, logit))
        return float(1.0 / (1.0 + np.exp(-logit)))

    def threshold(self, energy: float) -> float:
        """
        Compute the adaptive alert threshold for a given energy level.

        Parameters
        ----------
        energy : float
            Current energy budget [0, 1]. 1.0 = full, 0.0 = depleted.

        Returns
        -------
        float: Alert threshold for this energy level.
        """
        energy = float(np.clip(energy, 0.0, 1.0))
        t = energy ** self.energy_exponent
        return self.low_threshold + t * (self.high_threshold - self.low_threshold)

    def calibrate(self, risk_score: float, energy: float) -> EnergyCalibrationResult:
        """
        Calibrate a risk score under a given energy budget.

        Parameters
        ----------
        risk_score : float
            Raw risk score [0, 1] from AgentGuard.
        energy : float
            Current energy level [0, 1].

        Returns
        -------
        EnergyCalibrationResult
        """
        energy = float(np.clip(energy, 0.0, 1.0))
        cal_risk = self._platt_scale(risk_score)
        thr = self.threshold(energy)

        if energy >= 0.7:
            regime = "high"
        elif energy >= 0.3:
            regime = "medium"
        else:
            regime = "low"

        return EnergyCalibrationResult(
            risk_score=float(risk_score),
            calibrated_risk=round(cal_risk, 4),
            threshold=round(thr, 4),
            alert=cal_risk >= thr,
            energy=round(energy, 4),
            energy_regime=regime,
        )

    def calibrate_batch(
        self, risk_scores: List[float], energies: List[float]
    ) -> List[EnergyCalibrationResult]:
        """
        Calibrate a batch of (risk_score, energy) pairs.

        Parameters
        ----------
        risk_scores : List[float]
        energies : List[float]
            Same length as risk_scores. Each entry is the energy for that chain.

        Returns
        -------
        List[EnergyCalibrationResult]
        """
        return [
            self.calibrate(r, e) for r, e in zip(risk_scores, energies)
        ]

    def optimal_energy_for_recall(self, target_recall: float) -> float:
        """
        Return the energy level that achieves a target recall rate.

        Lower recall target → higher energy (stricter precision mode is OK).
        Higher recall target → lower energy (more aggressive threshold needed).

        This is an approximation: assumes risk scores are uniformly distributed.
        Use calibrate_batch for precision on real data.

        Parameters
        ----------
        target_recall : float
            Desired recall [0, 1]. 1.0 = catch all failures (low threshold needed).

        Returns
        -------
        float: Energy level [0, 1] that achieves approximate target_recall.
        """
        target_recall = float(np.clip(target_recall, 0.0, 1.0))
        # recall ↑ as threshold ↓ as energy ↓
        # recall=1.0 → energy=0.0 (lowest threshold)
        # recall=0.0 → energy=1.0 (highest threshold)
        return float(np.clip(1.0 - target_recall, 0.0, 1.0))
