import torch

import numpy as np

import matplotlib.pyplot as plt

import pandas as pd

import seaborn as sns

import sys

sys.path.append("../..")

from models.gpsa_vi_lmc import VariationalWarpGP

from data.simulated.generate_oned_data import (

    generate_oned_data_affine_warp,

    generate_oned_data_gp_warp,

)

from plotting.callbacks import callback_oned

from sklearn.gaussian_process import GaussianProcessRegressor

from sklearn.gaussian_process.kernels import WhiteKernel, RBF

import matplotlib

LATEX_FONTSIZE = 30

font = {"size": LATEX_FONTSIZE}

matplotlib.rc("font", **font)

matplotlib.rcParams["text.usetex"] = True

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

n_spatial_dims = 1

n_views = 2

n_outputs = 10

n_samples_per_view = 100

m_G = 20

m_X_per_view = 20

N_EPOCHS = 2000

PRINT_EVERY = 100

N_LATENT_GPS = {"expression": 3}

NOISE_VARIANCE = 0.1

N_REPEATS = 2

errors_union, errors_separate, errors_gpsa = [], [], []

for repeat_idx in range(N_REPEATS):

    # X, Y, n_samples_list, view_idx = generate_oned_data_affine_warp(

    #     n_views,

    #     n_outputs,

    #     n_samples_per_view,

    #     noise_variance=NOISE_VARIANCE,

    #     n_latent_gps=N_LATENT_GPS,

    #     scale_factor=1.,

    #     additive_factor=1.0,

    # )

    X, Y, n_samples_list, view_idx = generate_oned_data_gp_warp(

        n_views,

        n_outputs,

        n_samples_per_view,

        noise_variance=NOISE_VARIANCE,

        n_latent_gps=N_LATENT_GPS["expression"],

        kernel_variance=0.25,

        kernel_lengthscale=5.0,

    )

    ## Drop part of the second view (this is the part we'll try to predict)

    second_view_idx = view_idx[1]

    n_drop = int(1.0 * n_samples_per_view // 2.0)

    test_idx = np.random.choice(second_view_idx, size=n_drop, replace=False)

    keep_idx = np.setdiff1d(second_view_idx, test_idx)

    train_idx = np.concatenate([np.arange(n_samples_per_view), keep_idx])

    X_train = X[train_idx]

    Y_train = Y[train_idx]

    n_samples_list_train = n_samples_list

    n_samples_list_train[1] -= n_drop

    n_samples_list_test = [[0], [n_drop]]

    X_test = X[test_idx]

    Y_test = Y[test_idx]

    x_train = torch.from_numpy(X_train).float().clone()

    y_train = torch.from_numpy(Y_train).float().clone()

    x_test = torch.from_numpy(X_test).float().clone()

    y_test = torch.from_numpy(Y_test).float().clone()

    data_dict_train = {

        "expression": {

            "spatial_coords": x_train,

            "outputs": y_train,

            "n_samples_list": n_samples_list_train,

        }

    }

    data_dict_test = {

        "expression": {

            "spatial_coords": x_test,

            "outputs": y_test,

            "n_samples_list": n_samples_list_test,

        }

    }

    model = VariationalWarpGP(

        data_dict_train,

        n_spatial_dims=n_spatial_dims,

        m_X_per_view=m_X_per_view,

        m_G=m_G,

        data_init=False,

        minmax_init=False,

        grid_init=True,

        n_latent_gps=N_LATENT_GPS,

        mean_function="identity_fixed",

        # fixed_kernel_variances=np.ones(n_views) * 2,

        # fixed_kernel_lengthscales=np.ones(n_views) * 2,

        # mean_function="identity_initialized",

        # fixed_view_idx=0,

    ).to(device)

    view_idx_train, Ns_train, _, _ = model.create_view_idx_dict(data_dict_train)

    view_idx_test, Ns_test, _, _ = model.create_view_idx_dict(data_dict_test)

    ## Make predictions for naive alignment

    gpr_union = GaussianProcessRegressor(kernel=RBF() + WhiteKernel())

    gpr_union.fit(X=X_train, y=Y_train)

    preds = gpr_union.predict(X_test)

    error_union = np.mean((preds - Y_test) ** 2)

    errors_union.append(error_union)

    print("MSE, union: {}".format(round(error_union, 5)))

    ## Make predictons for each view separately

    preds, truth = [], []

    for vv in range(n_views):

        gpr_separate = GaussianProcessRegressor(kernel=RBF() + WhiteKernel())

        curr_trainX = X_train[view_idx_train["expression"][vv]]

        curr_trainY = Y_train[view_idx_train["expression"][vv]]

        curr_testX = X_test[view_idx_test["expression"][vv]]

        curr_testY = Y_test[view_idx_test["expression"][vv]]

        if len(curr_testX) == 0:

            continue

        gpr_separate.fit(X=curr_trainX, y=curr_trainY)

        curr_preds = gpr_separate.predict(curr_testX)

        preds.append(curr_preds)

        truth.append(curr_testY)

    preds = np.concatenate(preds, axis=0)

    truth = np.concatenate(truth, axis=0)

    error_separate = np.mean((preds - truth) ** 2)

    errors_separate.append(error_separate)

    print("MSE, separate: {}".format(round(error_separate, 5)))

    # preds = gpr_union.predict(X_test)

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

    def train(model, loss_fn, optimizer):

        model.train()

        # Forward pass

        G_means, G_samples, F_latent_samples, F_samples = model.forward(

            X_spatial={"expression": x_train}, view_idx=view_idx_train, Ns=Ns_train

        )

        # Compute loss

        loss = loss_fn(data_dict_train, F_samples)

        # Compute gradients and take optimizer step

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        return loss.item(), G_means

    # Set up figure.

    fig = plt.figure(figsize=(18, 7), facecolor="white", constrained_layout=True)

    ax_dict = fig.subplot_mosaic(

        [

            ["data", "preds"],

            ["latent", "preds"],

        ],

    )

    plt.show(block=False)

    for t in range(N_EPOCHS):

        loss, G_means = train(model, model.loss_fn, optimizer)

        if t % PRINT_EVERY == 0:

            print("Iter: {0:<10} LL {1:1.3e}".format(t, -loss))

            G_means_test, _, F_samples_test, _, = model.forward(

                X_spatial={"expression": x_test},

                view_idx=view_idx_test,

                Ns=Ns_test,

                prediction_mode=True,

                S=10,

            )

            curr_preds = torch.mean(F_samples_test["expression"], dim=0)

            callback_oned(

                model,

                X_train,

                Y_train,

                data_expression_ax=ax_dict["data"],

                latent_expression_ax=ax_dict["latent"],

                prediction_ax=ax_dict["preds"],

                X_aligned=G_means,

                X_test=X_test,

                Y_test_true=Y_test,

                Y_pred=curr_preds,

                X_test_aligned=G_means_test,

            )

            error_gpsa = np.mean((Y_test - curr_preds.detach().numpy()) ** 2)

            print("MSE, GPSA: {}".format(round(error_gpsa, 5)))

    errors_gpsa.append(error_gpsa)

    plt.close()

results_df = pd.DataFrame(

    {"Union": errors_union, "Separate": errors_separate, "GPSA": errors_gpsa}

)

results_df_melted = pd.melt(results_df)

plt.figure(figsize=(7, 5))

sns.boxplot(data=results_df_melted, x="variable", y="value", color="gray")

plt.xlabel("")

plt.ylabel("MSE")

plt.tight_layout()

plt.savefig("../../plots/one_d_prediction_comparison.png")

plt.show()

import ipdb

ipdb.set_trace()

# 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_oned(model, X, Y, data_expression_ax, latent_expression_ax)

# plt.tight_layout()

# plt.savefig("../../plots/one_d_simulation.png")

# plt.show()

# import ipdb

# ipdb.set_trace()