import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

import seaborn as sns
import sys
from gpsa.util import rbf_kernel_numpy as rbf_covariance
from scipy.stats import multivariate_normal as mvnpy


def generate_oned_data_affine_warp(
    n_views,
    n_outputs,
    n_samples_per_view,
    noise_variance=0.0,
    n_latent_gps=None,
    scale_factor=1.1,
    additive_factor=0.3,
):

    kernel = rbf_covariance
    kernel_params_true = np.array([np.log(1.0), np.log(1.0)])
    n_latent_gps = 2
    n_spatial_dims = 1

    X_orig_single = np.random.uniform(-10, 10, size=(n_samples_per_view, 1))
    # X_orig_single = np.linspace(-10, 10, n_samples_per_view)[:, :-1]
    X_orig = np.concatenate([X_orig_single.copy(), X_orig_single.copy()], 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)

    nY = n_outputs if n_latent_gps is None else n_latent_gps

    Y_orig = np.vstack(
        [
            mvnpy.rvs(
                mean=np.zeros(X_orig_single.shape[0]),
                cov=kernel(X_orig_single, X_orig_single, kernel_params_true),
            )
            for _ in range(nY)
        ]
    ).T

    if n_latent_gps is not None:
        W_mat = np.random.normal(size=(n_latent_gps, n_outputs))
        # W_mat = np.expand_dims(np.array([1, -1]), 0)
        Y_orig = Y_orig @ W_mat

    Y = np.concatenate([Y_orig, Y_orig], axis=0)
    Y += np.random.normal(scale=np.sqrt(noise_variance), size=(Y.shape))
    X = X_orig.copy()
    X[n_samples_per_view:] = X[n_samples_per_view:] * scale_factor + additive_factor

    return X, Y, n_samples_list, view_idx


def generate_oned_data_gp_warp(
    n_views,
    n_outputs,
    n_samples_per_view,
    noise_variance=0.0,
    n_latent_gps=None,
    kernel_variance=1.0,
    kernel_lengthscale=1.0,
    mean_slope=1.,
    mean_intercept=0.,
):

    kernel = rbf_covariance
    kernel_params_true = np.array([np.log(1.0), np.log(1.0)])
    n_spatial_dims = 1

    # X_orig_single = np.random.uniform(-10, 10, size=(n_samples_per_view, 1))
    X_orig_single = np.linspace(-10, 10, n_samples_per_view).reshape(-1, 1)
    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)

    nY = n_outputs if n_latent_gps is None else n_latent_gps

    Y_orig = np.vstack(
        [
            mvnpy.rvs(
                mean=np.zeros(X_orig_single.shape[0]),
                cov=kernel(X_orig_single, X_orig_single, kernel_params_true),
            )
            for _ in range(nY)
        ]
    ).T

    if n_latent_gps is not None:
        if n_outputs == 2:
            W_mat = np.expand_dims(np.array([1, -1]), 0)
        else:
            W_mat = np.random.normal(size=(n_latent_gps, n_outputs))
        
        Y_orig = Y_orig @ W_mat

    Y = np.concatenate([Y_orig] * n_views, axis=0)
    Y += np.random.normal(scale=np.sqrt(noise_variance), size=(Y.shape))
    X = X_orig.copy()

    # X_view1 = X[:n_samples_per_view]
    # X_view2 = X[n_samples_per_view : n_samples_per_view * 2]

    # 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):
        X_curr_view_warped = mvnpy.rvs(
            mean=X_orig_single.squeeze() * mean_slope + mean_intercept,
            cov=kernel(X_orig_single, X_orig_single, warp_kernel_params_true),
        )
        X_curr_view_warped = np.expand_dims(X_curr_view_warped, 1)
        # import ipdb; ipdb.set_trace()
        X[n_samples_per_view * vv : n_samples_per_view * (vv + 1)] = X_curr_view_warped

    # X += np.random.normal(scale=np.sqrt(0.1), size=(X.shape))
    # X_view1_warped = mvnpy.rvs(
    #     mean=X_view2.squeeze() * mean_slope + mean_intercept,
    #     cov=kernel(X_view1, X_view1, warp_kernel_params_true),
    # )
    # X_view1_warped = np.expand_dims(X_view1_warped, 1)
    # X[n_samples_per_view : n_samples_per_view * 2] = X_view1_warped

    # X_view2_warped = mvnpy.rvs(
    #     mean=X_view2.squeeze() * mean_slope + mean_intercept,
    #     cov=kernel(X_view2, X_view2, warp_kernel_params_true),
    # )
    # X_view2_warped = np.expand_dims(X_view2_warped, 1)
    # X[n_samples_per_view : n_samples_per_view * 2] = X_view2_warped

    return X, Y, n_samples_list, view_idx


if __name__ == "__main__":

    n_views = 2
    n_samples_per_view = 100
    X, Y, n_samples_list, view_idx = generate_oned_data_gp_warp(
        n_views=n_views, n_outputs=1, n_samples_per_view=n_samples_per_view
    )

    for vv in range(n_views):
        curr_start_idx = vv * n_samples_per_view
        curr_end_idx = vv * n_samples_per_view + n_samples_per_view
        plt.scatter(X[curr_start_idx:curr_end_idx], Y[curr_start_idx:curr_end_idx])
    plt.show()

    import ipdb

    ipdb.set_trace()