import torch
import numpy as np
import matplotlib.pyplot as plt

import seaborn as sns
import sys

sys.path.append("../..")
from models.gpsa_mle import WarpGPMLE

sys.path.append("../../data")
from simulated.generate_twod_data import generate_twod_data
from plotting.callbacks import callback_twod
from util import ConvergenceChecker


device = "cuda" if torch.cuda.is_available() else "cpu"

LATEX_FONTSIZE = 50

n_spatial_dims = 2
n_views = 2
# n_outputs = 10

N_EPOCHS = 3000
PRINT_EVERY = 25
N_LATENT_GPS = 1


def two_d_gpsa(n_outputs, n_epochs, warp_kernel_variance=0.1, plot_intermediate=True):

    X, Y, n_samples_list, view_idx = generate_twod_data(
        n_views,
        n_outputs,
        grid_size=15,
        n_latent_gps=None,
        kernel_lengthscale=10.0,
        kernel_variance=warp_kernel_variance,
        noise_variance=1e-4,
    )
    n_samples_per_view = X.shape[0] // n_views

    ## Fit GP on one view to get initial estimates of data kernel parameters
    from sklearn.gaussian_process.kernels import RBF, WhiteKernel
    from sklearn.gaussian_process import GaussianProcessRegressor

    kernel = RBF(length_scale=1.0) + WhiteKernel()
    gpr = GaussianProcessRegressor(kernel=kernel)
    gpr.fit(X[view_idx[0]], Y[view_idx[0]])
    data_lengthscales_est = gpr.kernel_.k1.theta[0]

    x = torch.from_numpy(X).float().clone()
    y = torch.from_numpy(Y).float().clone()

    data_dict = {
        "expression": {
            "spatial_coords": x,
            "outputs": y,
            "n_samples_list": n_samples_list,
        }
    }

    model = WarpGPMLE(
        data_dict,
        n_spatial_dims=n_spatial_dims,
        n_latent_gps=None,
        # n_latent_gps=None,
        mean_function="identity_fixed",
        fixed_warp_kernel_variances=np.ones(n_views) * 0.01,
        fixed_warp_kernel_lengthscales=np.ones(n_views) * 10,
        # fixed_data_kernel_lengthscales=np.exp(gpr.kernel_.k1.theta.astype(np.float32)),
        # fixed_data_kernel_lengthscales=np.exp(data_lengthscales_est),
        # mean_function="identity_initialized",
        fixed_view_idx=0,
    ).to(device)

    view_idx, Ns, _, _ = model.create_view_idx_dict(data_dict)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

    def train(model, loss_fn, optimizer):
        model.train()

        # Forward pass
        model.forward({"expression": x}, view_idx=view_idx, Ns=Ns)

        # Compute loss
        loss = loss_fn(
            X_spatial={"expression": x}, view_idx=view_idx, data_dict=data_dict
        )

        # Compute gradients and take optimizer step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        return loss.item()

    # Set up figure.
    fig = plt.figure(figsize=(14, 7), facecolor="white", constrained_layout=True)
    data_expression_ax = fig.add_subplot(122, frameon=False)
    latent_expression_ax = fig.add_subplot(121, frameon=False)
    plt.show(block=False)

    convergence_checker = ConvergenceChecker(span=100)

    loss_trace = []
    error_trace = []

    for t in range(n_epochs):
        loss = train(model, model.loss_fn, optimizer)
        loss_trace.append(loss)
        # print(model.G["expression"][-1])
        # print(torch.exp(model.warp_kernel_variances))
        if t >= convergence_checker.span - 1:
            rel_change = convergence_checker.relative_change(loss_trace)
            is_converged = convergence_checker.converged(loss_trace, tol=1e-4)
            if is_converged:
                convergence_counter += 1

                if convergence_counter == 2:
                    print("CONVERGED")
                    break

            else:
                convergence_counter = 0

        if plot_intermediate and t % PRINT_EVERY == 0:
            print("Iter: {0:<10} LL {1:1.3e}".format(t, -loss))
            model.forward({"expression": x}, view_idx=view_idx, Ns=Ns)

            callback_twod(
                model,
                X,
                Y,
                data_expression_ax=data_expression_ax,
                latent_expression_ax=latent_expression_ax,
                X_aligned=model.G,
                is_mle=True,
            )
            plt.draw()
            plt.pause(1 / 60.0)
            err = np.mean(
                (
                    model.G["expression"]
                    .detach()
                    .numpy()
                    .squeeze()[:n_samples_per_view]
                    - model.G["expression"]
                    .detach()
                    .numpy()
                    .squeeze()[n_samples_per_view:]
                )
                ** 2
            )
            print("Error: {}".format(err))

            if t >= convergence_checker.span - 1:
                print(rel_change)

        # G_means, G_samples, F_latent_samples, F_samples = model.forward(
        #     {"expression": x}, view_idx=view_idx, Ns=Ns
        # )

    print("Done!")

    plt.close()

    return X, Y, model.G, model


if __name__ == "__main__":

    n_outputs = 10
    X, Y, G_means, model = two_d_gpsa(n_epochs=N_EPOCHS, n_outputs=n_outputs)

    import matplotlib

    font = {"size": LATEX_FONTSIZE}
    matplotlib.rc("font", **font)
    matplotlib.rcParams["text.usetex"] = True

    fig = plt.figure(figsize=(10, 10))
    data_expression_ax = fig.add_subplot(211, frameon=False)
    latent_expression_ax = fig.add_subplot(212, frameon=False)
    callback_twod(
        model,
        X,
        Y,
        data_expression_ax=data_expression_ax,
        latent_expression_ax=latent_expression_ax,
        X_aligned=G_means,
    )

    plt.tight_layout()
    plt.savefig("../../plots/two_d_simulation.png")
    plt.show()

    import ipdb

    ipdb.set_trace()