import h5py
import numpy as np

from TfELM.Optimizers.ELMOptimizer import ELMOptimizer
from TfELM.Resources.ActivationFunction import ActivationFunction
import tensorflow as tf

from TfELM.Resources.Kernel import CombinedProductKernel, Kernel, CombinedSumKernel
from TfELM.Resources.gram_schmidt import gram_schmidt
from TfELM.Resources.kernel_distances import calculate_pairwise_distances_vector, calculate_pairwise_distances


class SSKELMLayer:
    """
        Semi-Supervised Kernel Extreme Learning Machine (SSKELM) layer.

        This layer implements a semi-supervised version of the Kernel Extreme Learning Machine (KELM) algorithm. It is
        capable of handling both labeled and unlabeled data for training, utilizing kernel-based feature mapping and
        Laplacian graph regularization.

        Parameters:
        -----------
        - kernel (Kernel): Kernel function to be used for feature mapping.
        - activation (str): Name of the activation function. Default is 'tanh'.
        - act_params (dict): Parameters for the activation function. Default is None.
        - C (float): Regularization parameter. Default is 1.0.
        - lam (float): Laplacian graph regularization parameter. Default is 0.5.
        - nystrom_approximation (bool): Whether to use Nystrom approximation for kernel matrix computation. Default is False.
        - landmark_selection_method (str): Method for landmark selection in Nystrom approximation. Default is 'random'.

        Attributes:
        -----------
        - error_history (None or array-like): History of errors during training.
        - feature_map (None or array-like): Feature map generated by the layer.
        - name (str): Name of the layer.
        - beta (None or tf.Tensor): Weight matrix.
        - input (None or tf.Tensor): Input data.
        - output (None or tf.Tensor): Output data.
        - activation_name (str): Name of the activation function.
        - activation (callable): Activation function.
        - C (float): Regularization parameter.
        - kernel (Kernel): Kernel function.
        - nystrom_approximation (bool): Flag indicating Nystrom approximation usage.
        - landmark_selection_method (str): Method for landmark selection in Nystrom approximation.
        - denoising (None or str): Denoising method.
        - denoising_param (None or float): Denoising parameter.

        Methods:
        -----------
        - build(input_shape): Build the layer by initializing necessary variables.
        - fit(x_labeled, x_unlabeled, y_labeled, y_unlabeled): Train the layer on labeled and unlabeled data.
        - predict(x): Predict output for the given input data.
        - predict_proba(x): Predict class probabilities for the given input data.
        - calc_output(x): Calculate the output of the layer for the given input data.
        - __str__(): Return a string representation of the layer.
        - count_params(): Count the number of trainable and non-trainable parameters in the layer.
        - to_dict(): Convert the layer's attributes to a dictionary.
        - load(attributes): Load a layer instance from a dictionary of attributes.

        Notes:
        -----------
        - This layer supports both labeled and unlabeled data for training.
        - It utilizes kernel-based feature mapping and Laplacian graph regularization for enhanced learning.
        - Nystrom approximation can be employed for efficient kernel matrix computation.

        Example:
        -----------
        Initialize a Kernel (it can be instanced as Kernel class and its subclasses like CombinedProductKernel)

        >>> kernel = CombinedProductKernel([Kernel("rational_quadratic"), Kernel("exponential")])

        Initializing the Semi-Supervised Kernel Extreme Learning Machine (SS-KELM) layer.

        >>> layer = SSKELMLayer(kernel=kernel, lam=0.001)

        Initializing the SS-KELM model with the defined layer

        >>> model = SSELMModel(layer)
    """
    def __init__(self,
                 kernel,
                 activation='tanh',
                 act_params=None,
                 C=1.0,
                 lam=0.5,
                 nystrom_approximation=False,
                 landmark_selection_method='random',
                 **params):
        self.error_history = None
        self.feature_map = None
        self.name = "sselm"
        self.beta = None
        self.input = None
        self.output = None
        self.lam = lam
        self.act_params = act_params
        if act_params is None:
            act = ActivationFunction(1.0)
        elif "act_param" in act_params and "act_param2" in act_params:
            act = ActivationFunction(act_param=act_params["act_param"], act_param2=act_params["act_param2"])
        elif "act_param" in act_params:
            act = ActivationFunction(act_param=act_params["act_param"])
        elif "knots" in act_params:
            act = ActivationFunction(knots=act_params["knots"])
        else:
            raise Exception("TypeError: Wrong specified activation function parameters")
        self.activation_name = activation
        self.activation = eval("act." + activation)
        self.C = C
        self.kernel = kernel
        self.nystrom_approximation = nystrom_approximation
        self.landmark_selection_method = landmark_selection_method

        if "K" in params:
            self.K = params.pop("K")
        if "input" in params:
            self.input = params.pop("input")
        if "beta" in params:
            self.beta = params.pop("beta")

        if "denoising" in params:
            self.denoising = params.pop("denoising")
        else:
            self.denoising = None

        if "denoising_param" in params:
            self.denoising_param = params.pop("denoising_param")
        else:
            self.denoising_param = None

    def build(self, input_shape):
        """
          Build the layer by initializing necessary variables.

          Parameters:
          -----------
          - input_shape (tuple): Shape of the input data.

          Returns:
          -----------
          None

          Example:
          -----------
            >>> kernel = CombinedProductKernel([Kernel("rational_quadratic"), Kernel("exponential")])
            >>> layer = SSKELMLayer(kernel=kernel, lam=0.001)
            >>> layer.build(x.shape)
        """
        observations = input_shape[0]
        self.K = tf.Variable(
            tf.zeros(shape=(observations, observations)),
            dtype=tf.float32,
            trainable=False
        )

    def fit(self, x_labeled, x_unlabeled, y_labeled, y_unlabeled):
        """
            Train the layer on labeled and unlabeled data.

            Parameters:
            -----------
            - x_labeled (np.ndarray or tf.Tensor): Labeled input data.
            - x_unlabeled (np.ndarray or tf.Tensor): Unlabeled input data.
            - y_labeled (np.ndarray or tf.Tensor): Labeled target data.
            - y_unlabeled (np.ndarray or tf.Tensor): Unlabeled target data.

            Returns:
            -----------
            None

            Example:
            -----------
            >>> kernel = CombinedProductKernel([Kernel("rational_quadratic"), Kernel("exponential")])
            >>> layer = SSKELMLayer(kernel=kernel, lam=0.001)
            >>> layer.build(x.shape)
            >>> layer.fit(X_labeled, X_unlabeled, y_labeled, y_unlabeled)
        """
        y_un_zero = np.zeros(shape=np.shape(y_unlabeled))
        n_labeled = np.shape(x_labeled)[0]
        X_combined = np.vstack([x_labeled, x_unlabeled])
        Y_combined = np.vstack([y_labeled, y_un_zero])

        X_combined = tf.cast(X_combined, dtype=tf.float32)
        Y_combined = tf.cast(Y_combined, dtype=tf.float32)

        self.input = X_combined

        # Laplacian Graph
        squared_norms = tf.reduce_sum(tf.square(X_combined), axis=1, keepdims=True)
        dot_product = tf.matmul(X_combined, X_combined, transpose_b=True)
        distances = squared_norms - 2 * dot_product + tf.transpose(squared_norms)
        distances = tf.maximum(distances, 0.0)
        sigma = 1.0
        W = tf.exp(-distances / (2.0 * sigma ** 2))
        D = tf.linalg.diag(tf.reduce_sum(W, axis=1))
        L_unnormalized = D - W
        D_sqrt_inv = tf.linalg.inv(tf.linalg.sqrtm(D))
        L = tf.matmul(tf.matmul(D_sqrt_inv, L_unnormalized), D_sqrt_inv)

        ni = self.C / tf.reduce_sum(Y_combined, axis=0)
        C = tf.linalg.diag(tf.reduce_sum(tf.multiply(Y_combined, ni), axis=1))
        m, n = tf.shape(C)[0], tf.shape(C)[1]
        mask = tf.math.logical_and(tf.range(m) > n_labeled, tf.range(n) > n_labeled)
        C = tf.where(mask, tf.zeros_like(C), C)

        if self.nystrom_approximation:
            num_rows = tf.shape(X_combined)[0]
            shuffled_indices = tf.random.shuffle(tf.range(num_rows))
            selected_indices = shuffled_indices[:100]
            L = tf.gather(X_combined, selected_indices)
            C = calculate_pairwise_distances_vector(X_combined, L, self.kernel.ev)
            W = calculate_pairwise_distances(L, self.kernel.ev)
            K = tf.matmul(tf.matmul(C, tf.linalg.inv(W)), C, transpose_b=True)
        else:
            K = calculate_pairwise_distances(X_combined, self.kernel.ev)

        CHHt = tf.matmul(tf.matmul(C, K), K, transpose_b=True)
        lLHHt = self.lam * tf.matmul(tf.matmul(L, K), K, transpose_b=True)
        i = tf.eye(tf.shape(CHHt)[0])
        inv = tf.linalg.inv(i + CHHt + lLHHt)
        beta = tf.matmul(tf.matmul(tf.matmul(K, inv, transpose_a=True), C), Y_combined)

        self.beta = beta
        self.K = K

    def predict(self, x):
        """
            Predict output for the given input data.

            Parameters:
            -----------
            - x (np.ndarray or tf.Tensor): Input data.

            Returns:
            -----------
            tf.Tensor: Predicted output tensor.

            Example:
            -----------
            >>> kernel = CombinedProductKernel([Kernel("rational_quadratic"), Kernel("exponential")])
            >>> layer = SSKELMLayer(kernel=kernel, lam=0.001)
            >>> layer.build(x.shape)
            >>> layer.fit(X_labeled, X_unlabeled, y_labeled, y_unlabeled)
            >>> pred = layer.predict(test_data)
        """
        x = tf.cast(x, dtype=tf.float32)
        k = calculate_pairwise_distances_vector(x, self.input, self.kernel.ev)
        kpKT = tf.matmul(k, self.K)
        output = tf.matmul(kpKT, self.beta)
        self.output = output
        return output

    def predict_proba(self, x):
        """
            Predict class probabilities for the given input data.

            Parameters:
            -----------
            - x (np.ndarray or tf.Tensor): Input data.

            Returns:
            -----------
            tf.Tensor: Predicted class probabilities' tensor.

            Example:
            -----------
            >>> kernel = CombinedProductKernel([Kernel("rational_quadratic"), Kernel("exponential")])
            >>> layer = SSKELMLayer(kernel=kernel, lam=0.001)
            >>> layer.build(x.shape)
            >>> layer.fit(X_labeled, X_unlabeled, y_labeled, y_unlabeled)
            >>> pred = layer.predict_proba(test_data)
        """
        x = tf.cast(x, dtype=tf.float32)
        pred = self.predict(x)
        return tf.keras.activations.softmax(pred)

    def calc_output(self, x):
        """
            Calculate the output of the layer for the given input data.

            Parameters:
            -----------
            - x (np.ndarray or tf.Tensor): Input data.

            Returns:
            -----------
            tf.Tensor: Output tensor.
        """
        x = tf.cast(x, dtype=tf.float32)
        k = calculate_pairwise_distances_vector(x, self.input, self.kernel.ev)
        out = self.activation(tf.matmul(k, self.beta))
        self.output = out
        return out

    def __str__(self):
        """
        Returns a string representation of the ELM layer.

        Returns:
        -----------
        str: String representation.
        """
        return f"{self.name}, kernel: {self.kernel.__class__.__name__}"

    def count_params(self):
        """
        Counts the number of trainable and non-trainable parameters in the ELM layer.

        Returns:
        -----------
        dict: Dictionary containing counts for trainable, non-trainable, and total parameters.
        """
        if self.beta is None:
            trainable = 0
        else:
            trainable = self.beta.shape[0] * self.beta.shape[1]

        non_trainable = 0
        return {'trainable': trainable, 'non_trainable': non_trainable, 'all': trainable + non_trainable}

    def to_dict(self):
        """
            Convert the layer's attributes to a dictionary.

            Returns:
            -----------
            dict: Dictionary containing the layer's attributes.
        """
        attributes = {
            'name': 'SSKELMLayer',
            'C': self.C,
            "beta": self.beta,
            "kernel": self.kernel.kernel_name,
            "kernel_param": self.kernel.kernel_param,
            "kernel_type": self.kernel.__class__.__name__,
            "nystrom_approximation": self.nystrom_approximation,
            "landmark_selection_method": self.landmark_selection_method,
            "input": self.input,
            "K": self.K,
            "denoising": self.denoising,
            "denoising_param": self.denoising_param,
            "lam": self.lam
        }
        filtered_attributes = {key: value for key, value in attributes.items() if value is not None}
        return filtered_attributes

    @classmethod
    def load(cls, attributes):
        """
            Load a layer instance from a dictionary of attributes.

            Parameters:
            -----------
            - attributes (dict): Dictionary containing the layer's attributes.

            Returns:
            -----------
            SSKELMLayer: Loaded layer instance.
        """
        if "kernel" in attributes:
            k_n = attributes.pop("kernel")
            k_p = attributes.pop("kernel_param")
            k_t = attributes.pop("kernel_type")
            if k_t == "Kernel":
                k = Kernel(kernel_name=k_n, param=k_p)
            elif k_t == "CombinedSumKernel":
                kernels = []
                for n, p in zip(k_n, k_p):
                    kernels.append(Kernel(kernel_name=n.decode('utf-8'), param=p))
                k = CombinedSumKernel(kernels)
            else:
                kernels = []
                for n, p in zip(k_n, k_p):
                    kernels.append(Kernel(kernel_name=n.decode('utf-8'), param=p))
                k = CombinedProductKernel(kernels)
        else:
            k = Kernel()

        attributes.update({"kernel": k})
        layer = cls(**attributes)
        return layer
