
import numpy as np
from phylokrr.core import PhylogeneticRegressor


def dim_red(X, n_com):
    # extracting averages
    X_centered = (X - np.mean(X, axis=0))

    # U,E,Vt (SVD)
    _,_,Vt = np.linalg.svd(X_centered)
    # PCA components
    W = Vt.T[:, :n_com]
    # projection into n_com dimensions
    return X_centered.dot(W)



# get data
data = np.loadtxt('../data/test_data.csv', delimiter=',')
X,y = data[:,:-1], data[:,-1]


X1D = dim_red(X, 1)

import matplotlib.pyplot as plt
plt.scatter(X1D[:,0], y,)


# get covariance matrix
cov = np.loadtxt('../data/test_cov.csv',delimiter=',')

# set model
model = PhylogeneticRegressor(X, y, cov, kernel='rbf')

# set hyperparamter space
params = {
    'lambda' : np.logspace(-5, 15, 200, dtype = float, base=10),
    'gamma' : np.logspace(-15, 3, 200, dtype = float, base=10),
}
model.set_hyperparameter_space(params=params)
    
# fit the model
model.fit(seed=12, sample=200)
model


# also you can set new set of hyperparameters
# and then fit
model.set_params(
    {'lambda': 0.0001, 'gamma': 0.0755},
    fit=True
)

# plotting
import matplotlib.pyplot as plt


feature_names = ['feature 1', 'feature 2', 'feature 3']

FI = model.FeatureImportance(seed=None)
plt.bar(feature_names, FI)
plt.ylabel('RMSE increase')

# len(model.P_inv)
# partial depedence plot

pdp_0 = model.pdp(feature=0)
pdp_1 = model.pdp(feature=1)
pdp_2 = model.pdp(feature=2)



plt.rcParams['figure.figsize'] = [10, 4]

fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
fig.suptitle('Partial Dependence Plot')

ax1.scatter(pdp_0[:,0], pdp_0[:,1])
ax1.set_ylabel('Predicted y')
ax1.set_xlabel('feature 1')

ax2.scatter(pdp_1[:,0], pdp_1[:,1])
ax2.set_xlabel('feature 2')

ax3.scatter(pdp_2[:,0], pdp_2[:,1])
ax3.set_xlabel('feature 3')



f1_idx = 2
f2_idx = 1

main_data = model.X

f1 = model.sample_feature(f1_idx, quantiles=[0,1], sample=30)
f2 = model.sample_feature(f2_idx, quantiles=[0,1], sample=30)


# make a meshgrid from f1 and f2
X1, X2 = np.meshgrid(f1, f2)
OUT = np.zeros(X1.shape)

for i in range(X1.shape[0]):
    for j in range(X1.shape[1]):
        # i,j
        X_tmp = main_data.copy()
        X_tmp[:,f1_idx] = X1[i,j]
        X_tmp[:,f2_idx] = X2[i,j]
        OUT[i,j] = np.mean( model.P_inv @ model.model.predict(X_tmp))

# can you add contour lines?

plt.rcParams['figure.figsize'] = [5, 5]
plt.contourf(X1, X2, OUT)
plt.colorbar()
plt.xlabel('feature 1')
plt.ylabel('feature 2')


plt.title('2D Partial Dependence Plot')



# ## 
# data = np.loadtxt('../data/test_log_data.csv', delimiter=',')
# X,y = data[:,:-1], data[:,-1]
# import matplotlib.pyplot as plt
# plt.scatter(X[:,0], X[:,2], c = y, alpha=0.6)


# file_path = "/Users/ulises/Desktop/ABL/software/phylokrr/data/test.tree"

# import dendropy
# tree = dendropy.Tree.get( path = 'path/to/tree', schema = 'newick')

# for i in tree.postorder_node_iter():
#     print( i.distance_from_root() )


