import torch

import numpy as np

import matplotlib.pyplot as plt

from torch.utils.data import Dataset, DataLoader

import time

import pandas as pd

from scipy.stats import pearsonr

from matplotlib.lines import Line2D

import seaborn as sns

SCATTER_POINT_SIZE = 50

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

def callback_oned(

    model,

    X,

    Y,

    X_aligned,

    data_expression_ax,

    latent_expression_ax,

    prediction_ax=None,

    X_test=None,

    Y_pred=None,

    Y_test_true=None,

    X_test_aligned=None,

    F_samples=None,

):

    model.eval()

    markers = list(Line2D.markers.keys())

    colors = ["blue", "orange"]

    if model.fixed_view_idx is not None:

        curr_idx = model.view_idx["expression"][model.fixed_view_idx]

        X_aligned["expression"][curr_idx] = torch.tensor(X[curr_idx].astype(np.float32))

    data_expression_ax.cla()

    latent_expression_ax.cla()

    data_expression_ax.set_title("Observed data")

    latent_expression_ax.set_title("Aligned data")

    data_expression_ax.set_xlabel("Spatial coordinate")

    latent_expression_ax.set_xlabel("Spatial coordinate")

    data_expression_ax.set_ylabel("Outcome")

    latent_expression_ax.set_ylabel("Outcome")

    data_expression_ax.set_xlim([X.min(), X.max()])

    latent_expression_ax.set_xlim([X.min(), X.max()])

    for vv in range(model.n_views):

        view_idx = model.view_idx["expression"]

        data_expression_ax.scatter(

            X[view_idx[vv], 0],

            Y[view_idx[vv], 0],

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

            marker=markers[vv],

            s=SCATTER_POINT_SIZE,

            c="blue",

        )

        if Y.shape[1] > 1:

            data_expression_ax.scatter(

                X[view_idx[vv], 0],

                Y[view_idx[vv], 1],

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

                marker=markers[vv],

                s=SCATTER_POINT_SIZE,

                c="orange",

            )

        latent_expression_ax.scatter(

            # model.G_means["expression"].detach().cpu().numpy()[view_idx[vv], 0],

            X_aligned["expression"].detach().cpu().numpy()[view_idx[vv], 0],

            Y[view_idx[vv], 0],

            c="blue",

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

            marker=markers[vv],

            s=SCATTER_POINT_SIZE,

        )

        if Y.shape[1] > 1:

            latent_expression_ax.scatter(

                # model.G_means["expression"].detach().cpu().numpy()[view_idx[vv], 0],

                X_aligned["expression"].detach().cpu().numpy()[view_idx[vv], 0],

                Y[view_idx[vv], 1],

                c="orange",

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

                marker=markers[vv],

                s=SCATTER_POINT_SIZE,

            )

        # latent_expression_ax.scatter(

        #   model.Xtilde.detach().cpu().numpy()[vv, :, 0],

        #   model.delta_list.detach().cpu().numpy()[vv][:, 0],

        #   c="red",

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

        #   marker="^",

        #   s=100,

        # )

        if F_samples is not None:

            latent_expression_ax.scatter(

                X_aligned["expression"].detach().cpu().numpy()[view_idx[vv], 0],

                F_samples.detach().cpu().numpy()[view_idx[vv], 0],

                c="red",

                marker=markers[vv],

                s=SCATTER_POINT_SIZE,

            )

            if Y.shape[1] > 1:

                latent_expression_ax.scatter(

                    X_aligned["expression"].detach().cpu().numpy()[view_idx[vv], 0],

                    F_samples.detach().cpu().numpy()[view_idx[vv], 1],

                    c="green",

                    marker=markers[vv],

                    s=SCATTER_POINT_SIZE,

                )

    if prediction_ax is not None:

        prediction_ax.cla()

        prediction_ax.set_title("Predictions")

        prediction_ax.set_xlabel("True outcome")

        prediction_ax.set_ylabel("Predicted outcome")

        ### Plots the warping function

        # prediction_ax.scatter(

        #   X[view_idx[vv], 0],

        #   X_aligned["expression"].detach().cpu().numpy()[view_idx[vv], 0],

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

        #   marker=markers[vv],

        #   s=100,

        #   c="blue",

        # )

        # prediction_ax.scatter(

        #   model.Xtilde.detach().cpu().numpy()[vv, :, 0],

        #   model.delta_list.detach().cpu().numpy()[vv][:, 0],

        #   c="red",

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

        #   marker="^",

        #   s=100,

        # )

        latent_expression_ax.scatter(

            X_test_aligned["expression"].detach().cpu().numpy()[:, 0],

            Y_pred.detach().cpu().numpy()[:, 0],

            c="blue",

            label="Prediction",

            marker="^",

            s=SCATTER_POINT_SIZE,

        )

        latent_expression_ax.scatter(

            X_test_aligned["expression"].detach().cpu().numpy()[:, 0],

            Y_pred.detach().cpu().numpy()[:, 1],

            c="orange",

            label="Prediction",

            marker="^",

            s=SCATTER_POINT_SIZE,

        )

        prediction_ax.scatter(

            Y_test_true[:, 0],

            Y_pred.detach().cpu().numpy()[:, 0],

            c="black",

            s=SCATTER_POINT_SIZE,

        )

        prediction_ax.scatter(

            Y_test_true[:, 1],

            Y_pred.detach().cpu().numpy()[:, 1],

            c="black",

            s=SCATTER_POINT_SIZE,

            marker="^",

        )

    data_expression_ax.legend()

    plt.draw()

    plt.pause(1 / 60.0)

def callback_twod(

    model,

    X,

    Y,

    X_aligned,

    data_expression_ax,

    latent_expression_ax,

    is_mle=False,

    gene_idx=0,

    s=200,

    include_legend=False,

):

    if model.fixed_view_idx is not None:

        if is_mle:

            pass

        else:

            curr_idx = model.view_idx["expression"][model.fixed_view_idx]

            X_aligned["expression"][curr_idx] = torch.tensor(

                X[curr_idx].astype(np.float32), device=device

            )

    model.eval()

    markers = [".", "+", "^"]

    colors = ["blue", "orange"]

    data_expression_ax.cla()

    latent_expression_ax.cla()

    data_expression_ax.set_title("Observed data")

    latent_expression_ax.set_title("Aligned data")

    curr_view_idx = model.view_idx["expression"]

    latent_Xs = []

    Xs = []

    Ys = []

    markers_list = []

    viewname_list = []

    for vv in range(model.n_views):

        ## Data

        Xs.append(X[curr_view_idx[vv]])

        ## Latents

        curr_latent_Xs = X_aligned["expression"].detach().cpu().numpy()[curr_view_idx[vv]]

        latent_Xs.append(curr_latent_Xs)

        Ys.append(Y[curr_view_idx[vv], gene_idx])

        markers_list.append([markers[vv]] * curr_latent_Xs.shape[0])

        viewname_list.append(

            ["Observation {}".format(vv + 1)] * curr_latent_Xs.shape[0]

        )

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

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

    Ys = np.concatenate(Ys)

    markers_list = np.concatenate(markers_list)

    viewname_list = np.concatenate(viewname_list)

    data_df = pd.DataFrame(

        {

            "X1": Xs[:, 0],

            "X2": Xs[:, 1],

            "Y": Ys,

            "marker": markers_list,

            "view": viewname_list,

        }

    )

    latent_df = pd.DataFrame(

        {

            "X1": latent_Xs[:, 0],

            "X2": latent_Xs[:, 1],

            "Y": Ys,

            "marker": markers_list,

            "view": viewname_list,

        }

    )

    plt.sca(data_expression_ax)

    g = sns.scatterplot(

        data=data_df,

        x="X1",

        y="X2",

        hue="Y",

        style="view",

        ax=data_expression_ax,

        s=s,

        linewidth=1.8,

        edgecolor="black",

        palette="viridis",

    )

    if not include_legend:

        g.legend_.remove()

    # plt.colorbar()

    # plt.axis("off")

    # plt.scatter(model.Xtilde.detach().cpu().numpy()[0, :, 0], model.Xtilde.detach().cpu().numpy()[0, :, 1], color="red")

    # plt.scatter(model.Xtilde.detach().cpu().numpy()[1, :, 0], model.Xtilde.detach().cpu().numpy()[1, :, 1], color="red")

    # plt.scatter(model.Gtilde.detach().cpu().numpy()[:, 0], model.Gtilde.detach().cpu().numpy()[:, 1], color="red")

    # plt.axis("off")

    plt.sca(latent_expression_ax)

    g = sns.scatterplot(

        data=latent_df,

        x="X1",

        y="X2",

        hue="Y",

        style="view",

        ax=latent_expression_ax,

        s=s,

        linewidth=1.8,

        edgecolor="black",

        palette="viridis",

    )

    if not include_legend:

        g.legend_.remove()

    # plt.colorbar()

    # import ipdb; ipdb.set_trace()

    # for vv in range(model.n_views):

    #     # import ipdb; ipdb.set_trace()

    #     data_expression_ax.scatter(

    #         X[curr_view_idx[vv], 0],

    #         X[curr_view_idx[vv], 1],

    #         c=Y[curr_view_idx[vv], 0],

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

    #         marker=markers[vv],

    #         s=400,

    #     )

    # latent_expression_ax.scatter(

    #     X_aligned["expression"].detach().cpu().numpy()[curr_view_idx[vv], 0],

    #     X_aligned["expression"].detach().cpu().numpy()[curr_view_idx[vv], 1],

    #     c=Y[curr_view_idx[vv], 0],

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

    #     marker=markers[vv],

    #     s=400,

    # )

    # plt.axis("off")

def callback_twod_aligned_only(

    model,

    X,

    Y,

    X_aligned,

    latent_expression_ax1,

    latent_expression_ax2,

    is_mle=False,

    gene_idx=0,

):

    if model.fixed_view_idx is not None:

        if is_mle:

            pass

        else:

            curr_idx = model.view_idx["expression"][model.fixed_view_idx]

            X_aligned["expression"][curr_idx] = torch.tensor(

                X[curr_idx].astype(np.float32)

            )

    model.eval()

    markers = [".", "+", "^"]

    colors = ["blue", "orange"]

    latent_expression_ax1.cla()

    latent_expression_ax2.cla()

    latent_expression_ax1.set_title("Observed data")

    latent_expression_ax2.set_title("Aligned data")

    curr_view_idx = model.view_idx["expression"]

    latent_Xs = []

    Xs = []

    Ys = []

    markers_list = []

    viewname_list = []

    aligned_coords = X_aligned["expression"].detach().cpu().numpy()

    for vv in range(model.n_views):

        ## Data

        Xs.append(X[curr_view_idx[vv]])

        ## Latents

        curr_latent_Xs = aligned_coords[curr_view_idx[vv]]

        latent_Xs.append(curr_latent_Xs)

        Ys.append(Y[curr_view_idx[vv], gene_idx])

        markers_list.append([markers[vv]] * curr_latent_Xs.shape[0])

        viewname_list.append(["View {}".format(vv + 1)] * curr_latent_Xs.shape[0])

    latent_expression_ax1.scatter(

        aligned_coords[curr_view_idx[0]][:, 0],

        aligned_coords[curr_view_idx[0]][:, 1],

        c=Y[curr_view_idx[0]][:, gene_idx].squeeze(),

        s=24,

        marker="h",

    )

    latent_expression_ax2.scatter(

        aligned_coords[curr_view_idx[1]][:, 0],

        aligned_coords[curr_view_idx[1]][:, 1],

        c=Y[curr_view_idx[1]][:, gene_idx].squeeze(),

        s=24,

        marker="h",

    )

    # latent_expression_ax1.scatter(model.Xtilde.detach().cpu().numpy()[0, :, 0], model.Xtilde.detach().cpu().numpy()[0, :, 1], color="red")

    # latent_expression_ax2.scatter(model.Xtilde.detach().cpu().numpy()[1, :, 0], model.Xtilde.detach().cpu().numpy()[1, :, 1], color="red")

    plt.axis("off")

def callback_twod_multimodal(

    model, data_dict, X_aligned, axes, rgb=False, scatterpoint_size=100

):

    # if model.fixed_view_idx is not None:

    #     if is_mle:

    #         pass

    #     else:

    #         curr_idx = model.view_idx["expression"][model.fixed_view_idx]

    #         X_aligned["expression"][curr_idx] = torch.tensor(X[curr_idx].astype(np.float32))

    model.eval()

    markers = [".", "+", "^"]

    colors = ["blue", "orange"]

    [ax.cla() for ax in axes]

    axes[0].set_title("Observed expression")

    axes[1].set_title("Aligned expression")

    axes[2].set_title("Observed histology")

    axes[3].set_title("Aligned histology")

    axis_counter = 0

    n_mods = 2

    for mod in ["expression", "histology"]:

        curr_view_idx = model.view_idx[mod]

        for vv in range(model.n_views):

            # import ipdb; ipdb.set_trace()

            curr_coords = data_dict[mod]["spatial_coords"]

            if mod == "histology" and rgb:

                curr_outputs = data_dict[mod]["outputs"][curr_view_idx[vv], :]

            else:

                curr_outputs = data_dict[mod]["outputs"][curr_view_idx[vv], 0]

            axes[axis_counter].scatter(

                curr_coords[curr_view_idx[vv], 0],

                curr_coords[curr_view_idx[vv], 1],

                c=curr_outputs,

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

                marker=markers[vv],

                s=scatterpoint_size,

            )

            axes[axis_counter + 1].scatter(

                X_aligned[mod].detach().cpu().numpy()[curr_view_idx[vv], 0],

                X_aligned[mod].detach().cpu().numpy()[curr_view_idx[vv], 1],

                c=curr_outputs,

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

                marker=markers[vv],

                s=scatterpoint_size,

            )

        axis_counter += n_mods