import math
from typing import List, Union

import numpy as np
from numba import cuda
from sklearn.preprocessing import MinMaxScaler, StandardScaler


class Transformations:
    """

    Class for storing and applying data transformations

    """

    def standardscaler(self, data: np.ndarray) -> np.ndarray:
        """

        Standardises the data by utilising sklearn.preprocessing.StandardScaler to transform the data such that its
        distribution will have a mean value 0 and standard deviation of 1. For multivariate data, this is done
        feature-wise, i.e. independently for each column of the data

        Parameters
        ----------
        data: np.ndarray
            np.ndarray with shape (n_samples, n_features) containing the data used to compute the mean and standard
            deviation used for scaling along the features axis

        Returns
        -------
        data_standardised : np.ndarray
            np.ndarray with shape (n_samples, n_features) of standardised data

        """
        if data.ndim == 1:
            data = data.reshape(-1, 1)
        self.scaler = StandardScaler()
        data_standardised = self.scaler.fit_transform(data)

        return data_standardised

    def standardscaler_transform_only(self, data: np.ndarray) -> np.ndarray:
        """

        Standardises the data using the previously determined mean and standard deviation. For multivariate data,
        this is done feature-wise, i.e. independently for each column of the data

        Parameters
        ----------
        data: np.ndarray
            np.ndarray with shape (n_samples, n_features) containing the data to be transformed by standardisation
            along the features axis

        Returns
        -------
        data_standardised : np.ndarray
            np.ndarray with shape (n_samples, n_features) of standardised data

        """
        if data.ndim == 1:
            data = data.reshape(-1, 1)
        data_standardised = self.scaler.transform(data)

        return data_standardised

    def inverse_standardscaler(self, data_standardised: np.ndarray) -> np.ndarray:
        """

        Un-standardises the data using the previously determined mean and standard deviation. For multivariate data,
        this is done feature-wise, i.e. independently for each column of the data

        Parameters
        ----------
        data_standardised: np.ndarray
            np.ndarray with shape (n_samples, n_features) containing the data to be transformed by un-standardising
            along the features axis

        Returns
        -------
        data : np.ndarray
            np.ndarray with shape (n_samples, n_features) of un-standardised data

        """
        if data_standardised.ndim == 1:
            data_standardised = data_standardised.reshape(-1, 1)
        data = self.scaler.inverse_transform(data_standardised)

        return data

    def minmaxscaler(self, data: np.ndarray) -> np.ndarray:
        """

        Min-max normalises the data by utilising sklearn.preprocessing.MinMaxScaler to transform the data by scaling
        and translating each feature individually such that it is in the given range on the training set, e.g. between
        zero and one. For multivariate data, this is done feature-wise, i.e. independently for each column of the data

        Parameters
        ----------
        data : array-like with shape (n_samples, n_features) containing the data used to identify the minimum and
        maximum values to be used for min-max normalisation scaling along the features axis.

        Returns
        -------
        data_normalised : numpy.ndarray type of normalised data

        """

        if data.ndim == 1:
            data = data.reshape(-1, 1)
        self.scaler = MinMaxScaler()
        data_normalised = self.scaler.fit_transform(data)

        return data_normalised

    def minmaxscaler_transform_only(self, data: np.ndarray) -> np.ndarray:
        """

        Min-max normalises the data using the previously identified minimum and maximum values. For multivariate data,
        this is done feature-wise, i.e. independently for each column of the data

        Parameters
        ----------
        data: np.ndarray
            np.ndarray with shape (n_samples, n_features) containing the data to be transformed by min-max normalisation
            along the features axis

        Returns
        -------
        data_normalised : np.ndarray
            np.ndarray with shape (n_samples, n_features) of normalised data

        """
        if data.ndim == 1:
            data = data.reshape(-1, 1)
        data_normalised = self.scaler.transform(data)

        return data_normalised

    def inverse_minmaxscaler(self, data_normalised: np.ndarray) -> np.ndarray:
        """

        Un-normalises the data using the previously identified minimum and maximum values. For multivariate data,
        this is done feature-wise, i.e. independently for each column of the data

        Parameters
        ----------
        data_standardised: np.ndarray
            np.ndarray with shape (n_samples, n_features) containing the data to be transformed by un-normalising
            along the features axis

        Returns
        -------
        data : np.ndarray
            np.ndarray with shape (n_samples, n_features) of un-normalised data

        """
        if data_normalised.ndim == 1:
            data_normalised = data_normalised.reshape(-1, 1)
        data = self.scaler.inverse_transform(data_normalised)

        return data

    def numba_cuda_1d_kernel(self, array: np.ndarray) -> None:
        # Kernel declaration and invocation for a 1D array and grid of threads
        # The number of threads to assign to each block when instantiating the kernel
        self.threadsperblock1d = 1024
        # Defines number of thread blocks on the grid
        self.blockspergrid1d = (array.shape[0] + (self.threadsperblock1d - 1)) // self.threadsperblock1d

    def numba_cuda_2d_kernel(self, array: np.ndarray) -> None:
        # Kernel declaration and invocation for a 2D array and grid of threads
        # The number of threads to assign to each block when instantiating the kernel. Each value of 32 is equal to the square root of the maximum number of threads per block on the computer. In this case the square root of 1024
        self.threadsperblock2d = (
            32,
            32,
        )
        # Number of thread blocks in the x direction. FYI: The product of the number of thread blocks (or “blocks per grid”) and a number of threads per block will give the total number of threads launched.
        blockspergrid_x = math.ceil(array.shape[0] / self.threadsperblock2d[0])
        # Number of thread blocks in the y direction.
        blockspergrid_y = math.ceil(array.shape[1] / self.threadsperblock2d[1])
        # Defines number of thread blocks on the xy grid
        self.blockspergrid2d = (
            blockspergrid_x,
            blockspergrid_y,
        )

    def X_min_max_norm_cuda_run(
        self,
        X: np.ndarray,
        x_max: Union[List[Union[int, float]], np.ndarray],
        x_min: Union[List[Union[int, float]], np.ndarray],
    ) -> np.ndarray:
        X_norm = np.zeros_like(X)
        self.numba_cuda_2d_kernel(X)
        self.X_min_max_norm_cuda[self.blockspergrid2d, self.threadsperblock2d](X, X_norm, x_max, x_min)

        return X_norm

    @staticmethod
    @cuda.jit
    def X_min_max_norm_cuda(
        X: np.ndarray,
        X_norm: np.ndarray,
        x_max: Union[List[Union[int, float]], np.ndarray],
        x_min: Union[List[Union[int, float]], np.ndarray],
    ) -> None:
        # Normalise the features by applying min-max feature scaling
        # Increments through a 2D grid
        x, y = cuda.grid(2)
        if x < X.shape[0] and y < X.shape[1]:
            X_norm[x, y] = (X[x][y] - x_min[y]) / (x_max[y] - x_min[y])

    def X_inverse_min_max_norm_cuda_run(
        self,
        X_norm: np.ndarray,
        x_max: Union[List[Union[int, float]], np.ndarray],
        x_min: Union[List[Union[int, float]], np.ndarray],
    ) -> np.ndarray:
        X = np.zeros_like(X_norm)
        self.numba_cuda_2d_kernel(X_norm)
        self.X_inverse_min_max_norm_cuda[self.blockspergrid2d, self.threadsperblock2d](X_norm, X, x_max, x_min)

        return X

    @staticmethod
    @cuda.jit
    def X_inverse_min_max_norm_cuda(
        X_norm: np.ndarray,
        X: np.ndarray,
        x_max: Union[List[Union[int, float]], np.ndarray],
        x_min: Union[List[Union[int, float]], np.ndarray],
    ):
        # Un-normalise the features by reversing the min-max feature scaling
        # Increments through a 2D grid
        x, y = cuda.grid(2)
        if x < X_norm.shape[0] and y < X_norm.shape[1]:
            X[x, y] = (X_norm[x][y] * (x_max[y] - x_min[y])) + x_min[y]

    def X_standardisation_cuda_run(
        self,
        X: np.ndarray,
        X_mean: Union[List[Union[int, float]], np.ndarray],
        X_stddev: Union[List[Union[int, float]], np.ndarray],
    ) -> np.ndarray:
        X_standardised = np.zeros_like(X)
        self.numba_cuda_2d_kernel(X)
        self.X_standardisation_cuda[self.blockspergrid2d, self.threadsperblock2d](X, X_standardised, X_mean, X_stddev)

        return X_standardised

    @staticmethod
    @cuda.jit
    def X_standardisation_cuda(
        X: np.ndarray,
        X_standardised: np.ndarray,
        X_mean: Union[List[Union[int, float]], np.ndarray],
        X_stddev: Union[List[Union[int, float]], np.ndarray],
    ) -> None:
        # Increments through a 2D grid
        x, y = cuda.grid(2)
        if x < X.shape[0] and y < X.shape[1]:
            X_standardised[x, y] = (X[x][y] - X_mean[y]) / X_stddev[y]

    def Y_standardisation_cuda_run(
        self, Y: np.ndarray, Y_mean: Union[int, float], Y_stddev: Union[int, float]
    ) -> np.ndarray:
        if len(Y.shape) > 1:
            Y = Y.flatten()

        Y_standardised = np.zeros_like(Y)
        self.numba_cuda_1d_kernel(Y)
        self.Y_standardisation_cuda[self.blockspergrid1d, self.threadsperblock1d](Y, Y_standardised, Y_mean, Y_stddev)

        return Y_standardised

    @staticmethod
    @cuda.jit
    def Y_standardisation_cuda(
        Y: np.ndarray, Y_standardised: np.ndarray, Y_mean: Union[int, float], Y_stddev: Union[int, float]
    ) -> None:
        # Increments through a 1D grid
        x = cuda.grid(1)
        if x < Y.shape[0]:
            Y_standardised[x] = (Y[x] - Y_mean) / Y_stddev

    def Y_inverse_standardisation_cuda_run(
        self, Y_standardised: np.ndarray, Y_mean: Union[int, float], Y_stddev: Union[int, float]
    ) -> np.ndarray:
        if len(Y_standardised.shape) > 1:
            Y_standardised = Y_standardised.flatten()

        Y = np.zeros_like(Y_standardised)
        self.numba_cuda_1d_kernel(Y_standardised)
        self.Y_inverse_standardisation_cuda[self.blockspergrid1d, self.threadsperblock1d](
            Y_standardised, Y, Y_mean, Y_stddev
        )

        return Y

    @staticmethod
    @cuda.jit
    def Y_inverse_standardisation_cuda(
        Y_standardised: np.ndarray, Y: np.ndarray, Y_mean: Union[int, float], Y_stddev: Union[int, float]
    ) -> None:
        # Increments through a 1D grid
        x = cuda.grid(1)
        if x < Y.shape[0]:
            Y[x] = (Y_standardised[x] * Y_stddev) + Y_mean
