#!/usr/bin/env python3
"""
Confidence Calibration Module
==============================
Calibrates model confidence scores to match empirical accuracy.

Implements multiple calibration methods:
1. Isotonic Regression - Non-parametric, monotonic calibration
2. Platt Scaling - Parametric sigmoid calibration
3. Temperature Scaling - Simple division by learned temperature

Includes evaluation metrics:
- Expected Calibration Error (ECE)
- Maximum Calibration Error (MCE)
- Reliability diagrams

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from typing import Dict, List, Optional, Tuple, Any, Union
from dataclasses import dataclass
from loguru import logger
from scipy.optimize import minimize
from sklearn.isotonic import IsotonicRegression
import warnings


@dataclass
class CalibrationMetrics:
    """Metrics for evaluating calibration quality."""

    ece: float  # Expected Calibration Error
    mce: float  # Maximum Calibration Error
    ace: float  # Average Calibration Error
    brier_score: float  # Brier Score (for probabilistic accuracy)
    reliability_data: Dict[str, np.ndarray]  # Data for reliability diagram

    def to_dict(self) -> Dict[str, Any]:
        return {
            'ece': float(self.ece),
            'mce': float(self.mce),
            'ace': float(self.ace),
            'brier_score': float(self.brier_score)
        }


class ConfidenceCalibrator:
    """
    Calibrates confidence scores to match empirical accuracy.

    Supports multiple calibration methods:
    - isotonic: Non-parametric isotonic regression
    - platt: Parametric Platt scaling (sigmoid)
    - temperature: Temperature scaling (division by T)
    - beta: Beta calibration (generalized sigmoid)

    Usage:
        calibrator = ConfidenceCalibrator(method='isotonic')
        calibrator.fit(y_true, y_prob)

        calibrated_probs = calibrator.calibrate(new_probs)

        metrics = calibrator.evaluate(y_true, y_prob)
        print(f"ECE: {metrics.ece:.4f}")
    """

    SUPPORTED_METHODS = ['isotonic', 'platt', 'temperature', 'beta', 'ensemble']

    def __init__(
        self,
        method: str = 'isotonic',
        n_bins: int = 15,
        device: str = 'cpu'
    ):
        """
        Initialize the calibrator.

        Args:
            method: Calibration method ('isotonic', 'platt', 'temperature', 'beta', 'ensemble')
            n_bins: Number of bins for ECE calculation
            device: PyTorch device for temperature scaling
        """
        if method not in self.SUPPORTED_METHODS:
            raise ValueError(f"Unknown method: {method}. Supported: {self.SUPPORTED_METHODS}")

        self.method = method
        self.n_bins = n_bins
        self.device = device

        # Calibration models (initialized during fit)
        self._isotonic: Optional[IsotonicRegression] = None
        self._platt_params: Optional[Tuple[float, float]] = None  # (a, b) for sigmoid
        self._temperature: Optional[float] = None
        self._beta_params: Optional[Tuple[float, float, float]] = None  # (a, b, c)

        self._is_fitted = False

        logger.info(f"ConfidenceCalibrator initialized with method={method}")

    def fit(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray,
        sample_weight: Optional[np.ndarray] = None
    ) -> 'ConfidenceCalibrator':
        """
        Fit the calibration model.

        Args:
            y_true: True binary labels (0 or 1)
            y_prob: Predicted probabilities (0 to 1)
            sample_weight: Optional sample weights

        Returns:
            Self for chaining
        """
        y_true = np.asarray(y_true).ravel()
        y_prob = np.asarray(y_prob).ravel()

        # Clip probabilities to avoid numerical issues
        y_prob = np.clip(y_prob, 1e-8, 1 - 1e-8)

        if len(y_true) != len(y_prob):
            raise ValueError("y_true and y_prob must have same length")

        if self.method == 'isotonic':
            self._fit_isotonic(y_true, y_prob, sample_weight)
        elif self.method == 'platt':
            self._fit_platt(y_true, y_prob, sample_weight)
        elif self.method == 'temperature':
            self._fit_temperature(y_true, y_prob)
        elif self.method == 'beta':
            self._fit_beta(y_true, y_prob)
        elif self.method == 'ensemble':
            self._fit_ensemble(y_true, y_prob, sample_weight)

        self._is_fitted = True
        logger.info(f"Calibrator fitted with {len(y_true)} samples")

        return self

    def _fit_isotonic(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray,
        sample_weight: Optional[np.ndarray]
    ):
        """Fit isotonic regression calibrator."""
        self._isotonic = IsotonicRegression(
            y_min=0.0,
            y_max=1.0,
            out_of_bounds='clip'
        )
        self._isotonic.fit(y_prob, y_true, sample_weight=sample_weight)

    def _fit_platt(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray,
        sample_weight: Optional[np.ndarray]
    ):
        """
        Fit Platt scaling (sigmoid calibration).

        P(y=1|f) = 1 / (1 + exp(a*f + b))
        where f = log(p / (1-p)) is the logit of original probability
        """
        # Convert to logits
        logits = np.log(y_prob / (1 - y_prob))

        # Optimize sigmoid parameters using NLL loss
        def nll_loss(params):
            a, b = params
            calibrated = 1.0 / (1.0 + np.exp(a * logits + b))
            calibrated = np.clip(calibrated, 1e-8, 1 - 1e-8)

            loss = -np.mean(
                y_true * np.log(calibrated) +
                (1 - y_true) * np.log(1 - calibrated)
            )

            if sample_weight is not None:
                loss = -np.average(
                    y_true * np.log(calibrated) +
                    (1 - y_true) * np.log(1 - calibrated),
                    weights=sample_weight
                )

            return loss

        # Initialize and optimize
        result = minimize(
            nll_loss,
            x0=[1.0, 0.0],
            method='L-BFGS-B',
            bounds=[(-10, 10), (-10, 10)]
        )

        self._platt_params = (result.x[0], result.x[1])
        logger.debug(f"Platt params: a={result.x[0]:.4f}, b={result.x[1]:.4f}")

    def _fit_temperature(self, y_true: np.ndarray, y_prob: np.ndarray):
        """
        Fit temperature scaling.

        calibrated_prob = softmax(logits / T)
        For binary case: calibrated = sigmoid(logit / T)
        """
        # Convert to logits
        logits = np.log(y_prob / (1 - y_prob))

        # Find optimal temperature using NLL
        def nll_loss(temperature):
            t = temperature[0]
            if t <= 0:
                return 1e10

            calibrated = 1.0 / (1.0 + np.exp(-logits / t))
            calibrated = np.clip(calibrated, 1e-8, 1 - 1e-8)

            loss = -np.mean(
                y_true * np.log(calibrated) +
                (1 - y_true) * np.log(1 - calibrated)
            )
            return loss

        result = minimize(
            nll_loss,
            x0=[1.0],
            method='L-BFGS-B',
            bounds=[(0.01, 10.0)]
        )

        self._temperature = result.x[0]
        logger.debug(f"Temperature: T={self._temperature:.4f}")

    def _fit_beta(self, y_true: np.ndarray, y_prob: np.ndarray):
        """
        Fit beta calibration.

        calibrated = 1 / (1 + 1/(exp(c) * (p^a / (1-p)^b)))

        This is a more flexible version of Platt scaling.
        """
        # Avoid log(0)
        p = np.clip(y_prob, 1e-8, 1 - 1e-8)

        def nll_loss(params):
            a, b, c = params

            # Beta calibration formula
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                ratio = np.exp(c) * np.power(p, a) / np.power(1 - p, b)

            calibrated = ratio / (1 + ratio)
            calibrated = np.clip(calibrated, 1e-8, 1 - 1e-8)

            loss = -np.mean(
                y_true * np.log(calibrated) +
                (1 - y_true) * np.log(1 - calibrated)
            )
            return loss

        # Try multiple initializations
        best_result = None
        best_loss = float('inf')

        for init in [(1, 1, 0), (0.5, 0.5, 0), (2, 2, 0)]:
            try:
                result = minimize(
                    nll_loss,
                    x0=init,
                    method='L-BFGS-B',
                    bounds=[(0.01, 10), (0.01, 10), (-5, 5)]
                )
                if result.fun < best_loss:
                    best_loss = result.fun
                    best_result = result
            except Exception:
                continue

        if best_result is not None:
            self._beta_params = tuple(best_result.x)
            logger.debug(f"Beta params: a={best_result.x[0]:.4f}, "
                         f"b={best_result.x[1]:.4f}, c={best_result.x[2]:.4f}")
        else:
            # Fallback to identity
            self._beta_params = (1.0, 1.0, 0.0)

    def _fit_ensemble(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray,
        sample_weight: Optional[np.ndarray]
    ):
        """Fit ensemble of calibrators and use best one."""
        # Fit all methods
        self._fit_isotonic(y_true, y_prob, sample_weight)
        self._fit_platt(y_true, y_prob, sample_weight)
        self._fit_temperature(y_true, y_prob)

        # Evaluate each on training data
        methods = ['isotonic', 'platt', 'temperature']
        eces = []

        for method in methods:
            calibrated = self._calibrate_with_method(y_prob, method)
            ece = self._compute_ece(y_true, calibrated)
            eces.append(ece)

        # Select best method
        best_idx = np.argmin(eces)
        self._ensemble_best = methods[best_idx]
        logger.info(f"Ensemble selected: {self._ensemble_best} (ECE={eces[best_idx]:.4f})")

    def calibrate(self, probabilities: np.ndarray) -> np.ndarray:
        """
        Calibrate probabilities using fitted model.

        Args:
            probabilities: Uncalibrated probabilities (0 to 1)

        Returns:
            Calibrated probabilities
        """
        if not self._is_fitted:
            raise RuntimeError("Calibrator must be fitted before calibrating")

        probabilities = np.asarray(probabilities).ravel()
        probabilities = np.clip(probabilities, 1e-8, 1 - 1e-8)

        if self.method == 'ensemble':
            return self._calibrate_with_method(probabilities, self._ensemble_best)

        return self._calibrate_with_method(probabilities, self.method)

    def _calibrate_with_method(
        self,
        probabilities: np.ndarray,
        method: str
    ) -> np.ndarray:
        """Calibrate using specified method."""
        probabilities = np.clip(probabilities, 1e-8, 1 - 1e-8)

        if method == 'isotonic':
            if self._isotonic is None:
                return probabilities
            return self._isotonic.predict(probabilities)

        elif method == 'platt':
            if self._platt_params is None:
                return probabilities
            a, b = self._platt_params
            logits = np.log(probabilities / (1 - probabilities))
            return 1.0 / (1.0 + np.exp(a * logits + b))

        elif method == 'temperature':
            if self._temperature is None:
                return probabilities
            logits = np.log(probabilities / (1 - probabilities))
            return 1.0 / (1.0 + np.exp(-logits / self._temperature))

        elif method == 'beta':
            if self._beta_params is None:
                return probabilities
            a, b, c = self._beta_params
            p = probabilities
            ratio = np.exp(c) * np.power(p, a) / np.power(1 - p, b)
            return ratio / (1 + ratio)

        else:
            return probabilities

    def _compute_ece(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray
    ) -> float:
        """Compute Expected Calibration Error."""
        bin_boundaries = np.linspace(0, 1, self.n_bins + 1)
        bin_lowers = bin_boundaries[:-1]
        bin_uppers = bin_boundaries[1:]

        ece = 0.0
        for bin_lower, bin_upper in zip(bin_lowers, bin_uppers):
            in_bin = (y_prob > bin_lower) & (y_prob <= bin_upper)
            prop_in_bin = in_bin.mean()

            if prop_in_bin > 0:
                avg_confidence = y_prob[in_bin].mean()
                avg_accuracy = y_true[in_bin].mean()
                ece += np.abs(avg_confidence - avg_accuracy) * prop_in_bin

        return ece

    def compute_expected_calibration_error(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray,
        calibrated: bool = True
    ) -> float:
        """
        Compute ECE (Expected Calibration Error).

        ECE = sum_b (|B_b|/n) * |acc(B_b) - conf(B_b)|

        Args:
            y_true: True labels
            y_prob: Predicted probabilities
            calibrated: If True, calibrate probabilities first

        Returns:
            ECE value (lower is better)
        """
        y_true = np.asarray(y_true).ravel()
        y_prob = np.asarray(y_prob).ravel()

        if calibrated and self._is_fitted:
            y_prob = self.calibrate(y_prob)

        return self._compute_ece(y_true, y_prob)

    def compute_reliability_diagram(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray,
        calibrated: bool = True
    ) -> Dict[str, np.ndarray]:
        """
        Compute data for reliability diagram.

        Args:
            y_true: True labels
            y_prob: Predicted probabilities
            calibrated: If True, calibrate probabilities first

        Returns:
            Dictionary with:
            - bin_centers: Center of each bin
            - bin_accuracies: Average accuracy in each bin
            - bin_confidences: Average confidence in each bin
            - bin_counts: Number of samples in each bin
            - perfect_calibration: Diagonal line (for reference)
        """
        y_true = np.asarray(y_true).ravel()
        y_prob = np.asarray(y_prob).ravel()

        if calibrated and self._is_fitted:
            y_prob = self.calibrate(y_prob)

        bin_boundaries = np.linspace(0, 1, self.n_bins + 1)
        bin_lowers = bin_boundaries[:-1]
        bin_uppers = bin_boundaries[1:]
        bin_centers = (bin_lowers + bin_uppers) / 2

        bin_accuracies = []
        bin_confidences = []
        bin_counts = []

        for bin_lower, bin_upper in zip(bin_lowers, bin_uppers):
            in_bin = (y_prob > bin_lower) & (y_prob <= bin_upper)
            count = in_bin.sum()

            if count > 0:
                bin_accuracies.append(y_true[in_bin].mean())
                bin_confidences.append(y_prob[in_bin].mean())
            else:
                bin_accuracies.append(np.nan)
                bin_confidences.append(np.nan)

            bin_counts.append(count)

        return {
            'bin_centers': bin_centers,
            'bin_accuracies': np.array(bin_accuracies),
            'bin_confidences': np.array(bin_confidences),
            'bin_counts': np.array(bin_counts),
            'perfect_calibration': np.linspace(0, 1, 100)
        }

    def evaluate(
        self,
        y_true: np.ndarray,
        y_prob: np.ndarray,
        calibrated: bool = True
    ) -> CalibrationMetrics:
        """
        Evaluate calibration quality with multiple metrics.

        Args:
            y_true: True labels
            y_prob: Predicted probabilities
            calibrated: If True, evaluate calibrated probabilities

        Returns:
            CalibrationMetrics object
        """
        y_true = np.asarray(y_true).ravel()
        y_prob = np.asarray(y_prob).ravel()

        if calibrated and self._is_fitted:
            y_prob_cal = self.calibrate(y_prob)
        else:
            y_prob_cal = y_prob

        # ECE
        ece = self._compute_ece(y_true, y_prob_cal)

        # MCE (Maximum Calibration Error)
        bin_boundaries = np.linspace(0, 1, self.n_bins + 1)
        bin_lowers = bin_boundaries[:-1]
        bin_uppers = bin_boundaries[1:]

        bin_errors = []
        for bin_lower, bin_upper in zip(bin_lowers, bin_uppers):
            in_bin = (y_prob_cal > bin_lower) & (y_prob_cal <= bin_upper)
            if in_bin.sum() > 0:
                avg_conf = y_prob_cal[in_bin].mean()
                avg_acc = y_true[in_bin].mean()
                bin_errors.append(np.abs(avg_conf - avg_acc))

        mce = max(bin_errors) if bin_errors else 0.0

        # ACE (Average Calibration Error - unweighted ECE)
        ace = np.mean(bin_errors) if bin_errors else 0.0

        # Brier Score
        brier_score = np.mean((y_prob_cal - y_true) ** 2)

        # Reliability diagram data
        reliability_data = self.compute_reliability_diagram(y_true, y_prob, calibrated)

        return CalibrationMetrics(
            ece=ece,
            mce=mce,
            ace=ace,
            brier_score=brier_score,
            reliability_data=reliability_data
        )


class TemperatureScalingNN(nn.Module):
    """
    Temperature scaling implemented as a PyTorch module.

    Useful for end-to-end training with neural networks.
    """

    def __init__(self, init_temperature: float = 1.5):
        """
        Initialize temperature scaling module.

        Args:
            init_temperature: Initial temperature value
        """
        super().__init__()
        self.temperature = nn.Parameter(torch.tensor([init_temperature]))

    def forward(self, logits: torch.Tensor) -> torch.Tensor:
        """
        Apply temperature scaling to logits.

        Args:
            logits: Raw logits from model

        Returns:
            Temperature-scaled logits
        """
        return logits / self.temperature

    def calibrate_probs(self, probs: torch.Tensor) -> torch.Tensor:
        """
        Calibrate probabilities (convert to logits, scale, convert back).

        Args:
            probs: Probability tensor

        Returns:
            Calibrated probabilities
        """
        # Convert to logits
        probs = torch.clamp(probs, 1e-8, 1 - 1e-8)
        logits = torch.log(probs / (1 - probs))

        # Scale
        scaled_logits = self.forward(logits)

        # Convert back to probabilities
        return torch.sigmoid(scaled_logits)


def calibrate_ensemble_confidence(
    predictions: List[Dict[str, float]],
    y_true: np.ndarray,
    method: str = 'isotonic'
) -> ConfidenceCalibrator:
    """
    Convenience function to calibrate ensemble confidence scores.

    Args:
        predictions: List of prediction dictionaries with 'confidence' key
        y_true: True binary outcomes
        method: Calibration method

    Returns:
        Fitted ConfidenceCalibrator
    """
    confidences = np.array([p.get('confidence', 0.5) for p in predictions])
    calibrator = ConfidenceCalibrator(method=method)
    calibrator.fit(y_true, confidences)
    return calibrator


if __name__ == "__main__":
    # Test calibration module
    np.random.seed(42)

    print("Testing ConfidenceCalibrator")
    print("=" * 60)

    # Generate synthetic uncalibrated probabilities
    # Overconfident model
    n_samples = 1000
    true_probs = np.random.beta(2, 2, n_samples)
    y_true = (np.random.rand(n_samples) < true_probs).astype(int)

    # Simulate overconfident predictions
    y_prob = np.clip(true_probs * 1.3, 0.01, 0.99)

    print(f"Samples: {n_samples}")
    print(f"True positive rate: {y_true.mean():.3f}")
    print(f"Mean predicted probability: {y_prob.mean():.3f}")

    # Test each calibration method
    methods = ['isotonic', 'platt', 'temperature', 'beta', 'ensemble']

    for method in methods:
        print(f"\n--- {method.upper()} Calibration ---")

        calibrator = ConfidenceCalibrator(method=method)
        calibrator.fit(y_true[:800], y_prob[:800])

        # Evaluate on held-out data
        y_test = y_true[800:]
        prob_test = y_prob[800:]

        # Before calibration
        before_metrics = calibrator.evaluate(y_test, prob_test, calibrated=False)
        print(f"Before: ECE={before_metrics.ece:.4f}, Brier={before_metrics.brier_score:.4f}")

        # After calibration
        after_metrics = calibrator.evaluate(y_test, prob_test, calibrated=True)
        print(f"After:  ECE={after_metrics.ece:.4f}, Brier={after_metrics.brier_score:.4f}")

        improvement = (before_metrics.ece - after_metrics.ece) / before_metrics.ece * 100
        print(f"ECE Improvement: {improvement:.1f}%")

    # Test reliability diagram
    print("\n--- Reliability Diagram Data ---")
    calibrator = ConfidenceCalibrator(method='isotonic')
    calibrator.fit(y_true[:800], y_prob[:800])

    diagram_data = calibrator.compute_reliability_diagram(y_true[800:], y_prob[800:])
    print(f"Bin centers: {diagram_data['bin_centers'][:5]}...")
    print(f"Bin accuracies: {diagram_data['bin_accuracies'][:5]}...")
    print(f"Bin counts: {diagram_data['bin_counts'][:5]}...")

    print("\nConfidenceCalibrator test complete!")
