import torch
import torch.nn as nn
import numpy as np
from sklearn.preprocessing import StandardScaler

from contextlib import contextmanager
import logging
from typing import Callable, Dict, List, Optional, Union

from .utils import fixpoint_iteration

logger = logging.getLogger(__name__)

__activation_types__ = (
    nn.CELU,
    nn.ELU,
    nn.GELU,
    nn.GLU,
    nn.Hardshrink,
    nn.Hardsigmoid,
    nn.Hardswish,
    nn.Hardtanh,
    nn.LeakyReLU,
    nn.LogSigmoid,
    nn.LogSoftmax,
    nn.Mish,
    nn.MultiheadAttention,
    nn.PReLU,
    nn.ReLU,
    nn.ReLU6,
    nn.RReLU,
    nn.SELU,
    nn.Sigmoid,
    nn.SiLU,
    nn.Softmax,
    nn.Softmax2d,
    nn.Softmin,
    nn.Softplus,
    nn.Softshrink,
    nn.Softsign,
    nn.Tanh,
    nn.Tanhshrink,
    nn.Threshold,
)

class HookManager:
    def __init__(self):
        """
        HookManager class to collect pre- and post-activations from a model.
        """
        self.pre_activations = []
        self.post_activations = []
        self.hooks = []

    def hook_fn(self, module : torch.nn.Module, input : torch.Tensor, output : torch.Tensor):
        """
        Hook function to collect pre- and post-activations from a model.
        """
        if input[0].size() == output.size():
            self.pre_activations.append(input[0].detach())
            self.post_activations.append(output.detach())

    def _apply_all_hooks(self, module : torch.nn.Module):
        """
        Recursively apply hooks to all activation layers in the model.
        """
        children = list(module.children())
        if len(children) == 0:  # It's a leaf module
            if isinstance(module, __activation_types__): # check if the module is an activation layer
                hook = module.register_forward_hook(self.hook_fn)
                self.hooks.append(hook)
        else:
            for child in module.children():
                self._apply_all_hooks(child)  # Recursively apply to children
    
    def _check_modules(self, module : torch.nn.Module, list_relu : List[torch.nn.Module]):
        """
        Recursively look for activation layers inside the module.
        """
        children = list(module.children())
        if len(children) == 0:  # It's a leaf module
            if isinstance(module, __activation_types__):
                list_relu.append(module)
        else:
            for child in module.children():
                self._check_modules(child, list_relu)  # Recursively apply to children

    def _apply_n_hooks(self, module : torch.nn.Module, n : int = 3):
        """
        Skip some activations and apply hooks only to the last activation layer of each layer block.

        Args:
            module (torch.nn.Module): The module to apply hooks to.
            n (int): The threshold to define a block of layers. Only skip layers when the number of children is greater than n.
        """
        children_length = len(list(module.children()))
        if children_length == 0:
            if isinstance(module, __activation_types__):
                hook = module.register_forward_hook(self.hook_fn)
                self.hooks.append(hook)
        elif children_length <= n: # if the number of children is less than n, apply hooks to all children
            for child in module.children():
                self._apply_n_hooks(child, n)
        else: # if the number of children is greater than n, apply hooks to the last activation of the layer block
            for child in module.children():
                relu_layers = []
                self._check_modules(child, relu_layers)
                if relu_layers:
                    final_relu = relu_layers[-1] # get only the last activation of a block of layers
                    hook = final_relu.register_forward_hook(self.hook_fn)
                    self.hooks.append(hook)

    @contextmanager
    def register_hooks(self, model : torch.nn.Module, skip_layers : Optional[int] = None):
        """
        Register hooks to all activation layers in the model.
        """
        if skip_layers is None:
            self._apply_all_hooks(model)
        else:
            self._apply_n_hooks(model, skip_layers)
        try:
            yield
        finally:
            for hook in self.hooks:
                hook.remove()
            self.pre_activations.clear()
            self.post_activations.clear()
            self.hooks.clear()


class SIM():
    r"""
    SIM base class.

    Args:
        activation_fn (Callable): Activation function used in the implicit model.
        kappa (float): Parameter to ensure convergence of the fixed-point iteration. Default is 0.99.
        atol (float): Absolute tolerance for the fixed-point iteration. Default is 1e-6.
        skip_layers (int, optional): If not None, only the last activation of each layer block will be used. The block size is controlled by the parameter. Defaults to None.
        standardize (bool, optional): Whether to standardize the input data using scipy StandardScaler. Defaults to False.
        device (str or torch.device, optional): Device to use for SIM. Defaults to "cpu".
        dtype (str or torch.dtype, optional): Data type for SIM. Defaults to torch.float32.
    """
    def __init__(
        self,
        activation_fn : Callable = nn.functional.relu,
        kappa : float = 0.99,
        atol : float = 1e-6,
        skip_layers : Optional[int] = None,
        standardize : bool = False,
        device : Optional[Union[str, torch.device]] = "cpu", 
        dtype : Optional[Union[str, torch.dtype]] = torch.float32,
    ) -> None:
        self.activation_fn = activation_fn
        self.atol = atol
        self.kappa = kappa
        self.skip_layers = skip_layers
        self.device = device
        self.dtype = dtype
        self.weights = {
            'A': None,
            'B': None,
            'C': None,
            'D': None,
        }
        self.standardize = standardize
        if standardize:
            self.scaler = StandardScaler()

    def to(
        self,
        device : Optional[Union[str, torch.device]] = None, 
        dtype : Optional[Union[str, torch.dtype]] = None
    ) -> None:
        if device is not None:
            self.device = device
        if dtype is not None:
            self.dtype = dtype
        for weight in self.weights.keys():
            if self.weights[weight] is not None:
                self.weights[weight] = self.weights[weight].to(self.device, self.dtype)
        return self
    
    def __call__(
        self, 
        input : torch.Tensor,
    ) -> torch.Tensor:
        """
        Forward pass of a standard implicit model.

        Args:
            input (torch.Tensor): Input data with shape (batch_size, input_dim).
        
        Returns:
            output (torch.Tensor): Output data with shape (batch_size, output_dim).
        """
        for weight in self.weights.keys():
            assert self.weights[weight] is not None, f"Weight matrix {weight} is not trained"

        X = fixpoint_iteration(self.weights['A'], self.weights['B'], input, self.activation_fn, self.device, atol=self.atol)
        Y = X @ self.weights['C'] + self.weights['D'] @ input

        return Y
    
    def get_states(
        self, 
        model : torch.nn.Module, 
        dataloader : torch.utils.data.DataLoader,
    ) -> Dict[str, np.ndarray]:
        """
        Extract the states data (pre-activations, post-activations, inputs, outputs) from the explicit model.
        The dataloader should only load a small amount of data to avoid memory issues.

        Args:
            model (torch.nn.Module): Explicit model (teacher) to extract state data.
            dataloader (torch.utils.data.DataLoader): Data loader.
        
        Returns:
            states_data (dict): Dictionary containing the states data:
                - U (np.ndarray): Input data with shape (batch_size, input_dim).
                - Z (np.ndarray): Pre-activations with shape (batch_size, hidden_dim).
                - X (np.ndarray): Post-activations with shape (batch_size, hidden_dim).
                - Y (np.ndarray): Output data with shape (batch_size, output_dim).
        """
        model.requires_grad_(False)
        model.eval()
        model.to(self.device)
        
        hooks = HookManager()

        outputs_accumulated = []
        inputs_accumulated = []
        pre_activations_accumulated = []
        post_activations_accumulated = []

        for input_samples, _ in dataloader:
            input_samples = input_samples.to(self.device)

            # Run the model with hooks and no gradient calculations
            with hooks.register_hooks(model, skip_layers=self.skip_layers), torch.no_grad():
                outputs = model(input_samples)

                # Accumulate outputs
                outputs_accumulated.append(
                    outputs.cpu().numpy()
                )  # Convert to NumPy array and store

                # Construct and accumulate implicit representation
                U = input_samples.flatten(1).cpu().numpy()
                inputs_accumulated.append(U)

                # Accumulate activations
                num_layers = len(hooks.pre_activations)

                Z = np.hstack(
                    [
                        hooks.pre_activations[i].flatten(1).cpu().numpy()
                        for i in range(num_layers - 1, -1, -1)
                    ]
                )
                pre_activations_accumulated.append(Z)

                X = np.hstack(
                    [
                        hooks.post_activations[i].flatten(1).cpu().numpy()
                        for i in range(num_layers - 1, -1, -1)
                    ]
                )
                post_activations_accumulated.append(X)

        states_data = {}
        states_data['Y'] = np.vstack(outputs_accumulated).T
        states_data['U'] = np.vstack(inputs_accumulated).T
        states_data['Z'] = np.vstack(pre_activations_accumulated).T
        states_data['X'] = np.vstack(post_activations_accumulated).T

        if self.standardize:
            states_data['U'] = self.scaler.fit_transform(states_data['U'].T).T
            states_data['Z'] = self.scaler.fit_transform(states_data['Z'].T).T
            states_data['X'] = self.scaler.fit_transform(states_data['X'].T).T
            states_data['Y'] = self.scaler.fit_transform(states_data['Y'].T).T

        logger.info("===== State Extraction Finished =====")
        logger.info(f"Input data shape: {states_data['U'].shape}")
        logger.info(f"Pre-activation data shape: {states_data['Z'].shape}")
        logger.info(f"Post-activation data shape: {states_data['X'].shape}")
        logger.info(f"Output data shape: {states_data['Y'].shape}")
        return states_data
    
    def train(
        self,
        solver : Callable,
        model : Optional[torch.nn.Module],
        dataloader : Optional[torch.utils.data.DataLoader] = None,
        states_data_path : Optional[str] = None,
        save_states_path : Optional[str] = None,
    ):
        """
        Train the SIM model.

        Args:
            solver (Callable): Solver to use for training.
            model (torch.nn.Module): Explicit model (teacher) to extract state data.
            dataloader (torch.utils.data.DataLoader, optional): Training data loader.
            states_data_path (str, optional): Path to the states data file.
            save_states_path (str, optional): Path to save the states data file.
        """
        # check function arguments for states extraction
        if model is None and states_data_path is None:
            raise ValueError("Either model or states_data_path must be provided")
        elif model is not None and dataloader is None:
            raise ValueError("dataloader must be provided for states extraction")

        if states_data_path is not None and save_states_path is not None:
            logger.warning("Both states_data_path and save_states_path are provided. Only states_data_path will be used.")

        if states_data_path is None:
            logger.info("===== Collecting state matrices from the explicit model =====")
            states_data = self.get_states(model, dataloader)
            if save_states_path is not None:
                logger.info(f"===== Saving states to {save_states_path} =====")
                np.save(save_states_path, states_data)
        else:
            logger.info(f"===== Loading states from {states_data_path} =====")
            states_data = np.load(states_data_path)

        logger.info("===== Start training SIM =====")

        training_config = {
            'activation_fn': self.activation_fn,
            'device': self.device,
            'atol': self.atol,
            'kappa': self.kappa,
        }

        A, B, C, D = solver.solve(states_data['X'], states_data['U'], states_data['Z'], states_data['Y'], training_config)

        logger.info(f"===== Training finished =====")

        self.weights['A'] = A
        self.weights['B'] = B
        self.weights['C'] = C
        self.weights['D'] = D

        return self
    
    def evaluate(
        self,
        dataloader : torch.utils.data.DataLoader,
    ) -> float:
        """
        Evaluate the SIM model on the given test data.

        Args:
            dataloader (torch.utils.data.DataLoader): Test data loader.
        
        Returns:
            test_accuracy (float): Accuracy of the SIM model on the given test data.
        """
        U_test = dataloader.dataset.data.flatten(1).T

        if self.standardize:
            U_test = self.scaler.transform(U_test.T).T 

        X_test = fixpoint_iteration(self.weights['A'], self.weights['B'], U_test, self.activation_fn, self.device).cpu().numpy()
        Y_test_pred = self.weights['C'] @ X_test + self.weights['D'] @ U_test.cpu().numpy()
        test_accuracy = np.mean(
            np.argmax(Y_test_pred, axis=0) == dataloader.dataset.targets.cpu().numpy()
        )
        print(f"Test accuruacy: {(test_accuracy * 100):.2f}%")
        return test_accuracy
