
import numpy as np
import matplotlib.pyplot as plt

from phylokrr.utils import k_fold_cv_random, split_data, PGLS
from phylokrr.datasets import load_vcv, sim_data
from phylokrr.kernels import KRR
np.random.seed(12038)




vcv = load_vcv()
n = vcv.shape[0]
mean_vector = np.zeros(n)
num_test = round(0.40 * n)

# simulate phylogenetically weighted data
X_w, y_w = sim_data(vcv, mean_vector, b=1, n_var=1)


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



rbf_grid_params = {
    'lambda': np.logspace(-15, 15, 200, dtype=float, base=2),
    'gamma' : np.logspace(-15, 15, 200, dtype=float, base=2),
}

krr_rbf = KRR(kernel='rbf', fit_intercept=True)
pgls_model = PGLS(fit_intercept =  True)

best_params = k_fold_cv_random(X_train, y_train,
                                krr_rbf,
                                rbf_grid_params,
                                verbose = False,
                                folds = 5,
                                sample = 100)

krr_rbf.set_params(**best_params)
krr_rbf.fit(X_train, y_train)
y_pred_cv = krr_rbf.predict(X_test)

pgls_model.fit(X_train, y_train)
y_pred_pgls = pgls_model.predict(X_test)



fs = 10
fig, axs = plt.subplots(1, 2, figsize=(8, 4))
axs[0].scatter(X_test, y_test, color = 'blue', 
               alpha=0.5, label = 'Testing (unseen) data',)
axs[0].scatter(X_test, y_pred_pgls, color = 'red', 
               alpha=0.5, label = 'PGLS predictions',)
axs[0].legend()
axs[1].scatter(X_test, y_test, color = 'blue', alpha=0.5, 
               label = 'Testing (unseen) data',)
axs[1].scatter(X_test, y_pred_cv, color = 'green', alpha=0.5, 
               label = 'RBF kernel predictions',)
axs[1].legend()


r2_kernel = krr_rbf.score( X_test, y_test, metric='r2')
r2_pgls = pgls_model.score(X_test, y_test, metric='r2')

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