from statsmodels.compat.python import lzip

import numpy as np

from statsmodels.tools.validation import array_like

from . import kernels


# TODO: should this be a function?
class KDE:
    """
    Kernel Density Estimator

    Parameters
    ----------
    x : array_like
        N-dimensional array from which the density is to be estimated
    kernel : Kernel Class
        Should be a class from *
    """

    # TODO: amend docs for Nd case?
    def __init__(self, x, kernel=None):
        x = array_like(x, "x", maxdim=2, contiguous=True)
        if x.ndim == 1:
            x = x[:, None]

        nobs, n_series = x.shape

        if kernel is None:
            kernel = kernels.Gaussian()  # no meaningful bandwidth yet

        if n_series > 1:
            if isinstance(kernel, kernels.CustomKernel):
                kernel = kernels.NdKernel(n_series, kernels=kernel)

        self.kernel = kernel
        self.n = n_series  # TODO change attribute
        self.x = x

    def density(self, x):
        return self.kernel.density(self.x, x)

    def __call__(self, x, h="scott"):
        return np.array([self.density(xx) for xx in x])

    def evaluate(self, x, h="silverman"):
        density = self.kernel.density
        return np.array([density(xx) for xx in x])


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    from numpy import random

    import statsmodels.nonparametric.bandwidths as bw
    from statsmodels.sandbox.nonparametric.testdata import kdetest

    # 1-D case
    random.seed(142)
    x = random.standard_t(4.2, size=50)
    h = bw.bw_silverman(x)
    # NOTE: try to do it with convolution
    support = np.linspace(-10, 10, 512)

    kern = kernels.Gaussian(h=h)
    kde = KDE(x, kern)
    print(kde.density(1.015469))
    print(0.2034675)
    Xs = np.arange(-10, 10, 0.1)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(Xs, kde(Xs), "-")
    ax.set_ylim(-10, 10)
    ax.set_ylim(0, 0.4)

    # 2-D case
    x = lzip(kdetest.faithfulData["eruptions"], kdetest.faithfulData["waiting"])
    x = np.array(x)
    x = (x - x.mean(0)) / x.std(0)
    nobs = x.shape[0]
    H = kdetest.Hpi
    kern = kernels.NdKernel(2)
    kde = KDE(x, kern)
    print(kde.density(np.matrix([1, 2])))  # .T
    plt.figure()
    plt.plot(x[:, 0], x[:, 1], "o")

    n_grid = 50
    xsp = np.linspace(x.min(0)[0], x.max(0)[0], n_grid)
    ysp = np.linspace(x.min(0)[1], x.max(0)[1], n_grid)
    #    xsorted = np.sort(x)
    #    xlow = xsorted[nobs/4]
    #    xupp = xsorted[3*nobs/4]
    #    xsp = np.linspace(xlow[0], xupp[0], n_grid)
    #    ysp = np.linspace(xlow[1], xupp[1], n_grid)
    xr, yr = np.meshgrid(xsp, ysp)
    kde_vals = np.array(
        [kde.density(np.matrix([xi, yi])) for xi, yi in zip(xr.ravel(), yr.ravel())]
    )
    plt.contour(xsp, ysp, kde_vals.reshape(n_grid, n_grid))

    plt.show()

    # 5 D case
#    random.seed(142)
#    mu = [1.0, 4.0, 3.5, -2.4, 0.0]
#    sigma = np.matrix(
#        [[ 0.6 - 0.1*abs(i-j) if i != j else 1.0 for j in xrange(5)] for i in xrange(5)])
#    x = random.multivariate_normal(mu, sigma, size = 100)
#    kern = kernel.Gaussian()
#    kde = KernelEstimate( x, kern )
