"""This is a convolution kernel that first generates random
features on the input, then feeds this random feature profile
into an RBF kernel, analogous to a 3-layer CNN."""
from math import ceil

import numpy as np
from scipy.stats import chi

from xGPR.xgpr_cpu_rfgen_cpp_ext import cpuConv1dMaxpool
from xGPR.xgpr_cpu_rfgen_cpp_ext import cpuRBFFeatureGen, cpuRBFGrad
try:
    import cupy as cp
    from xGPR.xgpr_cuda_rfgen_cpp_ext import cudaConv1dMaxpool
    from xGPR.xgpr_cuda_rfgen_cpp_ext import cudaRBFFeatureGen, cudaRBFGrad
except:
    pass

from ..kernel_baseclass import KernelBaseclass




class Conv1dTwoLayer(KernelBaseclass):
    """This kernel first runs ReLU activation and global maxpool
    RF generation on the input, then feeds the output of that
    initial operation into an RBF kernel. This is equivalent
    to running the FastConv1d feature extractor on the input
    then running the output of that operation into an xGPR
    model with an RBF kernel. Sometimes using the FastConv
    feature extractor may be more convenient, and in other cases
    using this single-call kernel may be more convenient.

    Attributes:
        hyperparams (np.ndarray): This kernel has two
            hyperparameters: lambda_ (noise)
            and sigma (inverse mismatch tolerance).
        bounds (np.ndarray): The builtin bounds for hyperparameter optimization,
            which can be overriden / reset by user.
        conv_width (int): The width of the convolution kernel.
            This hyperparameter can be set based on experimentation
            or domain knowledge.
        init_rffs (int): The number of random features generated by
            the first layer.
        radem_diag1: The diagonal matrices for the SORF transform for the
            first layer.
        chi_arr1: A diagonal array whose elements are drawn from the chi
            distribution. Ensures the marginals of the matrix resulting
            from S H D1 H D2 H D3 are correct.
        radem_diag2: The diagonal matrices for the SORF transform for the
            second layer.
        chi_arr2: The diagonal array for the SORF transform for the
            second layer.
        num_threads (int): Number of threads to use if running on CPU;
            ignored if running on GPU.
    """

    def __init__(self, xdim, num_rffs, random_seed = 123,
                device = "cpu", num_threads = 2,
                double_precision = False,
                kernel_spec_parms = {}):
        """Constructor.

        Args:
            xdim (tuple): The dimensions of the input. Either (N, D) or (N, M, D)
                where N is the number of datapoints, D is number of features
                and M is number of timepoints or sequence elements (convolution
                kernels only).
            num_rffs (int): The user-requested number of random Fourier features.
                For sine-cosine kernels (RBF, Matern), this will be saved by the
                class as num_rffs.
            random_seed (int): The seed to the random number generator.
            device (str): One of 'cpu', 'cuda'. Indicates the starting device.
            num_threads (int): The number of threads to use for random feature generation
                if running on CPU. If running on GPU, this is ignored.
            double_precision (bool): If True, generate random features in double precision.
                Otherwise, generate as single precision.
            conv_width (int): The width of the convolution to perform.
            kernel_spec_parms (dict): A dictionary of additional kernel-specific
                attributes.

        Raises:
            RuntimeError: A RuntimeError is raised if the dimensions of the input are
                inappropriate given the conv_width.
        """
        if "conv_width" not in kernel_spec_parms:
            raise ValueError("conv_width must be included as a kernel-specific "
                    "parameter if using a sequence kernel.")
        if "init_rffs" not in kernel_spec_parms:
            raise ValueError("init_rffs must be included as a kernel-specific "
                    "parameter if using the 2 layer conv1d kernel.")
        if len(xdim) != 3:
            raise RuntimeError("Tried to initialize a Conv1d kernel with a 2d x-"
                    "array! x should be a 3d array for Conv1d.")

        self.init_rffs = kernel_spec_parms["init_rffs"]
        if self.init_rffs % 2 != 0:
            raise RuntimeError("Number of init rffs should be an even number.")

        super().__init__(num_rffs, xdim, num_threads = 2,
                sine_cosine_kernel = True, double_precision = double_precision,
                kernel_spec_parms = kernel_spec_parms)
        self.hyperparams = np.ones((2))
        self.bounds = np.asarray([[1e-3,5], [1e-6, 1e2]])

        # First, set up the information needed for the first layer.
        rng = np.random.default_rng(random_seed)
        self.conv_width = kernel_spec_parms["conv_width"]

        dim2_no_padding = self.conv_width * xdim[2]
        padded_dims = 2**ceil(np.log2(max(dim2_no_padding, 2)))
        init_calc_featsize = ceil(self.init_rffs / padded_dims) * \
                        padded_dims

        radem_array = np.asarray([-1,1], dtype=np.int8)

        self.radem_diag1 = rng.choice(radem_array, size=(3, 1, init_calc_featsize),
                                replace=True)
        self.chi_arr1 = chi.rvs(df=padded_dims, size=self.init_rffs,
                            random_state = random_seed)


        # Next, set up the information needed for the RBF output layer.
        padded_dims = 2**ceil(np.log2(max(self.init_rffs, 2)))

        if padded_dims < self.num_freqs:
            nblocks = ceil(self.num_freqs / padded_dims)
        else:
            nblocks = 1
        self.radem_diag2 = rng.choice(radem_array, size=(3, 1,
                nblocks * padded_dims), replace=True)
        self.chi_arr2 = chi.rvs(df=padded_dims, size=self.num_freqs,
                            random_state = random_seed)
        if not self.double_precision:
            self.chi_arr1 = self.chi_arr1.astype(np.float32)
            self.chi_arr2 = self.chi_arr2.astype(np.float32)

        self.num_threads = num_threads
        self.device = device


    def kernel_specific_set_device(self, new_device):
        """Called by parent class when device is changed. Moves
        some of the object parameters to the appropriate device."""
        if new_device == "cuda":
            self.radem_diag1 = cp.asarray(self.radem_diag1)
            self.chi_arr1 = cp.asarray(self.chi_arr1)
            self.radem_diag2 = cp.asarray(self.radem_diag2)
            self.chi_arr2 = cp.asarray(self.chi_arr2)
        elif not isinstance(self.radem_diag1, np.ndarray):
            self.radem_diag1 = cp.asnumpy(self.radem_diag1)
            self.chi_arr1 = cp.asnumpy(self.chi_arr1)
            self.radem_diag2 = cp.asnumpy(self.radem_diag2)
            self.chi_arr2 = cp.asnumpy(self.chi_arr2)



    def kernel_specific_set_hyperparams(self):
        """Provided for consistency with baseclass. This
        kernel has no kernel-specific properties that must
        be reset after hyperparameters are changed."""
        return


    def kernel_specific_transform(self, input_x, sequence_length):
        """Performs the feature generation."""
        if sequence_length is None:
            raise ValueError("sequence_length is required for convolution kernels.")
        if input_x.shape[2] != self._xdim[2]:
            raise RuntimeError("Unexpected input shape supplied.")

        sequence_length = sequence_length.astype(np.int32, copy=False)

        if self.device == "cpu":
            featurized_x = np.zeros((input_x.shape[0], self.init_rffs), np.float32)
            cpuConv1dMaxpool(input_x, featurized_x, self.radem_diag1, self.chi_arr1,
                    sequence_length, self.conv_width, self.num_threads)
            xtrans = np.zeros((featurized_x.shape[0], self.num_rffs), np.float64)
            featurized_x *= self.hyperparams[1]
            cpuRBFFeatureGen(featurized_x, xtrans, self.radem_diag2, self.chi_arr2,
                self.num_threads, self.fit_intercept)

        else:
            featurized_x = cp.zeros((input_x.shape[0], self.init_rffs), cp.float32)
            cudaConv1dMaxpool(input_x, featurized_x, self.radem_diag1, self.chi_arr1,
                    sequence_length, self.conv_width)

            xtrans = cp.zeros((featurized_x.shape[0], self.num_rffs), cp.float64)
            featurized_x *= self.hyperparams[1]
            cudaRBFFeatureGen(featurized_x, xtrans, self.radem_diag2, self.chi_arr2,
                self.fit_intercept)

        return xtrans



    def kernel_specific_gradient(self, input_x, sequence_length = None):
        """The gradient for kernel-specific hyperparameters.

        Args:
            input_x: A cupy or numpy array containing the raw input data.
            sequence_length: A numpy or cupy array containing the number of
                elements in each sequence -- so that zero padding can be masked.

        Returns:
            output_x: A cupy or numpy array containing the random feature
                representation of the input.
            dz_dsigma: A cupy or numpy array containing the derivative of
                output_x with respect to the kernel-specific hyperparameters.
        """
        if self.device == "cpu":
            featurized_x = np.zeros((input_x.shape[0], self.init_rffs), np.float32)
            cpuConv1dMaxpool(input_x, featurized_x, self.radem_diag1, self.chi_arr1,
                    sequence_length, self.conv_width, self.num_threads)

            output_x = np.zeros((input_x.shape[0], self.num_rffs), np.float64)
            dz_dsigma = np.zeros((input_x.shape[0], self.num_rffs, 1), np.float64)
            cpuRBFGrad(featurized_x, output_x, dz_dsigma, self.radem_diag2, self.chi_arr2,
                self.hyperparams[1], self.num_threads, self.fit_intercept)
        else:
            featurized_x = cp.zeros((input_x.shape[0], self.init_rffs), cp.float32)
            cudaConv1dMaxpool(input_x, featurized_x, self.radem_diag1, self.chi_arr1,
                    sequence_length, self.conv_width)

            output_x = cp.zeros((input_x.shape[0], self.num_rffs), cp.float64)
            dz_dsigma = cp.zeros((input_x.shape[0], self.num_rffs, 1), cp.float64)
            cudaRBFGrad(featurized_x, output_x, dz_dsigma, self.radem_diag2, self.chi_arr2,
                self.hyperparams[1], self.fit_intercept)
        return output_x, dz_dsigma
