
import numpy as np
import matplotlib.pyplot as plt


from phylokrr.kernels import KRR
from phylokrr.datasets import load_1d_data_example
from phylokrr.utils import weight_data

X_uw_uc,y_uw_uc,vcv = load_1d_data_example(return_cov=True)

# weighting the data by using the variance covariance matrix
X_w, y_w = weight_data(X_uw_uc, y_uw_uc, vcv, use_sd = False)
# second column is the absolute value of the first column
y_w = np.stack([y_w, np.abs(y_w)], axis = 1) 

# plotting data



# # region
fig, axs = plt.subplots(1, 2, figsize=(8, 4))
axs[0].scatter(X_w, y_w[:,0], color = 'blue', alpha=0.5, )
axs[0].set_title('Phylogenetically weighted data')
axs[0].set_xlabel('$x^*$')
axs[0].set_ylabel('$y^*_1$')

axs[1].scatter(X_w, y_w[:,1], color = 'blue', alpha=0.5, )
axs[1].set_title('Phylogenetically weighted data')
axs[1].set_xlabel('$x^*$')
axs[1].set_ylabel('$y^*_2$')
plt.tight_layout()
# endregion

"""## Simple model fitting without Cross-Validation (CV)"""

from phylokrr.utils import split_data

n,_ = X_w.shape
# split data into training and testing sets
num_test = round(0.5*n)

(X_train, X_test,
 y_train, y_test,) = split_data(X_w, y_w, num_test, seed = 12038) # seed defined above


# from phylokrr.kernels import KRR

#region #
from phylokrr.utils import k_fold_cv_random
# set model
pkrr_model = KRR(kernel='rbf', fit_intercept= True)

params = {
    'lambda' : np.logspace(-5, 10, 200, base=2),
    'gamma' : np.logspace(-5, 10, 200,  base=2),
}

best_params = k_fold_cv_random(X_train, y_train,
                                pkrr_model,
                                params,
                                verbose = False,
                                folds = 5,
                                sample = 100)

pkrr_model.set_params(**best_params)

# fit model with phylogenetic covariance matrix
pkrr_model.fit(X_train, y_train,)
# pkrr_model.alpha
y_pred_kernel = pkrr_model.predict(X_test)

"""Let's compare it with the standard phylogenetic regression (i.e., PGLS)"""

from phylokrr.utils import PGLS
# fit standard phylogenetic regression
pgls = PGLS(fit_intercept=True)
pgls.fit(X_train, y_train)
y_pred_pgls = pgls.predict(X_test)
# store figure


subplots =2
fig, axs = plt.subplots(1, subplots, figsize=(9, 4))
axs[0].scatter(X_test, y_test[:,0], color = 'blue' , alpha=0.5, label = 'Testing (unseen) data')
axs[0].scatter(X_test, y_pred_kernel[:,0], color = 'green', alpha=0.5, label = 'phyloKRR predictions')
axs[0].scatter(X_test, y_pred_pgls[:,0], color = 'red', alpha=0.5, label = 'PGLS predictions')
axs[0].set_xlabel('$x^*$')
axs[0].set_ylabel('$y^*_1$')
axs[0].legend()
axs[1].scatter(X_test, y_test[:,1], color = 'blue' , alpha=0.5, label = 'Testing (unseen) data')
axs[1].scatter(X_test, y_pred_kernel[:,1], color = 'green', alpha=0.5, label = 'phyloKRR predictions')
axs[1].scatter(X_test, y_pred_pgls[:,1], color = 'red', alpha=0.5, label = 'PGLS predictions')
axs[1].set_xlabel('$x^*$')
axs[1].set_ylabel('$y^*_2$')
axs[1].legend()
plt.tight_layout()
# plt.savefig('../../../writings/Kernels/revisions/multioutput.pdf')

r2_kernel = pkrr_model.score(X_test, y_test, metric='r2')
r2_pgls = pgls.score(X_test, y_test, metric='r2')

print(f"R2 for phyloKRR: {r2_kernel}")
print(f"R2 for PGLS: {r2_pgls}")
# endregion ####





# # X_train.shape
# K_train = RBF_kernel(X_train, X_train, pkrr_model.params['gamma'])
# K_test  = RBF_kernel(X_test, X_train,  pkrr_model.params['gamma'])



# cats = 2
# # ones = np.ones((cats,cats))
# ones = np.array([
#     [1, 0], 
#     [1, 1]])
# # repeat K_train in cats by cats blocks
# K_train_aug = np.kron(ones, K_train)
# K_test_aug = np.kron(ones, K_test)

# # stack y_train in a vector colum
# y_train_flat = y_train.flatten('F')
# y_test_flat = y_test.flatten('F')

# pkrr_mlt = KRR(kernel='precomputed', fit_intercept= True)


# pkrr_mlt.params['lambda'] = pkrr_model.params['lambda']
# pkrr_mlt.fit(K_train_aug,y_train_flat)

# y_pred_mlt = pkrr_mlt.predict(K_test_aug)


# subplots =2
# fig, axs = plt.subplots(1, subplots, figsize=(8, 4))
# axs[0].scatter(X_test, y_test_flat[:num_test], color = 'blue' , alpha=0.5, label = 'Testing (unseen) data')
# axs[0].scatter(X_test, y_pred_mlt[:num_test], color = 'green', alpha=0.5, label = 'phyloKRR predictions w\o CV')
# axs[0].scatter(X_test, y_pred_pgls[:,0], color = 'red', alpha=0.5, label = 'PGLS predictions')
# axs[0].set_xlabel('$x^*$')
# axs[0].set_ylabel('$y^*_1$')
# axs[0].legend()
# axs[1].scatter(X_test, y_test_flat[num_test:], color = 'blue' , alpha=0.5, label = 'Testing (unseen) data')
# axs[1].scatter(X_test, y_pred_mlt[num_test:], color = 'green', alpha=0.5, label = 'phyloKRR predictions w\o CV')
# axs[1].scatter(X_test, y_pred_pgls[:,1], color = 'red', alpha=0.5, label = 'PGLS predictions')
# axs[1].set_xlabel('$x^*$')
# axs[1].set_ylabel('$y^*_2$')
# axs[1].legend()
# plt.tight_layout()



