
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)



# write stacked generalization 
# which optimizes the weights such that its sum is 1
# using quadratic programming
import cvxpy as cp
import copy
from phylokrr.utils import P_mat_simple, weight_data, myfunc, P_inv_simple


def stacked_generalization(X, y, covs, num_test = 100, seed = 12038):
    import glob
    path_trees = '/Users/ulises/Desktop/ABL/software/phylokrr/data/r_coal_sim_500spps_100trees'
    vcvs = glob.glob(path_trees + '/*.csv')

    vcv = np.loadtxt(vcvs[0], delimiter=',')
    I_vcv = np.diag(np.diag(vcv))
    covs = [ I_vcv, vcv]

    np.random.seed(12038)

    # n,_ = X.shape
    # noise = np.random.multivariate_normal(np.zeros(n), vcv, 1).reshape(-1,1)
    # X = np.random.multivariate_normal(np.zeros(n), vcv, 1).reshape(-1,1)

    # y = myfunc(X, b=2, type='pol', 
    #         add_noise=False,
    #         noise=noise).reshape(-1,1)
    # y += noise
    # plt.scatter(X, y)

    n,_ = vcv.shape
    noise = np.random.normal(0,1, n).reshape(-1,1)
    X_w = np.random.normal(0,1,n).reshape(-1,1)

    y_w = myfunc(X_w, b=1, type='pol', 
            add_noise=False,
            noise=noise).reshape(-1,1)
    y_w += noise


    P_inv = P_inv_simple(vcv)
    X = P_inv @ X_w
    y = P_inv @ y_w

    subplots = 2

    fig, axs = plt.subplots(1, subplots, figsize=(8, 4))
    axs[0].scatter(  X[:,0], y, color = 'blue',alpha=0.5)
    axs[1].scatter(X_w[:,0], y_w, color = 'red', alpha=0.5)


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


    # take a random sample from range(n
    test_idx = np.random.choice(range(n), size = num_test, replace = False)
    train_idx = list( set(range(n)) - set(test_idx) )

    # X = np.ones(n).reshape(-1,1)
    X_train, X_test = X[train_idx,:], X[test_idx,:]
    y_train, y_test = y[train_idx]  , y[test_idx]

    errors = np.zeros((len(test_idx), len(covs)))
    for i,vcv in enumerate(covs):

        # h = 1
        # new_vcv  = np.eye(n) * h + vcv*(1-h)
        # np.max(np.linalg.eigvals(new_vcv))


        vcv_train, vcv_test = vcv[train_idx,:][:,train_idx], vcv[test_idx,:][:,test_idx]
        model = KRR(kernel='rbf', fit_intercept=True)

        X_w_train, y_w_train = weight_data(X_train, y_train, vcv_train)
        X_w_test, y_w_test = weight_data(X_test, y_test, vcv_test)
        # plt.scatter(X_w_train, y_w_train)

        best_params = k_fold_cv_random(X_w_train, y_w_train,
                                        model,
                                        params,
                                        verbose = False,
                                        folds = 5,
                                        sample = 100)

        model.set_params(**best_params)

        model.fit(X_w_train, y_w_train)
        tmp_errors = (y_w_test - model.predict(X_w_test)).ravel()
        
        # normalize errors
        errors[:,i] = tmp_errors /  (max(y_w_test) - min(y_w_test))

        # normalize errors with the first and third interquartile of y_w_test
        # errors[:,i] = tmp_errors / (np.percentile(y_w_test, 75) - np.percentile(y_w_test, 25))
        # errors[:,i] = tmp_errors / np.std(y_w_test)

    Q = errors.T @ errors
    print(Q)
    m = Q.shape[0]
    ones = np.ones(m)

    # Define and solve the CVXPY problem.
    w = cp.Variable(m)
    prob = cp.Problem(cp.Minimize( cp.quad_form(w, Q) ),
                    [ones.T @ w == 1, w >= 0])
    prob.solve()

    # Print result.
    print(np.round(w.value,3))

