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
