"""
Calibration Component for GRU Model Probabilities.

Provides methods for temperature scaling (Platt scaling) and generating 
action signals based on calibrated probabilities.
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize_scalar
from scipy.special import expit, logit
from typing import Tuple, Optional
import logging
import os

logger = logging.getLogger(__name__)

class Calibrator:
    """Handles probability calibration using temperature scaling."""

    def __init__(self, edge_threshold: float):
        """
        Initialize the Calibrator.

        Args:
            edge_threshold (float): Minimum calibrated edge magnitude for taking a trade (e.g., 0.55 means P(up) > 0.55 or P(down) > 0.55).
        """
        self.edge_threshold = edge_threshold
        self.optimal_T: Optional[float] = None # Stores the calculated temperature
        logger.info(f"Calibrator initialized with edge threshold: {self.edge_threshold}")

    def _nll_objective(self, T: float, logit_p: np.ndarray, y_true: np.ndarray) -> float:
        """Negative log-likelihood objective function for temperature optimization."""
        if T <= 0:
             return np.inf # Temperature must be positive
        p_cal = expit(logit_p / T)
        # Binary cross-entropy (NLL)
        eps = 1e-12 # Epsilon for numerical stability
        p_cal = np.clip(p_cal, eps, 1 - eps)
        # Ensure y_true is broadcastable if necessary (should be 1D)
        nll = -(y_true * np.log(p_cal) + (1 - y_true) * np.log(1 - p_cal))
        return float(np.mean(nll))

    def optimise_temperature(self, p_raw: np.ndarray, y_true: np.ndarray, bounds=(0.1, 10.0)) -> float:
        """
        Finds the optimal temperature `T` by minimizing NLL on validation data.

        Args:
            p_raw (np.ndarray): Raw model probabilities (validation set).
            y_true (np.ndarray): True binary labels (validation set).
            bounds (tuple): Bounds for the temperature search.

        Returns:
            float: The optimal temperature found.
        """
        logger.info(f"Optimizing calibration temperature using {len(p_raw)} samples...")
        # Ensure inputs are flat numpy arrays and correct type
        p_raw = np.asarray(p_raw).flatten().astype(float)
        y_true = np.asarray(y_true).flatten().astype(int)
        
        # Clip raw probabilities and compute logits for numerical stability
        eps = 1e-7
        p_clipped = np.clip(p_raw, eps, 1 - eps)
        logit_p = logit(p_clipped)
        
        # Handle cases where all predictions are the same (logit might be inf)
        if np.isinf(logit_p).any():
             logger.warning("Infinite values encountered in logits during temperature scaling. Clipping may be too aggressive or predictions are uniform. Returning T=1.0")
             self.optimal_T = 1.0
             return 1.0

        try:
            res = minimize_scalar(
                lambda T: self._nll_objective(T, logit_p, y_true),
                bounds=bounds,
                method="bounded",
            )
            
            if res.success:
                 optimal_T_found = float(res.x)
                 logger.info(f"Optimal temperature found: T = {optimal_T_found:.4f}")
                 self.optimal_T = optimal_T_found
                 return optimal_T_found
            else:
                 logger.warning(f"Temperature optimization failed: {res.message}. Returning T=1.0")
                 self.optimal_T = 1.0
                 return 1.0
        except Exception as e:
             logger.error(f"Error during temperature optimization: {e}", exc_info=True)
             logger.warning("Returning T=1.0 due to optimization error.")
             self.optimal_T = 1.0
             return 1.0

    def calibrate(self, p_raw: np.ndarray, T: Optional[float] = None) -> np.ndarray:
        """
        Applies temperature scaling to raw probabilities.

        Args:
            p_raw (np.ndarray): Raw model probabilities.
            T (Optional[float]): Temperature value. If None, uses the stored optimal_T.
                                 Defaults to 1.0 if optimal_T is also None.

        Returns:
            np.ndarray: Calibrated probabilities.
        """
        temp = T if T is not None else self.optimal_T
        if temp is None:
            logger.warning("Temperature T not provided and not optimized yet. Using T=1.0 for calibration.")
            temp = 1.0
        if temp <= 0:
            logger.error(f"Invalid temperature T={temp}. Using T=1.0 instead.")
            temp = 1.0
            
        # Clip raw probabilities and compute logits
        eps = 1e-7
        p_clipped = np.clip(np.asarray(p_raw).astype(float), eps, 1 - eps)
        logit_p = logit(p_clipped)
        
        # Apply temperature scaling
        p_cal = expit(logit_p / temp)
        return p_cal

    def reliability_curve(
        self,
        p_pred: np.ndarray, # Expects raw OR calibrated probabilities
        y_true: np.ndarray,
        n_bins: int = 10,
        plot_title: str = "Reliability Curve",
        save_path: Optional[str] = None
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Computes and optionally plots the reliability curve.

        Args:
            p_pred (np.ndarray): Predicted probabilities (raw or calibrated).
            y_true (np.ndarray): True binary labels.
            n_bins (int): Number of bins for the curve.
            plot_title (str): Title for the plot.
            save_path (Optional[str]): If provided, saves the plot to this path.

        Returns:
            Tuple[np.ndarray, np.ndarray]: (bin_centers, empirical_prob)
        """
        p_pred = np.asarray(p_pred).flatten()
        y_true = np.asarray(y_true).flatten()
        
        bins = np.linspace(0, 1, n_bins + 1)
        # Handle potential edge cases with digitize for values exactly 1.0
        bin_ids = np.digitize(p_pred, bins[1:], right=True) # Bin index 0 to n_bins-1

        bin_centres = 0.5 * (bins[:-1] + bins[1:])
        empirical_prob = np.zeros(n_bins)
        bin_counts = np.zeros(n_bins)
        
        for i in range(n_bins):
            idx = bin_ids == i
            bin_counts[i] = np.sum(idx)
            if bin_counts[i] > 0:
                empirical_prob[i] = y_true[idx].mean()
            else:
                empirical_prob[i] = np.nan # Use NaN for empty bins

        if save_path:
            try:
                plt.figure(figsize=(6, 6))
                plt.plot([0, 1], [0, 1], "k--", label="Perfect Calibration")
                # Only plot bins with counts
                valid_bins = bin_counts > 0
                plt.plot(bin_centres[valid_bins], empirical_prob[valid_bins], "o-", label="Model")
                plt.xlabel("Mean Predicted Probability (per bin)")
                plt.ylabel("Fraction of Positives (per bin)")
                plt.title(plot_title)
                plt.legend()
                plt.grid(True)
                plt.tight_layout()
                plt.savefig(save_path)
                plt.close()
                logger.info(f"Reliability curve saved to {save_path}")
            except Exception as e:
                 logger.error(f"Failed to generate or save reliability plot: {e}", exc_info=True)

        return bin_centres, empirical_prob

    def action_signal(self, p_cal: np.ndarray) -> np.ndarray:
        """
        Generates trading signal (1, -1, or 0) based on calibrated probability 
        and the instance's edge threshold.

        Args:
            p_cal (np.ndarray): Calibrated probabilities P(up).

        Returns:
            np.ndarray: Action signals (1 for long, -1 for short, 0 for neutral).
        """
        p_cal = np.asarray(p_cal)
        # Signal long if P(up) > threshold
        go_long = p_cal > self.edge_threshold
        # Signal short if P(down) > threshold, which is P(up) < 1 - threshold
        go_short = p_cal < (1.0 - self.edge_threshold)
        
        # Assign signals: 1 for long, -1 for short, 0 otherwise
        signal = np.where(go_long, 1, np.where(go_short, -1, 0))
        return signal 