import numpy as np
import matplotlib.pyplot as plt


from phylokrr.datasets import load_vcv
from phylokrr.utils import sim_data, split_data
from phylokrr.inspection import pdp_1d, pdp_2d, permutation_importance



np.random.seed(12037)

vcv = load_vcv()
mean_vector = np.zeros(vcv.shape[0])
b = 2
weights = [1,2,3]



# this function already produces the weighted data
X_w_uc, y_w_uc =  sim_data(vcv, mean_vector, b, extra_weights=weights)



n,_ = X_w_uc.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_uc, y_w_uc, num_test, seed = 12038) # seed defined above




from phylokrr.kernels import KRR
# set model
pkrr_model = KRR(kernel='linear', fit_intercept= True)


from phylokrr.utils import k_fold_cv_random

params = {
    'lambda' : np.logspace(-15, 15, 200, base=2),
    'c'  : np.logspace(-15, 15, 200,  base=2),
}

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

pkrr_model.set_params(**best_params)
pkrr_model.fit(X_train, y_train)

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



# permutation importance
PFI = permutation_importance(X_test, y_test, pkrr_model, num_test = 50, iterations=100)
np.savetxt('../../../writings/Kernels/revisions/pfi_linear.csv', PFI, delimiter=',', header='Trait1,Trait2,Trait3', comments='')
  

# plot bar chart
# adjust size of the plot
# plt.figure(figsize=(4,3))
# plt.bar(range(3), np.mean(PFI, axis=0), color='gray')
# plt.errorbar(range(3), np.mean(PFI, axis=0), np.std(PFI, axis=0), fmt='.', color='black')
# plt.xticks(range(3), ['Trait 1', 'Trait 2', 'Trait 3'])
# plt.ylabel('RMSE increase')
# plt.title('Permutation Feature Importance')
# plt.tight_layout()
# plt.savefig('../../../writings/Kernels/revisions/pfi_linear.pdf')




pdp_1d_f1 = pdp_1d(X_test, 0, pkrr_model, sample=100)
pdp_1d_f2 = pdp_1d(X_test, 1, pkrr_model, sample=100)
pdp_1d_f3 = pdp_1d(X_test, 2, pkrr_model, sample=100)

np.savetxt('../../../writings/Kernels/revisions/pdp_f1_linear.csv', pdp_1d_f1, delimiter=',', header='Trait1,PDvalue', comments='')
np.savetxt('../../../writings/Kernels/revisions/pdp_f2_linear.csv', pdp_1d_f2, delimiter=',', header='Trait2,PDvalue', comments='')
np.savetxt('../../../writings/Kernels/revisions/pdp_f3_linear.csv', pdp_1d_f3, delimiter=',', header='Trait3,PDvalue', comments='')





# all_pdp = [pdp_1d_f1, pdp_1d_f2, pdp_1d_f3]

# # get the y limits 
# y_lims = []
# for i in range(3):
#     y_lims.append(np.max(all_pdp[i][:,1]))
#     y_lims.append(np.min(all_pdp[i][:,1]))
# y_lims = [np.min(y_lims), np.max(y_lims)]



# fs = 18
# fig, ax = plt.subplots(2,3, figsize=(13/1.5, 7.5/1.5))
# # make a subplot with 1 row and 3 columns
# # fig, ax = plt.subplots(1,3, figsize=(13,5))

# # add a 'A' label to the subplot
# ax[0,0].text(-0.1, 1.1, 'A', transform=ax[0,0].transAxes, 
#              fontsize=fs, fontweight='bold', va='top', ha='right')
# for i in range(3):
#     tmp_pdp = all_pdp[i]
#     ax[0,i].plot(tmp_pdp[:,0], tmp_pdp[:,1])    
#     # ax[i].set_title(f'Trait {i+1}')
#     ax[0,i].set_xlabel(f'Trait {i+1}', size=fs)
#     ax[0,i].set_ylabel('PD value', size=fs)
#     ax[0,i].grid(True)
#     ax[0,i].set_ylim(y_lims)
#     # make space beteween subplots
#     fig.subplots_adjust(wspace=0.3)



# # make space beteween the two plots
# fig.subplots_adjust(hspace=0.3)

# # add a 'B' label to the subplot
# ax[1,0].text(-0.1, 1.1, 'B', transform=ax[1,0].transAxes, 
#              fontsize=fs, fontweight='bold', va='top', ha='right')

# 2d partial dependence plot
feature_combinations = [(0,1), (0,2), (1,2)]
# fig,ax=plt.subplots(1,3, figsize=(13,5))
for i,comb in enumerate(feature_combinations):
    X1, X2, Y = pdp_2d(X_test, comb[0], comb[1], pkrr_model, sample = 50) # 20

    out_file = f'../../../writings/Kernels/revisions/pdp_f{comb[0]}_f{comb[1]}_linear.csv'
    np.savetxt(out_file, np.c_[X1.flatten(), X2.flatten(), Y.flatten()], delimiter=',', header='Trait1,Trait2,PDvalue', comments='')


#     cp = ax[1,i].contourf(X1, X2, Y)
#     fig.colorbar(cp, ax=ax[1,i]) # Add a colorbar to a plot
#     # ax[i].set_title('2D ')
#     ax[1,i].set_xlabel(f'Trait {comb[0] + 1}', size=fs)
#     ax[1,i].set_ylabel(f'Trait {comb[1] + 1}', size=fs)
#     # ax[i].grid(True)
# plt.tight_layout()
# plt.savefig('../../../writings/Kernels/revisions/inspection.pdf')
