import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

from smartcal.calibration_algorithms.calibration_base import CalibratorBase
from smartcal.config.configuration_manager.configuration_manager import ConfigurationManager
from smartcal.config.enums.calibration_error_code_enum import CalibrationErrorCode
from smartcal.utils.timer import time_operation


config_manager = ConfigurationManager()

class MatrixScalingCalibrator(CalibratorBase):
    """
    Implements Matrix Scaling calibration.
    Applies a learned affine transformation ZW + b to logits before softmax.
    """

    def __init__(self, lr=config_manager.lr_matrixscaling, max_iter=config_manager.max_iter_matrixscaling,
                 seed=CalibratorBase.DEFAULT_SEED):
        super().__init__(seed)
        self.timing = {}
        self.lr = lr
        self.max_iter = max_iter
        self.W = None  # Weight matrix (C x C)
        self.b = None  # Bias vector (C)
        self.fitted = False

    @time_operation
    def fit(self, predictions, ground_truth):
        """
        Fits the matrix scaling model using LBFGS optimization.

        Args:
            predictions (np.ndarray): Model logits of shape (n_samples, n_classes).
            ground_truth (np.ndarray): True labels of shape (n_samples,).
        """
        self.set_random_seed(self.seed)
        predictions, ground_truth = self.validate_inputs(predictions, ground_truth)
        n_classes = predictions.shape[1]

        predictions = torch.tensor(predictions, dtype=torch.float32)
        ground_truth = torch.tensor(ground_truth, dtype=torch.long)

        # Initialize W as identity and b as zero
        self.W = nn.Parameter(torch.eye(n_classes, dtype=torch.float32))
        self.b = nn.Parameter(torch.zeros(n_classes, dtype=torch.float32))

        optimizer = optim.LBFGS([self.W, self.b], lr=self.lr, max_iter=self.max_iter)
        criterion = nn.CrossEntropyLoss()

        def closure():
            optimizer.zero_grad()
            scaled_logits = torch.matmul(predictions, self.W) + self.b
            loss = criterion(torch.nn.functional.log_softmax(scaled_logits, dim=1), ground_truth)
            loss.backward()
            return loss

        optimizer.step(closure)
        self.fitted = True

        # Store optimized values
        self.metadata["params"].update({
            "optimized_W": self.W.detach().numpy().tolist(),
            "optimized_b": self.b.detach().numpy().tolist()
        })
        self.metadata["dataset_info"].update({
            "n_samples": predictions.shape[0],
            "n_classes": predictions.shape[1],
        })

    @time_operation
    def predict(self, test_data):
        """
        Applies matrix scaling to test logits.

        Args:
            test_data (np.ndarray): Logits of shape (n_samples, n_classes).

        Returns:
            np.ndarray: Calibrated probabilities of shape (n_samples, n_classes).
        """
        if not self.fitted:
            raise RuntimeError("Model must be fitted before prediction.")

        if not isinstance(test_data, np.ndarray):
            raise ValueError(f"Expected input to be a NumPy array, but got {type(test_data)}.")

        if test_data.ndim != 2:
            raise ValueError(f"Expected input shape (N, C), but got {test_data.shape}.")

        if test_data.shape[1] != self.W.shape[0]:
            raise ValueError(f"Expected input shape (N, {self.W.shape[0]}), but got {test_data.shape}.")

        test_data = torch.tensor(test_data, dtype=torch.float32)
        scaled_logits = torch.matmul(test_data, self.W) + self.b
        probabilities = torch.nn.functional.softmax(scaled_logits, dim=1)

        if torch.isnan(probabilities).any():
            raise ValueError((CalibrationErrorCode.NANS_DETECTED, "Predicted probabilities contain NaN values. Try different parameters or initialization."))

        return probabilities.detach().numpy()
    
    def get_timing(self):
        return self.timing