import numpy as np

import matplotlib.pyplot as plt

import pandas as pd

import seaborn as sns

import sys

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

# from gp_functions import rbf_covariance

from gpsa.util import rbf_kernel_numpy as rbf_covariance

from gpsa import polar_warp

from scipy.stats import multivariate_normal as mvnpy

def apply_gp_warp(

    X_orig_single,

    Y_orig_single,

    n_views,

    # n_samples_per_view,

    noise_variance=0.0,

    kernel_variance=1.0,

    kernel_lengthscale=1.0,

    mean_slope=1.0,

    mean_intercept=0.0,

):

    n_samples_per_view = X_orig_single.shape[0]

    n_spatial_dims = X_orig_single.shape[1]

    kernel = rbf_covariance

    kernel_params_true = np.array([np.log(1.0), np.log(1.0)])

    X_orig = np.concatenate([X_orig_single.copy()] * n_views, axis=0)

    n_samples_list = [n_samples_per_view] * n_views

    cumulative_sums = np.cumsum(n_samples_list)

    cumulative_sums = np.insert(cumulative_sums, 0, 0)

    view_idx = np.array(

        [

            np.arange(cumulative_sums[ii], cumulative_sums[ii + 1])

            for ii in range(n_views)

        ]

    )

    n = np.sum(n_samples_list)

    X = X_orig.copy()

    # Draw warped coordinates from a GP

    warp_kernel_params_true = np.array(

        [np.log(kernel_variance), np.log(kernel_lengthscale)]

    )

    for vv in range(n_views):

        for ss in range(n_spatial_dims):

            X_curr_view_warped = mvnpy.rvs(

                mean=X_orig_single[:, ss] * mean_slope + mean_intercept,

                cov=kernel(X_orig_single, X_orig_single, warp_kernel_params_true),

            )

            # import ipdb; ipdb.set_trace()

            X[

                n_samples_per_view * vv : n_samples_per_view * (vv + 1), ss

            ] = X_curr_view_warped

    Y = np.concatenate([Y_orig_single] * n_views, axis=0)

    Y += np.random.normal(scale=np.sqrt(noise_variance), size=Y.shape)

    return X, Y, n_samples_list, view_idx

def apply_gp_warp_multimodal(

    X_orig_singles,

    Y_orig_singles,

    n_views,

    # n_samples_per_view,

    noise_variance=0.0,

    kernel_variance=1.0,

    kernel_lengthscale=1.0,

    mean_slope=1.0,

    mean_intercept=0.0,

):

    assert len(X_orig_singles) == len(Y_orig_singles)

    n_modalities = len(X_orig_singles)

    modality_idx = np.cumsum([x.shape[0] for x in X_orig_singles])

    modality_idx = np.insert(modality_idx, 0, 0)

    kernel = rbf_covariance

    kernel_params_true = np.array([np.log(1.0), np.log(1.0)])

    X_orig_single = np.concatenate(X_orig_singles, axis=0)

    X_orig_single = X_orig_single - X_orig_single.min(0)

    X_orig_single = X_orig_single / X_orig_single.max(0)

    X_orig_single *= 10

    X_orig = np.concatenate([X_orig_single.copy()] * n_views, axis=0)

    n_samples_per_view = X_orig_single.shape[0]

    n_spatial_dims = X_orig_single.shape[1]

    n_samples_list = [n_samples_per_view] * n_views

    cumulative_sums = np.cumsum(n_samples_list)

    cumulative_sums = np.insert(cumulative_sums, 0, 0)

    view_idx = np.array(

        [

            np.arange(cumulative_sums[ii], cumulative_sums[ii + 1])

            for ii in range(n_views)

        ]

    )

    n = np.sum(n_samples_list)

    X = X_orig.copy()

    # Draw warped coordinates from a GP

    warp_kernel_params_true = np.array(

        [np.log(kernel_variance), np.log(kernel_lengthscale)]

    )

    for vv in range(n_views):

        curr_idx = np.arange(n_samples_per_view * vv, n_samples_per_view * (vv + 1))

        for ss in range(n_spatial_dims):

            X_curr_view_warped = mvnpy.rvs(

                mean=X_orig_single[:, ss] * mean_slope + mean_intercept,

                cov=kernel(X_orig_single, X_orig_single, warp_kernel_params_true),

            )

            # import ipdb; ipdb.set_trace()

            X[curr_idx, ss] = X_curr_view_warped

    view_idx = np.cumsum([n_samples_per_view * vv for vv in range(n_views + 1)])

    X_warped = []

    Y_warped = []

    n_samples_list = []

    for mm in range(n_modalities):

        curr_modality_idx = np.concatenate(

            [

                view_idx[vv] + np.arange(modality_idx[mm], modality_idx[mm + 1])

                for vv in range(n_views)

            ]

        )

        X_warped.append(X[curr_modality_idx])

        Y_full_mm = np.concatenate([Y_orig_singles[mm]] * n_views, axis=0)

        Y_full_mm += np.random.normal(

            scale=np.sqrt(noise_variance), size=Y_full_mm.shape

        )

        Y_warped.append(Y_full_mm)

        n_samples_list.append([X_orig_singles[mm].shape[0]] * n_views)

    return X_warped, Y_warped, n_samples_list, view_idx

def apply_linear_warp(

    X_orig_single,

    Y_orig_single,

    n_views,

    linear_slope_variance=0.1,

    linear_intercept_variance=0.1,

    noise_variance=0.01,

    rotation=True,

):

    n_samples_per_view = X_orig_single.shape[0]

    n_spatial_dims = X_orig_single.shape[1]

    kernel = rbf_covariance

    kernel_params_true = np.array([np.log(1.0), np.log(1.0)])

    X_orig = np.concatenate([X_orig_single.copy()] * n_views, axis=0)

    n_samples_list = [n_samples_per_view] * n_views

    cumulative_sums = np.cumsum(n_samples_list)

    cumulative_sums = np.insert(cumulative_sums, 0, 0)

    view_idx = np.array(

        [

            np.arange(cumulative_sums[ii], cumulative_sums[ii + 1])

            for ii in range(n_views)

        ]

    )

    n = np.sum(n_samples_list)

    X = X_orig.copy()

    for vv in range(n_views):

        # curr_slopes = np.eye(n_spatial_dims) + np.random.normal(

        #     loc=0,

        #     scale=np.sqrt(linear_slope_variance),

        #     size=(n_spatial_dims, n_spatial_dims),

        # )

        # import ipdb; ipdb.set_trace()

        # curr_slopes /= np.linalg.norm(curr_slopes, ord=2, axis=0)

        # curr_slopes = np.linalg.svd(curr_slopes)[0]

        # curr_slopes = np.random.normal(

        #     loc=1,

        #     scale=np.sqrt(linear_slope_variance),

        #     size=n_spatial_dims,

        # )

        # curr_intercepts = np.random.normal(

        #     loc=0, scale=np.sqrt(linear_intercept_variance), size=n_spatial_dims

        # )

        curr_slopes = np.random.uniform(

            low=1 - linear_slope_variance,

            high=1 + linear_slope_variance,

            size=n_spatial_dims,

        )

        curr_intercepts = np.random.uniform(

            low=linear_intercept_variance,

            high=linear_intercept_variance,

            size=n_spatial_dims,

        )

        # print(curr_slopes, curr_intercepts)

        X_curr_view_warped = X_orig_single * curr_slopes + curr_intercepts

        X[

            n_samples_per_view * vv : n_samples_per_view * (vv + 1), :

        ] = X_curr_view_warped

    Y = np.concatenate([Y_orig_single] * n_views, axis=0)

    Y += np.random.normal(scale=np.sqrt(noise_variance), size=Y.shape)

    return X, Y, n_samples_list, view_idx

def apply_polar_warp(

    X_orig_single,

    Y_orig_single,

    n_views,

    linear_slope_variance=0.1,

    linear_intercept_variance=0.1,

    noise_variance=0.01,

    rotation=True,

):

    n_samples_per_view = X_orig_single.shape[0]

    n_spatial_dims = X_orig_single.shape[1]

    kernel = rbf_covariance

    kernel_params_true = np.array([np.log(1.0), np.log(1.0)])

    X_orig = np.concatenate([X_orig_single.copy()] * n_views, axis=0)

    n_samples_list = [n_samples_per_view] * n_views

    cumulative_sums = np.cumsum(n_samples_list)

    cumulative_sums = np.insert(cumulative_sums, 0, 0)

    view_idx = np.array(

        [

            np.arange(cumulative_sums[ii], cumulative_sums[ii + 1])

            for ii in range(n_views)

        ]

    )

    n = np.sum(n_samples_list)

    X = X_orig.copy()

    # no_distortion_B = np.array([

    #         [0, np.pi * 0.5],

    #         [0, 0]

    #     ])

    for vv in range(n_views):

        # B = np.random.normal(

        #     loc=0,

        #     scale=np.sqrt(linear_slope_variance),

        #     size=(n_spatial_dims, n_spatial_dims),

        # )

        B = np.random.uniform(

            low=-linear_slope_variance,

            high=linear_slope_variance,

            size=(n_spatial_dims, n_spatial_dims),

        )

        polar_params = X_orig_single @ B

        r, theta = polar_params[:, 0], polar_params[:, 1]

        X_curr_view_warped = np.array(

            [

                X_orig_single[:, 0] + r * np.cos(theta),

                X_orig_single[:, 1] + r * np.sin(theta),

            ]

        ).T

        # import ipdb; ipdb.set_trace()

        # additive_warp = B[:, 0] * X_orig_single * np.vstack([np.cos(X_orig_single[:, 0] * B[0, 1]), np.sin(X_orig_single[:, 1] * B[1, 1])]).T

        # X_curr_view_warped = X_orig_single + additive_warp

        # X_curr_view_warped = X_orig_single @ curr_slopes + curr_intercepts

        X[

            n_samples_per_view * vv : n_samples_per_view * (vv + 1), :

        ] = X_curr_view_warped

    Y = np.concatenate([Y_orig_single] * n_views, axis=0)

    Y += np.random.normal(scale=np.sqrt(noise_variance), size=Y.shape)

    return X, Y, n_samples_list, view_idx

if __name__ == "__main__":

    pass