"""
Calibration Component for GRU Model Probabilities.

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

import pandas as pd # Added for type hinting
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, List, Dict, Any
import logging
import os
from sklearn.metrics import confusion_matrix # Added for Youden's J

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 _nll_objective_regularized(self, T: float, logit_p: np.ndarray, y_true: np.ndarray, reg_lambda: float) -> float:
        """NLL objective with L2 regularization on T (deviation from 1)."""
        base_nll = self._nll_objective(T, logit_p, y_true)
        # Regularize T towards 1 (no scaling)
        l2_penalty = reg_lambda * ((T - 1.0) ** 2)
        return base_nll + l2_penalty

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

        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.
            reg_lambda (float): L2 regularization strength (penalty on (T-1)^2). Defaults to 0.0 (no reg).

        Returns:
            float: The optimal temperature found.
        """
        logger.info(f"Optimizing calibration temperature using {len(p_raw)} samples (L2 lambda={reg_lambda})...")
        # 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:
            # Select objective function based on regularization
            objective_func = self._nll_objective if reg_lambda == 0.0 else \
                             lambda T: self._nll_objective_regularized(T, logit_p, y_true, reg_lambda)

            res = minimize_scalar(
                objective_func,
                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 for binary classification.

        Args:
            p_pred (np.ndarray): Predicted probabilities for the positive class.
            y_true (np.ndarray): True binary labels (0 or 1).
            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()
        
        if not np.all((y_true == 0) | (y_true == 1)):
            # Handle potential soft labels by converting to hard labels for accuracy calculation
            logger.debug("Non-binary values detected in y_true for reliability curve. Converting > 0.5 to 1, else 0.")
            y_true = (y_true > 0.5).astype(int)

        bins = np.linspace(0, 1, n_bins + 1)
        bin_centers = 0.5 * (bins[:-1] + bins[1:])
        # Handle potential edge cases with digitize for values exactly 1.0
        # Ensure indices are within [0, n_bins-1]
        bin_ids = np.digitize(p_pred, bins[1:], right=False) 
        bin_ids = np.clip(bin_ids, 0, n_bins - 1) # Clip to handle edge case p_pred = 0

        empirical_prob = np.zeros(n_bins) * np.nan # Default to NaN
        avg_confidence = np.zeros(n_bins) * np.nan # Default to NaN
        bin_counts = np.zeros(n_bins, dtype=int)
        
        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()
                avg_confidence[i] = p_pred[idx].mean()

        # Filter out bins with no samples for plotting
        valid_mask = bin_counts > 0
        plot_centers = bin_centers[valid_mask]
        plot_probs = empirical_prob[valid_mask]

        # Calculate ECE (Expected Calibration Error)
        ece = np.sum(np.abs(empirical_prob[valid_mask] - avg_confidence[valid_mask]) * (bin_counts[valid_mask] / len(p_pred)))
        plot_title += f" (ECE = {ece:.3f})"

        if save_path:
            try:
                fig, ax = plt.subplots(1, 1, figsize=(6, 6))
                ax.plot([0, 1], [0, 1], "k--", label="Perfect Calibration")
                # Only plot bins with counts
                ax.plot(plot_centers, plot_probs, "o-", label="Model Calibration")
                
                # Add bar chart for confidence distribution underneath
                ax2 = ax.twinx()
                ax2.bar(bin_centers, bin_counts, width=(bins[1]-bins[0])*0.9, alpha=0.2, color='grey', label='Bin Counts')
                ax2.set_ylabel("Count per Bin", color='grey')
                ax2.tick_params(axis='y', labelcolor='grey')
                ax2.set_ylim(bottom=0)
                fig.legend(loc="upper left", bbox_to_anchor=(0.1, 0.9))
                
                ax.set_xlabel("Mean Predicted Probability (per bin)")
                ax.set_ylabel("Fraction of Positives (per bin)")
                ax.set_title(plot_title)
                ax.grid(True, alpha=0.5)
                ax.set_xlim([0, 1])
                ax.set_ylim([0, 1])
                
                plt.tight_layout()
                plt.savefig(save_path)
                plt.close(fig)
                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_centers, 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 

    def optimize_edge_threshold(self, p_cal: np.ndarray, y_true: np.ndarray, n_thresholds: int = 101) -> float:
        """
        Finds the optimal binary edge threshold by maximizing Youden's J statistic.

        Args:
            p_cal (np.ndarray): Calibrated probabilities for the positive class (validation set).
            y_true (np.ndarray): True binary labels (validation set).
            n_thresholds (int): Number of thresholds to test between 0 and 1.

        Returns:
            float: The threshold that maximizes Youden's J (TPR - FPR).
                   Defaults to 0.5 if calculation fails or yields no improvement.
        """
        logger.info(f"Optimizing edge threshold using {len(p_cal)} validation samples...")
        p_cal = np.asarray(p_cal).flatten()
        y_true = np.asarray(y_true).flatten()
        
        # Handle potential soft labels for threshold optimization
        if not np.all((y_true == 0) | (y_true == 1)):
            logger.debug("Non-binary values detected in y_true for threshold optimization. Converting > 0.5 to 1, else 0.")
            y_true = (y_true > 0.5).astype(int)

        thresholds = np.linspace(0, 1, n_thresholds)
        best_j = -1
        best_threshold = 0.5 # Default

        for thr in thresholds:
            # Predictions based on threshold
            y_pred = (p_cal >= thr).astype(int)
            
            try:
                # Calculate confusion matrix: [[TN, FP], [FN, TP]]
                tn, fp, fn, tp = confusion_matrix(y_true, y_pred, labels=[0, 1]).ravel()
                
                # Calculate TPR (Sensitivity) and FPR (1 - Specificity)
                tpr = tp / (tp + fn) if (tp + fn) > 0 else 0
                fpr = fp / (fp + tn) if (fp + tn) > 0 else 0
                
                # Youden's J statistic
                j_stat = tpr - fpr
                
                if j_stat > best_j:
                    best_j = j_stat
                    best_threshold = thr
                    
            except ValueError as e:
                # Handle cases where confusion_matrix might fail (e.g., only one class predicted/present)
                logger.debug(f"Could not calculate confusion matrix for threshold {thr:.3f}: {e}. Skipping.")
                continue
            except Exception as e:
                logger.error(f"Unexpected error calculating Youden's J for threshold {thr:.3f}: {e}. Skipping.", exc_info=True)
                continue

        if best_j <= 0:
            logger.warning(f"Could not find threshold yielding positive Youden's J (Max J = {best_j:.3f}). Defaulting to 0.5.")
            best_threshold = 0.5
        else:
            logger.info(f"Optimal edge threshold found: {best_threshold:.3f} (Max Youden's J = {best_j:.3f})")
            
        # Update the calibrator's internal threshold if desired?
        # self.edge_threshold = best_threshold # Or leave it to the pipeline to manage
        
        return best_threshold

    def fit_rolling(self, p_raw: np.ndarray, y_true: np.ndarray, every_n: int, window_m: int, indices: Optional[pd.Index] = None) -> List[Tuple[Any, float]]:
        """
        Fits temperature scaling parameters on rolling windows.

        Args:
            p_raw (np.ndarray): Raw model probabilities.
            y_true (np.ndarray): True binary labels.
            every_n (int): How often to refit the temperature (e.g., 5000 steps).
            window_m (int): The size of the rolling window to use for refitting (e.g., 20000 steps).
            indices (Optional[pd.Index]): Optional original indices corresponding to p_raw/y_true.
                                          If provided, the output schedule uses these indices.

        Returns:
            List[Tuple[Any, float]]: A schedule of (index, optimal_T) tuples.
                                    The index is the step/timestamp where the new T applies.
                                    Returns empty list if inputs are too short or invalid.
        """
        logger.info(f"Performing rolling temperature calibration: every {every_n} steps, window size {window_m}")
        n_samples = len(p_raw)
        if n_samples < window_m or n_samples < every_n or window_m <= 0 or every_n <= 0:
            logger.warning(f"Insufficient samples ({n_samples}) or invalid params (every={every_n}, window={window_m}) for rolling calibration. Skipping.")
            return []

        schedule = []
        last_fit_idx = -1 # Track the index of the last fit to avoid redundant calculations

        for i in range(window_m -1, n_samples, every_n):
            # Define the window for this fit
            start_idx = max(0, i - window_m + 1)
            end_idx = i + 1 # Slice is exclusive at the end
            
            # Check if this window significantly overlaps with the last one processed
            # This is a simple check; more sophisticated checks could be added
            # if start_idx <= last_fit_idx:
            #     continue 

            window_p = p_raw[start_idx:end_idx]
            window_y = y_true[start_idx:end_idx]
            
            # Use the standard optimise_temperature method on the window
            logger.debug(f"Rolling fit at step {i} (Window: [{start_idx}, {end_idx}))")
            optimal_t_window = self.optimise_temperature(window_p, window_y)
            
            # Determine the index/timestamp for this schedule entry
            schedule_idx = indices[i] if indices is not None and i < len(indices) else i
            
            schedule.append((schedule_idx, optimal_t_window))
            last_fit_idx = i 
            
            # Store the latest T as the primary optimal_T for the instance?
            self.optimal_T = optimal_t_window 

        logger.info(f"Generated rolling calibration schedule with {len(schedule)} entries.")
        if not schedule:
             logger.warning("Rolling calibration resulted in an empty schedule.")
             # Consider fitting once on the whole data as fallback?
             # optimal_t_full = self.optimise_temperature(p_raw, y_true)
             # schedule_idx_full = indices[-1] if indices is not None else n_samples - 1
             # schedule.append((schedule_idx_full, optimal_t_full))
             # self.optimal_T = optimal_t_full

        return schedule 