import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.gaussian_process.kernels import RBF
from scipy.stats import multivariate_normal as mvn

import seaborn as sns
import sys

sys.path.append("../../data")
from simulated.generate_twod_data import generate_twod_data
import os
from os.path import join as pjoin
import anndata

import matplotlib

font = {"size": 25}
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
markers = ["o", "X"]

lengthscale_list = [10**x for x in [-1, 0, 1]]
amplitude_list = [0.1, 1.0, 5.0]

xlimits = [0, 10]
ylimits = [0, 10]
grid_size = 10
x1s = np.linspace(*xlimits, num=grid_size)
x2s = np.linspace(*ylimits, num=grid_size)
X1, X2 = np.meshgrid(x1s, x2s)
X = np.vstack([X1.ravel(), X2.ravel()]).T
n = len(X)

## Lengthscales
plt.figure(figsize=(17, 15))
for ii, lengthscale in enumerate(lengthscale_list):
    for jj, amplitude in enumerate(amplitude_list):
        kernel = RBF(length_scale=lengthscale)
        K = amplitude * kernel(X) + 1e-8 * np.eye(n)
        X_warped = np.zeros((n, n_spatial_dims))
        for dd in range(n_spatial_dims):
            X_warped[:, dd] = mvn.rvs(mean=X[:, dd], cov=K)

        plt.subplot(
            len(amplitude_list),
            len(lengthscale_list),
            ii * len(lengthscale_list) + jj + 1,
        )
        plt.scatter(
            X[:, 0],
            X[:, 1],
            color="gray",
            marker=markers[0],
            label="Original" if (jj == len(amplitude_list) - 1) and (ii == 0) else None,
            s=50,
        )
        plt.scatter(
            X_warped[:, 0],
            X_warped[:, 1],
            color="red",
            marker=markers[1],
            label="Warped" if (jj == len(amplitude_list) - 1) and (ii == 0) else None,
            s=50,
        )
        plt.title(
            r"$\ell^2$ = "
            + str(round(lengthscale, 2))
            + ", $\sigma^2$ = "
            + str(round(amplitude, 2))
        )

        if (jj == len(amplitude_list) - 1) and (ii == 0):
            plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
        plt.xticks([])
        plt.yticks([])
plt.tight_layout()
plt.savefig("./out/warp_parameter_demo.png")
plt.show()
import ipdb

ipdb.set_trace()