import torch

import numpy as np

import matplotlib.pyplot as plt

import seaborn as sns

import sys

from gpsa import VariationalGPSA, matern12_kernel, rbf_kernel

from gpsa.plotting import callback_twod

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

from simulated.generate_twod_data import generate_twod_data

import matplotlib.animation as animation

import matplotlib.image as mpimg

import os

from os.path import join as pjoin

import anndata

import matplotlib

from matplotlib.lines import Line2D

font = {"size": 20}

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

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

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

LATEX_FONTSIZE = 35

n_spatial_dims = 2

n_views = 2

m_G = 50

m_X_per_view = 50

N_EPOCHS = 2000

PRINT_EVERY = 100

ONE_SAMPLE_FIXED = True

def two_d_gpsa(

    n_outputs,

    n_epochs,

    n_latent_gps,

    warp_kernel_variance=0.1,

    noise_variance=0.0,

    plot_intermediate=True,

    fixed_view_idx=None,

):

    X, Y, n_samples_list, view_idx = generate_twod_data(

        n_views,

        n_outputs,

        grid_size=10,

        n_latent_gps=n_latent_gps["expression"],

        kernel_lengthscale=5.0,

        kernel_variance=warp_kernel_variance,

        noise_variance=noise_variance,

        fixed_view_idx=0 if ONE_SAMPLE_FIXED else None,

    )

    n_samples_per_view = X.shape[0] // n_views

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

    markers = ["o", "X"]

    for vv in range(n_views):

        plt.scatter(

            X[view_idx[vv]][:, 0],

            X[view_idx[vv]][:, 1],

            c=Y[view_idx[vv]][:, 0],

            s=400,

            marker=markers[vv],

            label="View {}".format(vv + 1),

            edgecolor="black",

            linewidth=2,

        )

    plt.xlabel("X1")

    plt.ylabel("X2")

    plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))

    plt.tight_layout()

    plt.savefig("./../../examples/synthetic_data_example.png")

    plt.close()

    ## Save as anndata object

    data_obj = anndata.AnnData(Y)

    data_obj.obsm["spatial"] = X

    batch_id = np.concatenate([[xx] * n_samples_list[xx] for xx in range(n_views)])

    data_obj.obs["batch"] = batch_id

    data_obj.write("../../examples/synthetic_data.h5ad")

    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 = VariationalGPSA(

        data_dict,

        n_spatial_dims=n_spatial_dims,

        m_X_per_view=m_X_per_view,

        m_G=m_G,

        data_init=True,

        minmax_init=False,

        grid_init=False,

        n_latent_gps=n_latent_gps,

        mean_function="identity_fixed",

        kernel_func_warp=rbf_kernel,

        kernel_func_data=rbf_kernel,

        fixed_view_idx=fixed_view_idx,

    ).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

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

            {"expression": x}, view_idx=view_idx, Ns=Ns, S=5

        )

        # Compute loss

        loss = loss_fn(data_dict, F_samples)

        # Compute gradients and take optimizer step

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        return loss.item()

    # Set up figure.

    fig = plt.figure(figsize=(12.14, 5), facecolor="white", constrained_layout=True)

    data_expression_ax = fig.add_subplot(121, frameon=False)

    latent_expression_ax = fig.add_subplot(122, frameon=False)

    plt.show(block=False)

    loss_trace = []

    error_trace = []

    n_frames = 0

    for t in range(n_epochs):

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

        loss_trace.append(loss)

        if plot_intermediate and t % PRINT_EVERY == 0:

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

            G_means, G_samples, F_latent_samples, F_samples = 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=G_means,

                s=600,

            )

            legend_elements = [

                Line2D(

                    [0],

                    [0],

                    marker="o",

                    color="w",

                    label="Slice 1",

                    markerfacecolor="black",

                    markersize=20,

                ),

                Line2D(

                    [0],

                    [0],

                    marker="X",

                    color="w",

                    label="Slice 2",

                    markerfacecolor="black",

                    markersize=20,

                ),

            ]

            # Create the figure

            plt.legend(

                handles=legend_elements, loc="center left", bbox_to_anchor=(1, 0.5)

            )

            plt.tight_layout()

            # plt.draw()

            plt.savefig("./tmp/tmp{}".format(n_frames))

            n_frames += 1

            # plt.pause(1 / 60.0)

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

            {"expression": x}, view_idx=view_idx, Ns=Ns

        )

    print("Done!")

    plt.close()

    fig = plt.figure()

    ims = []

    for ii in range(n_frames):

        fname = "./tmp/tmp{}.png".format(ii)

        img = mpimg.imread(fname)

        im = plt.imshow(img)

        ax = plt.gca()

        ax.set_yticks([])

        ax.set_xticks([])

        ims.append([im])

        os.remove(fname)

    writervideo = animation.FFMpegWriter(fps=5)

    ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=500)

    if ONE_SAMPLE_FIXED:

        save_name = "alignment_animation_template.gif"

    else:

        save_name = "alignment_animation.gif"

    ani.save(

        pjoin("out", save_name),

        writer=writervideo,

        dpi=1000,

    )

    plt.close()

if __name__ == "__main__":

    n_outputs = 30

    two_d_gpsa(

        n_epochs=N_EPOCHS,

        n_outputs=n_outputs,

        warp_kernel_variance=0.5,

        noise_variance=0.001,

        n_latent_gps={"expression": 5},

        fixed_view_idx=0,

    )

    import ipdb

    ipdb.set_trace()