import time
import matplotlib.pyplot as plt
import corner
import numpy as np
from askcarl.lightgmm import LightBaggingGMM
from askcarl.plot import plot_gmm_corner, plot_gmm_corner_pdf


def test_fastplot():
    for ndim, scale in (2, 2.0), (5, 1.5), (10, 1.0), (25, 0.8):
        N = 3000
        for i, (corr, M) in enumerate([(0, 30), (0, 300), (0.8, N), (-0.2, N)]):
            print(f"[{i}] corr={corr} M={M}")
            cov = np.eye(ndim)
            cov[cov == 0] = corr

            for D in 0, 3:
                # Example GMM
                X = np.vstack([
                    np.random.multivariate_normal(0 + np.arange(ndim), cov, abs(M) * ndim),
                    np.random.multivariate_normal(D + np.arange(ndim), cov, abs(M) * ndim)
                ])

                t0 = time.time()
                lgmm = LightBaggingGMM(n_gmms=20, n_components=10)
                lgmm.fit(X)
                gmm = lgmm.to_sklearn()
                print('  LightBaggingGMM fit:', time.time() - t0)

                t1 = time.time()
                plot_gmm_corner_pdf(
                        f'plotcorner_{"mono" if D == 0 else "dual"}_{ndim}d_N{abs(M)}_corr{i}.pdf',
                        gmm, levels=[0.393, 0.675, 0.864], scale=scale,
                        truths=np.arange(ndim), linewidth=0.5 * scale, truthlinewidth=1.0 * scale,
                        margin=80. / scale, fontname="Times-Roman"
                        )
                print('  GMM plot:', time.time() - t1)

def test_plot():
    ndim = 10
    print(f'ndim={ndim}')

    # Example GMM
    X = np.vstack([
        np.random.multivariate_normal(0 + np.arange(ndim), np.eye(ndim), 3000 * ndim),
        np.random.multivariate_normal(3 + np.arange(ndim), np.eye(ndim), 3000 * ndim)
    ])
    X[:,0] /= 1000
    X[:,1] /= 100
    X[:,-1] = X[:,-1] / 10 + 100
    X[:,-2] *= 3

    t0 = time.time()
    lgmm = LightBaggingGMM(n_gmms=20, n_components=10)
    lgmm.fit(X)
    print(f'LightBaggingGMM fit of {len(X)} data points: {time.time() - t0:.2f}s')

    gmm = lgmm.to_sklearn()
    t0 = time.time()
    plot_gmm_corner_pdf('plotcorner_custom.pdf', gmm, levels=[0.393, 0.675, 0.864], scale=1.0)
    print(f'custom corner pdf: {time.time() - t0:.2f}s')

    # Corner-style plot with confidence contours

    t0 = time.time()
    Y, _ = lgmm.sample(100000)
    fig = corner.corner(Y, plot_datapoints=False, plot_density=False, levels=[0.393, 0.675, 0.864])
    corner.corner(X, plot_datapoints=False, plot_density=False, levels=[0.393, 0.675, 0.864], fig=fig)
    plt.savefig('plotcorner1.pdf')
    plt.close()
    print(f'sample + corner: {time.time() - t0:.2f}s')

    t0 = time.time()
    fig, axes = plot_gmm_corner(gmm, levels=[0.393, 0.675, 0.864])
    plt.savefig('plotcorner.pdf')
    plt.savefig('plotcorner.pdf')
    plt.close()
    print(f'GMM plot: {time.time() - t0:.2f}s')
