"""
Low-rank probability density model for continuous distributions

Author : Laurène DAVID
"""

import numpy as np
import warnings
import math
# from lowrankdensity.models import Discrete
from lowrankdensity.models.discrete import Discrete


class Continuous(Discrete):
    """Low-rank Bivariate Continuous Density Estimation

    Parameters
    ----------
    alpha : float, default = 1
    Level of precision of the density estimation


    Attributes :
    ---------
    density_function : function object
    A density function for bivariate distributions with inputs x,y each a 1d array


    Example :
    ----------
    >>> import numpy as np
    >>> from lowrankdensity.datasets import generate_lowrank_continuous
    >>> from lowrankdensity.models import Continuous

    # Generate low rank continuous samples
    >>> X = generate_lowrank_continuous()

    # Fit samples to the low rank Discrete model
    >>> model = Continuous(alpha=0.1)
    >>> model.fit(X)

    # Get the estimated probability density function f
    >>> lowrank_estimator = model.density_function

    # Compute the probability density function on two arrays x,y
    >>> x,y = np.linspace(0,1,100), np.linspace(0,1,100)
    >>> lowrank_matrix = model.pdf(x,y)

    # Generate new samples
    >>> new_samples = model.sample(n_samples=1000)

    """

    L = 1

    def __init__(self, alpha=1):
        self.alpha = alpha
        self.density_function = None

    def continuous_1d(self, Z, L=1):
        n = Z.shape[0]
        r = np.min(Z[:int(n / 2)])
        R = np.max(Z[:int(n / 2)])

        if (R - r < n**(-1 / 3) * L**(-1 / 2)):
            return lambda x: 1 / (R - r) if (x <= r and x >= R) else 0

        else:
            h = math.floor((R - r) * n**(1 / 3) * L**(1 / 2))**(-1) * (R - r)
            E = np.arange(-math.floor(r / h), math.ceil((1 - r) / h - 1))

            def f(x):
                N = np.zeros((len(E) + 1,))
                for i in range(int(n / 2) + 1, n):
                    N[int((Z[i] - r) / h)] += 1
                return (1 / h) * N[int((x - r) / h)]

            return f

    def fit(self, X):
        """
        Fit the Bivariate Continuous Density model on the data


        Parameters :
        ----------
        X : np.array of shape (n,2)
        A numpy array generated by a joint continuous distribution


        Return :
        ---------
        self : object
        Returns the instance itself.

        """

        if not isinstance(X, np.ndarray):
            raise TypeError(f"Input X should be a nd.array, not a {type(X)}")

        if X.shape[0] == 0:
            raise ValueError("X is an empty array")

        if self.alpha < 0:
            raise ValueError(f"alpha should be positive")

        if type(self.alpha) not in (int, float):
            raise TypeError(f"alpha should be a float, not {type(self.alpha)}")

        n = X.shape[0]
        L = self.L
        r1, R1 = np.min(X[:int(n / 2), 0]), np.max(X[:int(n / 2), 0])
        r2, R2 = np.min(X[:int(n / 2), 1]), np.max(X[:int(n / 2), 1])

        if R1 - r1 < n**(-1 / 3) * L**(-1 / 2):
            g = self.continuous_1d(Z=X[int(n / 2 + 1):, 1])
            self.density_function = lambda x, y: (1 / (R1 - r1)) * g(y) if (r1 <= x < R1) else 0
            return self

        if R2 - r2 < n**(-1 / 3) * L**(-1 / 2):
            g = self.continuous_1d(Z=X[int(n / 2 + 1):, 0])
            self.density_function = lambda x, y: (1 / (R2 - r2)) * g(x) if (r2 <= y < R2) else 0
            return self

        h1 = math.floor((R1 - r1) * n**(1 / 3) * L**(1 / 2))**(-1) * (R1 - r1)
        h2 = math.floor((R2 - r2) * n**(1 / 3) * L**(1 / 2))**(-1) * (R2 - r2)

        E1 = np.arange(-math.floor(r1 / h1), math.ceil((1 - r1) / h1 - 1))
        E2 = np.arange(-math.floor(r2 / h2), math.ceil((1 - r2) / h2 - 1))

        N1 = np.zeros((len(E1) + 1, len(E2) + 1))
        N2 = np.zeros((len(E1) + 1, len(E2) + 1))

        for i, j in zip(range(int(n / 2) + 1, int(3 * n / 4)), range(int(3 * n / 4) + 1, n)):
            x1, y1 = X[i, :]
            x2, y2 = X[j, :]

            N1[int((x1 - r1) / h1), int((y1 - r2) / h2)] = N1[int((x1 - r1) / h1), int((y1 - r2) / h2)] + 1
            N2[int((x2 - r1) / h1), int((y2 - r2) / h2)] = N2[int((x2 - r1) / h1), int((y2 - r2) / h2)] + 1

        P = super()._compute_matrix(n=int(n / 2), Y1=N1, Y2=N2, discrete=False)

        grid_dict = {"r1": r1, "r2": r2, "R1": R1, "R2": R2, "h1": h1, "h2": h2, "E1": E1, "E2": E2}
        self.grid_params = grid_dict

        m1 = math.floor((R1 - r1) * n**(1 / 3) * L**(1 / 2))
        m2 = math.floor((R2 - r2) * n**(1 / 3) * L**(1 / 2))

        def f(x, y):
            x1, y1 = int((x - r1) / h1), int((y - r2) / h2)
            if 0 <= x1 < m1 and 0 <= y1 < m2:
                return (1 / (h1 * h2)) * P[x1, y1]
            else:
                return (1 / (h1 * h2)) * (2 / n) * (N1[x1, y1] + N2[x1, y1])

        self.density_function = f
        return None

    def pdf(self, x, y):
        """
        Compute a matrix with the low-rank probability density function values

        Parameters
        ---------
        x : 1D nd.array
        The first dimension x to compute the density estimator f(x,y)

        y : 1D nd.array
        The second dimension x to compute the density estimator f(x,y)


        Return
        ---------
        mat : nd.array of shape (len(x),len(y))
        Matrix with computed density function values for inputs x and y

        """
        if (not isinstance(x, np.ndarray)) or (not isinstance(y, np.ndarray)):
            raise TypeError("x and y should be a numpy arrays")

        if (x.ndim > 1) or (y.ndim > 1):
            raise ValueError("x and y should be 1D arrays")

        if (np.all((x <= 1) & (x >= 0)) == False) or (np.all((y <= 1) & (y >= 0)) == False):
            warnings.warn("The low rank continuous estimator is sued for original densities with support on [0,1]x[0,1]")

        lowrank_estimator = self.density_function
        return np.array([[lowrank_estimator(i, j) for j in y] for i in x])

    def sample(self, n_samples=1000):
        """
        Sample continuous data with the low rank density function estimator


        Parameters
        -------
        n_samples : int, default=1
        Number of samples to draw from distribution


        Returns
        -------
        sample: nd.array of shape (n_samples,)
        Samples drawn from discrete distribution with probability matrix P

        """
        r1, r2, R1, R2, h1, h2, E1, E2 = self.grid_params.values()
        f = self.density_function

        # 1. Define P
        range_a1, range_a2 = r1 + E1[:len(E1) - 1] * h1, r1 + E1[1:] * h1
        range_b1, range_b2 = r2 + E2[:len(E2) - 1] * h2, r2 + E2[1:] * h2
        P2 = np.array([[(a2 - a1) * (b2 - b1) * f(a1, b1) for a1, a2 in zip(range_a1, range_a2)] for b1, b2 in zip(range_b1, range_b2)])
        P2 = P2 / np.sum(P2)

        # 2. Sample Z with a multinomial distribution with P
        p = P2.flatten()
        nrow2, ncol2 = P2.shape
        Z = np.random.multinomial(n=1, pvals=p, size=n_samples).reshape((n_samples, nrow2, ncol2))
        Z = np.argwhere(Z == 1)[:, 1:]

        # 3. Use Z to sample continuous data with uniform distribution
        Z1, Z2 = Z[:, 0], Z[:, 1]
        rangeZ_a1, rangeZ_a2 = r1 + Z1 * h1, r1 + (Z1 + 1) * h1
        rangeZ_b1, rangeZ_b2 = r2 + Z2 * h2, r2 + (Z2 + 1) * h2
        X = np.array([np.random.uniform(low=[a1, b1], high=[a2, b2], size=2).T for a1, a2, b1, b2 in zip(rangeZ_a1, rangeZ_a2, rangeZ_b1, rangeZ_b2)])

        return X
