
import numpy as np
import random
import numpy as np
import glob
# import matplotlib.pyplot as plt

# seed for reproducibility
np.random.seed(12038)
random.seed(12038)

def split_data(X,y, num_test, seed = 123):
    random.seed(seed)
    n,_ = X.shape

    test_idx = random.sample(range(n), k = num_test)
    train_idx = list( set(range(n)) - set(test_idx) )

    X_train, X_test = X[train_idx,:], X[test_idx,:]
    y_train, y_test = y[train_idx]  , y[test_idx]

    return X_train, X_test, y_train, y_test


def split_data_vcv(X,y,vcv, num_test, seed = 123):
    random.seed(seed)
    n,_ = X.shape

    test_idx = random.sample(range(n), k = num_test)
    train_idx = list( set(range(n)) - set(test_idx) )

    X_train, X_test = X[train_idx,:], X[test_idx,:]
    y_train, y_test = y[train_idx]  , y[test_idx]
    vcv_train, vcv_test = vcv[train_idx,:][:,train_idx], vcv[test_idx,:][:,test_idx]

    return X_train, X_test, y_train, y_test, vcv_train, vcv_test

def R2(y, r):
    """
    Coefficient of determination
    r = y - y_pred
    """

    # r = y - y_pred
    u = r.T @ r
    v = y.T @ y - (len(y) * np.mean(y)**2)
    return 1 - (u/v)

def rmse(r):
    """
    Root Mean Squared Error
    r = y - y_pred
    """
    # r = y - y_pred
    return np.sqrt( np.mean( r**2 ) )

def distance_matrix(a, b):
    """
    l2 norm squared matrix
    """
    return np.linalg.norm(a[:, None, :] - b[None, :, :], axis=-1)**2

def RBF_kernel(a, b, gamma):
    """
    Radial Basis Function
    """

    tmp_rbf = -gamma * distance_matrix(a, b)
    np.exp(tmp_rbf, tmp_rbf) # RBF kernel. Inplace exponentiation
    return tmp_rbf

def linear_kernel(a, b, c):
    """
    Linear Kernel
    """
    XXt = a.dot(b.T)
    C = c * np.ones(XXt.shape)

    return XXt + C

def P_mat_simple(vcv):
    """
    get the square root of the inverse of the
    """

    # Kr = np.diag(1/np.sqrt(np.diag(vcv)))
    # vcv = Kr @ vcv @ Kr

    if isinstance(vcv, type(None)):
        return None

    L,Q = np.linalg.eig( vcv )
    P = Q @ np.diag( L**(-1/2) ) @ Q.T

    return P

def _scaler(ref, dat, use_sd = False):
    """
    Center and scale data
    """

    u = np.mean(ref, axis=0)
    centered = dat - u

    if use_sd:
        sd = np.std(ref, axis=0)
        sted = centered / sd

        return sted
    else:

        return centered

def scale_weighted_data(X, vcv, use_sd = False):
    """
    scale data by the square root of the inverse of the
    covariance matrix

    X: np.array
        data

    vcv: np.array
        covariance matrix
    """

    P = P_mat_simple(vcv)
    X_w_uc = P @ X
    return _scaler(X_w_uc, X_w_uc, use_sd)

def k_fold_cv(X, y, model, num_folds):
    """
    k-fold cross-validation
    """

    n, p = X.shape
    fold_size = n // num_folds

    all_errors = []
    for i in range(num_folds):

        test_idx = list(range(i * fold_size, (i + 1) * fold_size))
        train_idx = list(set(range(n)) - set(test_idx))

        X_train, X_test = X[train_idx, :], X[test_idx, :]
        y_train, y_test = y[train_idx], y[test_idx]
    
        model.fit(X_train, y_train)
        tmp_err = model.score(X_test, y_test, metric='rmse')
        all_errors.append(tmp_err)

    return np.mean(all_errors)

def k_fold_cv_random(X, y, model, params, folds=3, sample=500, verbose=True, seed=123):
    """
    Random search for hyperparameter tuning using k-fold cross-validation
    """
    np.random.seed(seed=seed)

    # make random choice from the grid of hyperparameters
    all_params = params.keys()
    tested_params = np.zeros((sample, len(all_params)))

    for n, k in enumerate(all_params):
        tested_params[:, n] = np.random.choice(params[k], sample)

    if verbose:
        # check tested_params are unique
        tested_params = np.unique(tested_params, axis=0)
        print("Number of unique hyperparameters: ", tested_params.shape[0])

    # shuffle the data
    n = X.shape[0]
    idx = np.arange(n)
    np.random.shuffle(idx)
    X = X[idx]
    y = y[idx]

    all_errors = []
    for vec in tested_params:
        tmp_params = dict(zip(all_params, vec))
        model.set_params(**tmp_params)
        tmp_err = k_fold_cv(X, y, model, folds)
        all_errors.append([tmp_params, tmp_err])

        if verbose:
            print("CV score: %s, Parameters %s" % (tmp_err, tmp_params))

    # take the best hyperparameters
    best_ = sorted(all_errors, key=lambda kv: kv[1], reverse=False)[0]

    if verbose:
        print("Best CV score: ", best_[1])

    return best_[0]


class PGLS:
    def __init__(self, fit_intercept=False, weighted=False) -> None:
        self.fit_intercept = fit_intercept
        self.weighted = weighted
        self.intercept = 0
        self.beta = np.array([])

    def fit(self, X, y, vcv):
        n, p = X.shape

        if self.weighted:
            tmp_beta = np.linalg.solve(X.T @ X, X.T @ y)

        if self.fit_intercept:
            self.intercept = np.mean(y - X @ tmp_beta)

        self.beta = np.hstack((self.intercept, tmp_beta)) if self.fit_intercept else tmp_beta

    def predict(self, X):
        n, p = X.shape

        if self.fit_intercept:
            X = np.hstack((np.ones((n, 1)), X))

        return X @ self.beta

    def score(self, X_test, y_test, metric='rmse'):
        y_pred = self.predict(X_test)

        if metric == 'rmse':
            return np.sqrt(np.mean((y_pred - y_test) ** 2))
        else:
            u = ((y_test - y_pred) ** 2).sum()
            v = ((y_test - y_test.mean()) ** 2).sum()

            return 1 - (u / v)


class WKRR:
    def __init__(self, kernel='rbf', fit_intercept=True, check_cov=False) -> None:
        self.kernel = kernel

        if self.kernel == 'rbf':
            self.params = {'gamma': 0.1, 'lambda': 0.1}

        else:
            self.params = {'c': 0.1, 'lambda': 0.1}

        self.fit_intercept = fit_intercept
        self.check_cov = check_cov
        self.intercept = 0
        self.alpha = np.array([])
        self.X = np.array([])
        self.chol = False

    def set_params(self, **params):
        if self.kernel == 'rbf':
            self.params['gamma'] = params['gamma']

        else:
            self.params['c'] = params['c']

        self.params['lambda'] = params['lambda']

    def get_params(self):
        return self.params

    def P_mat(self, vcv):
        if isinstance(vcv, type(None)):
            return None

        if self.check_cov:
            self.assert_COV_sym(vcv)

        if self.chol:
            C = np.linalg.cholesky(vcv)
            P = np.linalg.inv(C)

        else:
            L, Q = np.linalg.eig(vcv)
            P = Q @ np.diag(L ** (-1 / 2)) @ Q.T

        if self.check_cov:
            self.assert_COV_decom(P, vcv)

        return P

    def assert_COV_sym(self, vcv, tol=1e-8):
        assert np.all(np.abs(vcv - vcv.T) < tol), 'not symmetric matrix'

    def assert_COV_decom(self, P, vcv):
        assert np.all(np.round(P.T @ P, 2) == np.round(np.linalg.inv(vcv), 2)), "P.T @ P != vcv^-1"

    def fit(self, X, y, vcv=None):
        self.X = X
        P = self.P_mat(vcv)

        if self.kernel == 'rbf':
            K_train = RBF_kernel(self.X, self.X, self.params['gamma'])
        else:
            K_train = linear_kernel(self.X, self.X, self.params['c'])

        self.alpha = self.opt_alpha(K_train, y, self.params['lambda'], P)

        if self.fit_intercept:
            self.intercept = np.mean(y - K_train @ self.alpha)

    def predict(self, X_test):
        assert len(self.alpha) > 0, "The model needs to be fitted first"

        if self.kernel == 'rbf':
            K_test = RBF_kernel(X_test, self.X, self.params['gamma'])
        else:
            K_test = linear_kernel(X_test, self.X, self.params['c'])

        return K_test @ self.alpha + self.intercept

    def score(self, X_test, y_test, vcv_test, metric='rmse'):
        y_pred = self.predict(X_test)

        if isinstance(vcv_test, type(None)):
            P = np.eye(X_test.shape[0])
        else:
            P = self.P_mat(vcv_test)

        werr = P @ (y_pred - y_test)  # weighted residuals

        if metric == 'rmse':
            return rmse(werr)
        else:
            Py = P @ y_test  # weighted targets
            return R2(Py, werr)

    def opt_alpha(self, K, y, reg_lam=None, P=None):

        n, _ = self.X.shape
        I = np.eye(K.shape[0])
        nlI = n * reg_lam * I

        if isinstance(P, type(None)):
            return np.linalg.solve(K + nlI, y)
        
        else:
            return P @ np.linalg.solve(P @ K @ P + nlI, P @ y)


class KRR(WKRR):

    def __init__(self, kernel='rbf', fit_intercept=True) -> None:
        super().__init__(kernel, fit_intercept)
        self.intercept = 0
        self.alpha = np.array([])
        self.X = np.array([])

    def fit(self, X, y):
        self.X = X
        self.y = y

        if self.kernel == 'rbf':
            K_train = RBF_kernel(self.X, self.X, self.params['gamma'])
            
        else:
            K_train = linear_kernel(self.X, self.X, self.params['c'])

        self.alpha = self.opt_alpha(K_train, self.y, self.params['lambda'])

        if self.fit_intercept:
            self.intercept = np.mean(y - K_train @ self.alpha)

    def opt_alpha(self, K, y, reg_lam=None):
        
        n, _ = self.X.shape
        I = np.eye(K.shape[0])
        nlI = n * reg_lam * I

        return np.linalg.solve(K + nlI, y)
        

    def score(self, X_test, y_test, metric='rmse'):
        y_pred = self.predict(X_test)

        if metric == 'rmse':
            return np.sqrt(np.mean((y_pred - y_test) ** 2))
        else:
            u = ((y_test - y_pred) ** 2).sum()
            v = ((y_test - y_test.mean()) ** 2).sum()

            return 1 - (u / v)


def myfunc(x, b=0, type='sin'):
    if type == 'sin':
        return np.sin( np.sum(x, axis=1) * 2).ravel()
    
    else:
        return np.sum(x, axis=1) ** b

def sim_data(vcv, mean_vector, b, 
             on_weighted=False, 
             add_noise = False, 
             noise_var = 1,
             n_var = 3, 
             extra_weights = [1,1,1], 
             type='pol'):
    """
    Simulate data for n_var independent variables

    Parameters
    ----------
    vcv: np.array
        covariance matrix
    
    mean_vector: np.array
        mean vector
    
    b: int
        power of the polynomial function

    on_weighted: bool
        whether to use weighted data or not to obtain the response variable.
        the weights are obtained using the 
        squared root of the inverse of 
        the covariance matrix

    add_noise: bool
        whether to add noise to the response variable

    noise_var: float
        variance of the noise

    extra_weights: list
        weights for the independent variables

    type: str
        type of the function to simulate the data
        if 'sin' then a sine function is used over the
        sum of the independent variables
        if 'pol' then a polynomial function is used over the
        sum of the independent variables
    
    Returns
    -------
    X_w_uc: np.array
        weighted independent variables (uncentered)
    
    y_w_uc: np.array
        weighted response variable (uncentered)
    """
    n_var = 3

    assert len(extra_weights) == n_var, "extra weights must have the same length as the number of predictors"

    X_uw_uc = np.zeros((vcv.shape[0], n_var))
    for j in range(n_var):
        X_uw_uc[:,j] = np.random.multivariate_normal(mean=mean_vector, cov=vcv)*extra_weights[j]


    P = P_mat_simple(vcv)
    noise = np.random.normal(0, noise_var, X_uw_uc.shape[0])

    if on_weighted:
        X_w_uc = P @ X_uw_uc
        y_w_uc = myfunc(X_w_uc, b=b, type=type).ravel()

        if add_noise:
            y_w_uc += noise

        # plt.scatter(np.sum(X_w_uc, 1), y_w_uc, color='blue', alpha=0.5)
    else:
        y_uw_uc = myfunc(X_uw_uc, b=b, type=type).ravel()

        if add_noise:
            y_uw_uc += noise        

        X_w_uc = P @ X_uw_uc
        y_w_uc = P @ y_uw_uc

    return X_w_uc, y_w_uc


def test_vcv(vcv, b=4, on_weighted=True, sample=50, add_noise=False, noise_var=1):

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

    X_w_uc, y_w_uc = sim_data(vcv, mean_vector, b=b, 
                              on_weighted=on_weighted,
                              add_noise=add_noise,
                              noise_var=noise_var)

    X_wc = _scaler(X_w_uc, X_w_uc, use_sd= False)
    y_wc = _scaler(y_w_uc, y_w_uc, use_sd= False)

    (X_wc_train, X_wc_test,
     y_wc_train, y_wc_test,
     vcv_train , vcv_test) = split_data_vcv(X_wc, y_wc, vcv, num_test=num_test, seed=12038)

    rbf_param = k_fold_cv_random(X_wc_train, y_wc_train, krr_rbf, rbf_grid_params, verbose=False, folds=3, sample=sample)
    linear_param = k_fold_cv_random(X_wc_train, y_wc_train, krr_linear, linear_grid_params, verbose=False, folds=3, sample=sample)

    krr_rbf.set_params(**rbf_param)
    krr_linear.set_params(**linear_param)

    krr_rbf.fit(X_wc_train, y_wc_train)
    krr_linear.fit(X_wc_train, y_wc_train)
    pgls_model.fit(X_wc_train, y_wc_train, vcv_train)

    rmse_wc_krr_rbf = krr_rbf.score(X_wc_test, y_wc_test, metric='rmse')
    rmse_wc_krr_linear = krr_linear.score(X_wc_test, y_wc_test, metric='rmse')
    rmse_wc_pgls = pgls_model.score(X_wc_test, y_wc_test, metric='rmse')

    error_sim = [b, rmse_wc_krr_rbf, rmse_wc_krr_linear, rmse_wc_pgls]
    print(error_sim)
    return error_sim

def permutation_importance(X, y, model, num_test = 100, seed = 12038, iterations = 1000):

    np.random.seed(seed)

    n,p = X.shape

    X_train, X_test, y_train,y_test = split_data(X, y, num_test, seed = seed)
    model.fit(X_train, y_train)
    error_orig = model.score(X_test, y_test)


    out = np.zeros((iterations, p))

    for i in range(iterations):

        for j in range(X_test.shape[1]):
            # Create a copy of X_test
            X_test_copy = X_test.copy()

            # Scramble the values of the given predictor
            X_test_copy[:,j] = np.random.permutation(X_test_copy[:,j])
                        
            # Calculate the new RMSE
            error_perm = model.score(X_test_copy, y_test)

            out[i,j] = error_perm - error_orig

    return out

def pdp(self, feature, integer = False, quantiles = [0,1], sample = 70):
    """
    partial depedence plot values.

    Parameters
    ----------
    feature : int

    Returns
    -------
    PDP of the transformed data
    """
    if not len(self.P_inv):
        self.P_inv = np.linalg.inv(self.P)

    Xq = self.sample_feature(feature, quantiles, sample, integer)
    
    pdp_values = []
    for n in Xq:
        # copy original values
        X_tmp = self.X.copy()
        # make original values with 
        # modified feature column
        X_tmp[:,feature] = n
        # weight the whole dataset as
        # the model learned from the 
        # weighted data
        X_tmp_test = self.P @ X_tmp

        pdp_values.append(
            np.mean(
                # make rows unweighted
                # (i.e., again 'correlated')
                self.P_inv @ self.model.predict(X_tmp_test)
            )
        )

    out = np.hstack((
        Xq.reshape(-1,1), 
        np.array(pdp_values).reshape(-1,1)
    ))
    return out

rbf_grid_params = {
    'lambda': np.logspace(-10, 5, 100, dtype=float, base=2),
    'gamma': np.logspace(-10, 5, 100, dtype=float, base=2),
}

linear_grid_params = {
    'c': np.logspace(-10, 5, 100, dtype=float, base=2),
    'lambda': np.logspace(-10, 5, 100, dtype=float, base=2)
}

krr_rbf    = KRR(kernel='rbf', fit_intercept=not True)
krr_linear = KRR(kernel='linear', fit_intercept=not True)
pgls_model = PGLS(fit_intercept=not True, weighted=True)

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=',')
b = 2
on_weighted = True
add_noise = True
noise_var = 1
sample = 100


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

X_w_uc, y_w_uc = sim_data(vcv, mean_vector, b=b, 
                            on_weighted=on_weighted,
                            add_noise=add_noise,
                            noise_var=noise_var,
                            extra_weights=[3,2,1]
                            )

X_wc = _scaler(X_w_uc, X_w_uc, use_sd= False)
y_wc = _scaler(y_w_uc, y_w_uc, use_sd= False)

(X_wc_train, X_wc_test,
    y_wc_train, y_wc_test,
    vcv_train , vcv_test) = split_data_vcv(X_wc, y_wc, vcv, num_test=num_test, seed=12038)

rbf_param = k_fold_cv_random(X_wc_train, y_wc_train, krr_rbf, rbf_grid_params, verbose=False, folds=3, sample=sample)
linear_param = k_fold_cv_random(X_wc_train, y_wc_train, krr_linear, linear_grid_params, verbose=False, folds=3, sample=sample)

krr_rbf.set_params(**rbf_param)
krr_rbf.fit(X_wc_train, y_wc_train)
# krr_linear.fit(X_wc_train, y_wc_train)

def sample_feature(X, feature, quantiles, sample, integer = False):
    """
    sample unweighted column
    """
    # X = X0
    # feature =0
    Xpdp = X[:,feature]

    if integer:
        return np.sort(np.unique(Xpdp))
    
    q = np.linspace(
        start = quantiles[0],
        stop  = quantiles[1],
        num   = sample
    )

    # sample from the original feature space
    Xq = np.quantile(Xpdp, q = q)

    return Xq    

def pdp_1d(X, feature, model, quantiles = [0,1], sample = 70, integer = False):
    """
    1D partial dependence plot
    """

    Xq = sample_feature(X, feature, quantiles, sample, integer)
    pdp_values = []
    for n in Xq:
        # copy original values
        X_tmp = X.copy()
        # make original values with 
        # modified feature column
        X_tmp[:,feature] = n

        pdp_values.append(
            np.mean(
                model.predict(X_tmp)
            )
        )

    out = np.hstack((
        Xq.reshape(-1,1), 
        np.array(pdp_values).reshape(-1,1)
        ))
    
    return out


import matplotlib.pyplot as plt


def pdp_2d(X, f1_idx, f2_idx, model, quantiles, sample, integer):
    """
    2D partial dependence plot
    """

    Xq1 = sample_feature(X, f1_idx, quantiles, sample, integer[0])
    Xq2 = sample_feature(X, f2_idx, quantiles, sample, integer[1])

    n_xq1 = len(Xq1)
    n_xq2 = len(Xq2)

    X1, X2 = np.meshgrid(Xq1, Xq2)

    Y = np.zeros((n_xq1, n_xq2))
    for i in range(n_xq1):
        for j in range(n_xq2):
            # j,i = 0,0
            X_tmp = X.copy()
            X_tmp[:,f1_idx] = X1[i,j]
            X_tmp[:,f2_idx] = X2[i,j]

            Y[i,j] = np.mean( model.predict(X_tmp) )

    return X1, X2, Y

# 2d partial dependence plot
sample = 10
f1_idx = 2
f2_idx = 1
X = X_wc
quantiles = [0,1]
integer = [False, False]
model = krr_rbf


X1, X2, Y = pdp_2d(X, f1_idx, f2_idx, model, quantiles, sample, integer)

fig,ax=plt.subplots(1,1)
cp = ax.contourf(X1, X2, Y)
fig.colorbar(cp) # Add a colorbar to a plot
ax.set_title('2D partial dependence plot')
ax.set_xlabel('feature 1')
ax.set_ylabel('feature 2')



# permutation importance
out = permutation_importance(X_wc, y_wc,krr_rbf, num_test = 100, iterations=5)

# plot bar chart
import matplotlib.pyplot as plt
plt.bar(range(3), np.mean(out, axis=0))
plt.errorbar(range(3), np.mean(out, axis=0), np.std(out, axis=0), fmt='.', color='black')
plt.xticks(range(3), ['feature 1', 'feature 2', 'feature 3'])
plt.ylabel('RMSE increase')


pdp_1d_out = pdp_1d(X_wc, 0, krr_rbf)

# make a line plot
plt.plot(pdp_1d_out[:,0], pdp_1d_out[:,1])




