"""
This example illustrates the nonparametric estimation of a
bivariate bi-modal distribution that is a mixture of two normal
distributions.

author: George Panterov
"""

from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np

import statsmodels.api as sm

if __name__ == "__main__":
    np.random.seed(123456)

    # generate the data
    nobs = 500
    BW = "cv_ml"

    mu1 = [3, 4]
    mu2 = [6, 1]
    cov1 = np.asarray([[1, 0.7], [0.7, 1]])
    cov2 = np.asarray([[1, -0.7], [-0.7, 1]])

    ix = np.random.uniform(size=nobs) > 0.5
    V = np.random.multivariate_normal(mu1, cov1, size=nobs)
    V[ix, :] = np.random.multivariate_normal(mu2, cov2, size=nobs)[ix, :]

    x = V[:, 0]
    y = V[:, 1]

    dens = sm.nonparametric.KDEMultivariate(
        data=[x, y],
        var_type="cc",
        bw=BW,
        defaults=sm.nonparametric.EstimatorSettings(efficient=True),
    )

    supportx = np.linspace(min(x), max(x), 60)
    supporty = np.linspace(min(y), max(y), 60)
    X, Y = np.meshgrid(supportx, supporty)

    edat = np.column_stack([X.ravel(), Y.ravel()])
    Z = dens.pdf(edat).reshape(X.shape)

    # plot
    fig = plt.figure(1)
    ax = fig.gca(projection="3d")
    surf = ax.plot_surface(
        X, Y, Z, rstride=1, cstride=1, cmap=cm.jet, linewidth=0, antialiased=False
    )

    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.figure(2)
    plt.imshow(Z)

    plt.show()
