import random
import numpy as np

from phylokrr.utils import P_inv_simple, P_mat_simple
from phylokrr.datasets import load_vcv

np.random.seed(12037)

vcv = load_vcv()
mean_vector = np.zeros(vcv.shape[0])

X_uw_uc = np.random.multivariate_normal(mean_vector + 4, vcv, 1).reshape(-1,1)
# Non-linear response variable (sine curve)
y_uw_uc = np.sin(X_uw_uc*1.5).reshape(-1,1)

# # Add noise to the response variable
# y_w_uc[::10] += 4 * (0.5 - np.random.rand(X_w_uc.shape[0] // 10))

# Add noise to the response variable
e = np.random.multivariate_normal(mean_vector, vcv, 1).reshape(-1,1)
# y_uw_uc[::2] += 2 * (0.5 + e[::2])
y_uw_uc += e

# we can attempt to unweight to the original space
# with the square root of the covariance matrix
# P_inv = P_inv_simple(vcv)
P = P_mat_simple(vcv)

X_w_uc, y_w_uc = P @ X_uw_uc, P @ y_uw_uc

import matplotlib.pyplot as plt
plt.scatter(X_uw_uc, y_uw_uc, color = 'blue', alpha=0.5, )
plt.title('Unweighted data')

plt.scatter(X_w_uc, y_w_uc, color = 'red', alpha=0.5, )