"""
Vector Scaling Calibration for Multi-Class Classifiers.

Ref: revisions.txt Section 4
Based on: https://arxiv.org/abs/1706.04599 (On Calibration of Modern Neural Networks)
"""

import numpy as np
import tensorflow as tf
from scipy.optimize import minimize
import logging
from typing import Optional, Tuple, List, Any
import pandas as pd

logger = logging.getLogger(__name__)

class VectorCalibrator:
    """
    Implements Vector Scaling calibration.
    
    Finds a diagonal matrix W and a vector b such that the calibrated
    probabilities p_cal = softmax(W * z + b) minimize the NLL, where z
    are the pre-softmax logits.
    For K classes, this involves optimizing 2*K parameters.
    """
    
    def __init__(self):
        """Initialize the calibrator."""
        self.W = None # Diagonal weight matrix (represented as a vector)
        self.b = None # Bias vector
        self.optimal_params = None # Store concatenated [W_diag, b]
        
    def _softmax(self, x: np.ndarray) -> np.ndarray:
        """Numerically stable softmax."""
        e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
        return e_x / np.sum(e_x, axis=-1, keepdims=True)
        
    def _nll_loss(self, params: np.ndarray, logits: np.ndarray, y_onehot: np.ndarray) -> float:
        """
        Negative Log-Likelihood loss function (without regularization).
        
        Args:
            params (np.ndarray): Concatenated vector [W_diag, b].
            logits (np.ndarray): Raw output logits from the model (shape: N x K).
            y_onehot (np.ndarray): One-hot encoded true labels (shape: N x K).
        
        Returns:
            float: The calculated NLL loss.
        """
        num_classes = logits.shape[1]
        if len(params) != 2 * num_classes:
            raise ValueError(f"Expected {2*num_classes} params, got {len(params)}")
            
        W_diag = params[:num_classes]
        b = params[num_classes:]
        
        # Apply scaling: W is diagonal, so element-wise multiplication works
        scaled_logits = logits * W_diag + b # Broadcasting W_diag and b
        
        # Calculate probabilities using softmax
        calibrated_probs = self._softmax(scaled_logits)
        
        # Avoid log(0) - clip probabilities
        eps = 1e-12
        calibrated_probs = np.clip(calibrated_probs, eps, 1.0 - eps)
        
        # Calculate NLL
        nll = -np.sum(y_onehot * np.log(calibrated_probs), axis=1)
        return np.mean(nll)

    def _nll_loss_regularized(self, params: np.ndarray, logits: np.ndarray, y_onehot: np.ndarray, reg_lambda: float) -> float:
        """NLL loss with L2 regularization on W and b (deviation from W=1, b=0)."""
        base_nll = self._nll_loss(params, logits, y_onehot)
        
        num_classes = logits.shape[1]
        W_diag = params[:num_classes]
        b = params[num_classes:]
        
        # L2 penalty: encourage W towards 1 and b towards 0
        l2_penalty_W = np.sum((W_diag - 1.0) ** 2)
        l2_penalty_b = np.sum(b ** 2)
        l2_penalty = reg_lambda * (l2_penalty_W + l2_penalty_b)
        
        return base_nll + l2_penalty

    def fit(self, logits: np.ndarray, y_onehot: np.ndarray, reg_lambda: float = 0.0) -> None:
        """
        Finds the optimal scaling parameters W (diagonal) and b, optionally with L2 regularization.

        Args:
            logits (np.ndarray): Raw output logits from the model (shape: N x K).
            y_onehot (np.ndarray): One-hot encoded true labels (shape: N x K).
            reg_lambda (float): L2 regularization strength. Defaults to 0.0.
        """
        if logits.shape[0] != y_onehot.shape[0]:
            raise ValueError("Logits and labels must have the same number of samples.")
        if len(logits.shape) != 2 or len(y_onehot.shape) != 2:
            raise ValueError("Logits and labels must be 2D arrays (N x K).")
        if logits.shape[1] != y_onehot.shape[1]:
            raise ValueError("Logits and labels must have the same number of classes.")
            
        num_classes = logits.shape[1]
        logger.info(f"Fitting Vector Scaling for {num_classes} classes (L2 lambda={reg_lambda})...")
        
        # Initial guess: W = identity (diag=1), b = zero vector
        initial_params = np.concatenate([np.ones(num_classes), np.zeros(num_classes)])
        
        # Define bounds (optional, but can help stability)
        # Allow W > 0, b can be anything
        bounds = [(1e-6, None)] * num_classes + [(None, None)] * num_classes
        
        # Select objective function based on regularization
        if reg_lambda == 0.0:
            objective_func = self._nll_loss
            extra_args = (logits, y_onehot)
        else:
            objective_func = self._nll_loss_regularized
            extra_args = (logits, y_onehot, reg_lambda)

        # Minimize the NLL loss
        # Using L-BFGS-B as it handles bounds well
        result = minimize(
            objective_func, 
            initial_params, 
            args=extra_args, # Pass logits, y_onehot (and lambda if needed)
            method='L-BFGS-B',
            bounds=bounds, # Use bounds
            options={'maxiter': 1000, 'ftol': 1e-8} # Example options
        )
        
        if result.success:
            self.optimal_params = result.x
            self.W = self.optimal_params[:num_classes]
            self.b = self.optimal_params[num_classes:]
            logger.info(f"Vector Scaling fit successful. Optimal NLL: {result.fun:.4f}")
            logger.info(f"  Optimal W (diag): {np.round(self.W, 3)}")
            logger.info(f"  Optimal b: {np.round(self.b, 3)}")
        else:
            logger.error(f"Vector Scaling optimization failed: {result.message}")
            # Handle failure: maybe use initial params or raise error?
            self.optimal_params = initial_params # Fallback to initial
            self.W = self.optimal_params[:num_classes]
            self.b = self.optimal_params[num_classes:]
            logger.warning("Using initial parameters (W=I, b=0) due to optimization failure.")

    def calibrate(self, logits: np.ndarray) -> np.ndarray:
        """
        Applies the learned scaling parameters to new logits.

        Args:
            logits (np.ndarray): Raw logits from the model (shape: N x K).

        Returns:
            np.ndarray: Calibrated probabilities (shape: N x K).
                      Returns uncalibrated softmax if fit() wasn't called or failed.
        """
        if self.W is None or self.b is None:
            logger.warning("Vector Scaling parameters not fitted. Returning uncalibrated softmax.")
            return self._softmax(logits)
            
        if logits.shape[1] != len(self.W):
            raise ValueError(f"Input logits have {logits.shape[1]} classes, but calibrator was fitted for {len(self.W)} classes.")
            
        scaled_logits = logits * self.W + self.b
        calibrated_probs = self._softmax(scaled_logits)
        return calibrated_probs

    def save_params(self, filepath: str) -> None:
        """Saves the optimal parameters (W_diag and b) to a .npy file."""
        if self.optimal_params is None:
            logger.error("No parameters to save. Call fit() first.")
            return
        try:
            np.save(filepath, self.optimal_params)
            logger.info(f"Vector Scaling parameters saved to {filepath}")
        except Exception as e:
            logger.error(f"Failed to save parameters to {filepath}: {e}")

    def load_params(self, filepath: str) -> bool:
        """Loads the optimal parameters from a .npy file."""
        try:
            params = np.load(filepath)
            num_params = len(params)
            if num_params % 2 != 0:
                 raise ValueError(f"Loaded params have odd length ({num_params}), expected 2*K.")
            num_classes = num_params // 2
            self.optimal_params = params
            self.W = self.optimal_params[:num_classes]
            self.b = self.optimal_params[num_classes:]
            logger.info(f"Vector Scaling parameters loaded successfully from {filepath} ({num_classes} classes).")
            return True
        except FileNotFoundError:
            logger.error(f"Parameter file not found: {filepath}")
            return False
        except Exception as e:
            logger.error(f"Failed to load parameters from {filepath}: {e}")
            # Reset parameters on load failure?
            self.W = None
            self.b = None
            self.optimal_params = None
            return False

    def optimize_edge_threshold_ternary(self, p_cal: np.ndarray, y_onehot: np.ndarray, 
                                        positive_class_indices: List[int] = [2], # Default: Class 2 (Up) is positive
                                        n_thresholds: int = 101) -> float:
        """
        Finds an optimal edge threshold for ternary classification using Youden's J.
        This simplifies the ternary problem to binary: Positive Class(es) vs Others.

        Args:
            p_cal (np.ndarray): Calibrated probabilities (N x K, e.g., N x 3).
            y_onehot (np.ndarray): True one-hot labels (N x K).
            positive_class_indices (List[int]): List of indices considered 'positive' (e.g., [2] for 'Up').
            n_thresholds (int): Number of thresholds to test for the summed positive class probability.

        Returns:
            float: The threshold on the summed positive class probability that maximizes Youden's J.
                   Defaults to 0.5 / len(positive_class_indices) if calculation fails.
        """
        logger.info(f"Optimizing ternary edge threshold (Positive classes: {positive_class_indices})...")
        if p_cal.shape[1] != y_onehot.shape[1]:
            raise ValueError("Calibrated probs and one-hot labels must have same number of classes.")
        if not positive_class_indices or any(i >= p_cal.shape[1] for i in positive_class_indices):
            raise ValueError(f"Invalid positive_class_indices: {positive_class_indices}")

        # Combine probabilities for positive classes
        p_positive_sum = np.sum(p_cal[:, positive_class_indices], axis=1)

        # Create binary true labels: 1 if true class is in positive_class_indices, 0 otherwise
        y_true_binary = np.sum(y_onehot[:, positive_class_indices], axis=1).astype(int)

        # Reuse the binary threshold optimization logic
        # Need an instance of the base Calibrator or copy the method here
        # For simplicity, let's assume a temporary base calibrator or call a static helper
        # This assumes a base Calibrator class exists with optimize_edge_threshold
        try:
            # Use a generic function or temporary instance approach if needed
            # Here, assuming we can call optimize_edge_threshold directly (or copy its logic)
            # Copied logic approach for self-containment:
            thresholds = np.linspace(0, 1, n_thresholds)
            best_j = -1
            num_pos_classes = len(positive_class_indices)
            default_threshold = 0.5 / num_pos_classes if num_pos_classes > 0 else 0.5 # Adjust default
            best_threshold = default_threshold

            for thr in thresholds:
                y_pred_binary = (p_positive_sum >= thr).astype(int)
                try:
                    tn, fp, fn, tp = confusion_matrix(y_true_binary, y_pred_binary, labels=[0, 1]).ravel()
                    tpr = tp / (tp + fn) if (tp + fn) > 0 else 0
                    fpr = fp / (fp + tn) if (fp + tn) > 0 else 0
                    j_stat = tpr - fpr
                    if j_stat > best_j:
                        best_j = j_stat
                        best_threshold = thr
                except ValueError:
                    continue # Skip thresholds where confusion matrix fails
                except Exception as e:
                    logger.error(f"Error calculating Youden's J for threshold {thr:.3f}: {e}")
                    continue
            
            if best_j <= 0:
                logger.warning(f"Could not find ternary threshold yielding positive Youden's J (Max J = {best_j:.3f}). Defaulting to {default_threshold:.3f}.")
                best_threshold = default_threshold
            else:
                logger.info(f"Optimal ternary edge threshold found: {best_threshold:.3f} (Max Youden's J = {best_j:.3f})")
            
            return best_threshold

        except Exception as e:
            logger.error(f"Error during ternary edge threshold optimization: {e}", exc_info=True)
            return default_threshold # Return default on error

    def fit_rolling(self, logits: np.ndarray, y_onehot: np.ndarray, every_n: int, window_m: int, indices: Optional[pd.Index] = None) -> List[Tuple[Any, np.ndarray]]:
        """
        Fits Vector Scaling parameters on rolling windows.

        Args:
            logits (np.ndarray): Raw model logits.
            y_onehot (np.ndarray): True one-hot labels.
            every_n (int): How often to refit the parameters (e.g., 5000 steps).
            window_m (int): The size of the rolling window for refitting (e.g., 20000 steps).
            indices (Optional[pd.Index]): Optional original indices corresponding to data.

        Returns:
            List[Tuple[Any, np.ndarray]]: A schedule of (index, optimal_params) tuples.
                                          optimal_params = [W_diag, b].
        """
        logger.info(f"Performing rolling vector scaling calibration: every {every_n} steps, window size {window_m}")
        n_samples = len(logits)
        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
        num_classes = logits.shape[1]
        initial_params = np.concatenate([np.ones(num_classes), np.zeros(num_classes)]) # Default params

        for i in range(window_m - 1, n_samples, every_n):
            start_idx = max(0, i - window_m + 1)
            end_idx = i + 1

            window_logits = logits[start_idx:end_idx]
            window_y = y_onehot[start_idx:end_idx]
            
            logger.debug(f"Rolling fit at step {i} (Window: [{start_idx}, {end_idx}))")
            
            # Fit parameters for this window
            # Use a temporary instance or call a static fit method? 
            # Re-using self._nll_loss and minimize directly seems cleaner here
            bounds = [(1e-6, None)] * num_classes + [(None, None)] * num_classes
            result = minimize(
                self._nll_loss, 
                initial_params, # Start from scratch or use previous params?
                args=(window_logits, window_y),
                method='L-BFGS-B',
                bounds=bounds,
                options={'maxiter': 500, 'ftol': 1e-7} # Faster options for rolling?
            )
            
            optimal_params_window = initial_params # Default if fit fails
            if result.success:
                optimal_params_window = result.x
            else:
                logger.warning(f"Rolling fit failed at step {i}: {result.message}. Using default params for this step.")

            schedule_idx = indices[i] if indices is not None and i < len(indices) else i
            schedule.append((schedule_idx, optimal_params_window))
            last_fit_idx = i 
            
            # Update the main instance parameters to the latest successful fit?
            if result.success:
                self.optimal_params = optimal_params_window
                self.W = self.optimal_params[:num_classes]
                self.b = self.optimal_params[num_classes:]

        logger.info(f"Generated rolling vector calibration schedule with {len(schedule)} entries.")
        if not schedule:
             logger.warning("Rolling vector calibration resulted in an empty schedule.")
             # Optionally fit once on the full data as fallback

        return schedule

    # --- Add Reliability Curve Plotting --- #
    def reliability_curve(
        self,
        probs: np.ndarray, # Calibrated probabilities (N, K)
        y_true: np.ndarray, # True labels (N,) or one-hot (N, K)
        n_bins: int = 10,
        plot_title: str = "Multi-Class Reliability Curve",
        save_path: Optional[str] = None
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Computes and optionally plots the reliability curve for multi-class classification.
        Uses the maximum predicted probability as confidence.

        Args:
            probs (np.ndarray): Calibrated probabilities (shape: N x K).
            y_true (np.ndarray): True labels (class index, shape: N) or one-hot (N, K).
            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, np.ndarray]: (bin_centers, accuracy_in_bin, avg_confidence_in_bin)
        """
        if len(probs.shape) != 2:
            raise ValueError("Input probabilities must be 2D (N x K).")
        
        n_samples, num_classes = probs.shape
        
        # Ensure y_true is class index (N,)
        if len(y_true.shape) == 2 and y_true.shape[1] == num_classes:
            y_true_idx = np.argmax(y_true, axis=1)
        elif len(y_true.shape) == 1:
            y_true_idx = np.asarray(y_true).astype(int)
        else:
            raise ValueError("y_true must be 1D class indices or 2D one-hot encoded.")

        if len(y_true_idx) != n_samples:
            raise ValueError("Number of samples mismatch between probs and y_true.")
            
        # Get confidence (max probability) and predicted class
        confidences = np.max(probs, axis=1)
        predictions = np.argmax(probs, axis=1)
        correctness = (predictions == y_true_idx).astype(float)

        bins = np.linspace(0, 1, n_bins + 1)
        bin_centers = 0.5 * (bins[:-1] + bins[1:])
        bin_ids = np.digitize(confidences, bins[1:], right=False)
        bin_ids = np.clip(bin_ids, 0, n_bins - 1)

        accuracy_in_bin = np.zeros(n_bins) * np.nan
        avg_confidence_in_bin = np.zeros(n_bins) * np.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:
                accuracy_in_bin[i] = np.mean(correctness[idx])
                avg_confidence_in_bin[i] = np.mean(confidences[idx])
                
        # Filter out bins with no samples for plotting
        valid_mask = bin_counts > 0
        plot_centers = bin_centers[valid_mask]
        plot_accuracy = accuracy_in_bin[valid_mask]
        plot_confidence = avg_confidence_in_bin[valid_mask]
        
        # Calculate ECE
        ece = np.sum(np.abs(plot_accuracy - plot_confidence) * (bin_counts[valid_mask] / n_samples))
        plot_title += f" (ECE = {ece:.3f})"

        if save_path:
            # Reuse plotting logic similar to binary Calibrator
            try:
                import matplotlib.pyplot as plt # Local import for plotting
                fig, ax = plt.subplots(1, 1, figsize=(6, 6))
                ax.plot([0, 1], [0, 1], "k--", label="Perfect Calibration")
                ax.plot(plot_confidence, plot_accuracy, "o-", label="Model Calibration") # Plot Acc vs Conf
                
                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)
                # Adjust legend position slightly for multi-class
                fig.legend(loc="upper left", bbox_to_anchor=(0.1, 0.9))
                
                ax.set_xlabel("Average Confidence (Max Probability per bin)")
                ax.set_ylabel("Accuracy (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"Multi-class reliability curve saved to {save_path}")
            except ImportError:
                logger.error("Matplotlib not found. Cannot generate reliability plot.")
            except Exception as e:
                 logger.error(f"Failed to generate or save multi-class reliability plot: {e}", exc_info=True)

        return bin_centers, accuracy_in_bin, avg_confidence_in_bin
    # --- End Reliability Curve Plotting --- #

# --- Set Availability Flag --- #
# Define this at the module level after the class definition
VECTOR_CALIBRATOR_AVAILABLE = True
logger.info("VectorCalibrator class defined and available.")
# --- End Set Flag --- # 