
import random
import numpy as np

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

from phylokrr.kernels import RBF_kernel, linear_kernel
from phylokrr.utils   import PGLS
from phylokrr.kernels import R2,rmse

def P_mat_simple(vcv):
    """
    get the square root of the inverse of the
    covariance matrix using spectral decomposition

    vcv: np.array, default = None
        covariance matrix
    """
    
    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

    ref: np.array
        reference data
    
    dat: np.array
        data to be scaled
    
    use_sd: bool, default = False
        use standard deviation to scale data
    
    Returns:
    --------
    np.array
        scaled 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)


class PGLS:
    def __init__(self, fit_intercept = False) -> None:
        self.fit_intercept = fit_intercept
        self.beta = np.array([])
    
    def fit(self, X, y, vcv):
        n,p = X.shape
        Oinv = np.linalg.inv(vcv)
        if self.fit_intercept:
            X_ = np.hstack((np.ones((n,1)), X))
        else:
            X_ = X

        self.beta = np.linalg.inv(X_.T @ Oinv @ X_) @ X_.T @ Oinv @ y
    
    def predict(self, X):
        if self.fit_intercept:
            X_ = np.hstack((np.ones((X.shape[0],1)), X))
        else:
            X_ = X
        return X_ @ self.beta

class WKRR:

    def __init__(self, kernel = 'rbf', fit_intercept = True, check_cov = False) -> None:
        """
        Weighted Kernel Ridge Regression

        kernel: str, default = 'rbf'
            kernel type

        fit_intercept: bool, default = True
            fit intercept term

        check_cov: bool, default = False
            Check covariance matrix, i.e., P.T @ P == vcv^-1
        """

        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
        
        # internal
        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):
        """
        get the square root of the inverse of the
        covariance matrix.

        vcv: np.array, default = None
            covariance matrix
        """
        
        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):
        """
        check if P.T @ P == vcv^-1
        expensive calculation for large matrices
        """
        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):
        """
        Fit the model

        vcv: np.array, default = None
            covariance matrix

        """
        # X = X_train
        # y = y_train
        # vcv = vcv_train
        
        self.X = X
        # self.y = y
        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)
        
        # intercept term 
        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'):
        """
        
        Calculates score considerint the covariance matrix
        This error is the same as the in objective function

        X_test: np.array
            testing data
        
        y_test: np.array
            testing targets
        
        vcv_test: np.array
            testing covariance matrix. If None, then it
            will return the standard RMSE and R^2

        metric: str, default = 'rmse'
            weighted root mean squared error (the same 
            that was used to construct the objective function).
            if metric = 'r2', then it returns the R^2 which
            also considers the covariance matrix, see Willett & Singer (1988)

        Returns:
        --------
        float
            score

        Reference:
        ---------
        Willett, J. B., & Singer, J. D. (1988). 
        Another cautionary note about R 2: Its use 
        in weighted least-squares regression analysis. 
        The American Statistician, 42(3), 236-238.
        """

        y_pred = self.predict(X_test)
        r = y_pred - y_test
        
        if isinstance(vcv_test, type(None)):
            P = np.eye(X_test.shape[0])

        else:
            P = self.P_mat(vcv_test)

        werr = P @ r # weighted residuals
        Py   = P @ y_test # weighted targets

        if metric == 'rmse':
            return rmse(werr)
        
        else:
            return R2(Py, werr)
            
        # if metric == 'rmse':
        #     return np.sqrt( np.mean( r**2 ) )
        
        # else:
        #     u = r.T @ r
        #     v = y_test.T @ y_test - n * np.mean(y_test)**2
        #     return 1 - (u/v)
    
    def opt_alpha(self, K, y, reg_lam = None, P = None):
        # Y = y
        # X = X
        # m = 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)
            # return np.linalg.solve(P.T @ P @ K + nlI,  P.T @ P @ y)

class KRR(WKRR):
    
        def __init__(self, kernel = 'rbf', fit_intercept = True) -> None:
            """
            Kernel Ridge Regression
    
            kernel: str, default = 'rbf'
                kernel type
    
            fit_intercept: bool, default = True
                fit intercept term
            """
    
            super().__init__(kernel, fit_intercept)
                    # internal
            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'])

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

        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:
                # r2
                u = ((y_test - y_pred)**2).sum()
                v = ((y_test - y_test.mean())** 2).sum()

                return 1 - (u/v)        
        
        def opt_alpha(self, K, y, reg_lam = None):
            # Y = y
            # X = X
            # m = None
            n,_ = self.X.shape
            K_Idx = K + n * reg_lam * np.diag( np.ones( K.shape[0] ) )
            return np.linalg.inv( K_Idx ).dot( y )


# cov. matrix obtained from the phylogenetic tree
vcv = np.loadtxt("../data/test_cov2.csv", delimiter=',') 

# Trait simulation under Brownian motion
n = vcv.shape[0]
mean = np.zeros(n)
X = np.random.multivariate_normal(cov=vcv, mean=mean).reshape(-1,1)
# Non-linear response variable (sine curve)

X = scale_weighted_data(X, vcv, use_sd = True)
y = np.sin(X*2).ravel() + 5 

# Add noise to the response variable
y[::10] += 4 * (0.5 - np.random.rand(X.shape[0] // 10)) 
# y = scale_weighted_data(y.reshape(-1,1), vcv, use_sd = False)

num_test = round(0.5*n)
from phylokrr.utils import split_data

(X_train, y_train, X_test, y_test) = split_data(X, y, num_test = num_test, seed = 12038)

model = KRR(kernel='rbf', fit_intercept= True)



# arbitrarily proposed hyperparameters
params = {'lambda': 1, 'gamma': 1}

# set hyperparamters
model.set_params(**params)


# fit model with phylogenetic covariance matrix
model.fit(X_train, y_train)
y_pred1 = model.predict(X_test)
rmse1 = model.score(X_test, y_test,  metric='rmse')
r1 = model.score(X_test, y_test,  metric='r2')







import matplotlib.pyplot as plt
# fit standard phylogenetic regression
b_wls = PGLS(X_train, y_train, vcv_train)
y_pred3 = np.hstack((np.ones((X_test.shape[0],1)), X_test)) @ b_wls
P = model.P_mat(vcv_test)
w_r_pgls = P @ (y_pred3 - y_test)
Py_test = P @ y_test
r2_pgls = R2(Py_test, w_r_pgls)
rmse_pgls = rmse(w_r_pgls)
# store figure

plt.scatter(X_test, y_test , color = 'blue' , alpha=0.5, label = 'Testing (unseen) data')
plt.scatter(X_test, y_pred1, color = 'green', alpha=0.5, label = 'phyloKRR predictions w\o CV')
# plt.scatter(X_test, y_pred2, color = 'gray', alpha=0.5, label = 'KRR predictions w\o CV')
plt.scatter(X_test, y_pred3, color = 'red', alpha=0.5, label = 'PGLS predictions')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.tight_layout()








from phylokrr.utils import split_data_vcv

# split data into training and testing sets 
num_test = round(0.5*n)

(X_train  , X_test,  
 y_train  , y_test,  
 vcv_train, vcv_test) = split_data_vcv(X, y, vcv, num_test, seed = 12038) # seed defined above

from phylokrr.kernels import KRR

# set model
model = KRR(kernel='rbf', fit_intercept= True)

# arbitrarily proposed hyperparameters
params = {'lambda': 1, 'gamma': 1}

# set hyperparamters
model.set_params(**params)

# fit model with phylogenetic covariance matrix
model.fit(X_train, y_train, vcv = vcv_train)
y_pred1 = model.predict(X_test)
rmse1 = model.score(X_test, y_test, vcv_test, metric='rmse')
r1 = model.score(X_test, y_test, vcv_test, metric='r2')

model.fit(X_train, y_train, vcv =None)
y_pred2 = model.predict(X_test)
rmse2 = model.score(X_test, y_test, None, metric='rmse')
r2 = model.score(X_test, y_test, None, metric='r2')


import matplotlib.pyplot as plt


# fit standard phylogenetic regression
b_wls = PGLS(X_train, y_train, vcv_train)
y_pred3 = np.hstack((np.ones((X_test.shape[0],1)), X_test)) @ b_wls
P = model.P_mat(vcv_test)
w_r_pgls = P @ (y_pred3 - y_test)
Py_test = P @ y_test
r2_pgls = R2(Py_test, w_r_pgls)
rmse_pgls = rmse(w_r_pgls)
# store figure

plt.scatter(X_test, y_test , color = 'blue' , alpha=0.5, label = 'Testing (unseen) data')
plt.scatter(X_test, y_pred1, color = 'green', alpha=0.5, label = 'phyloKRR predictions w\o CV')
# plt.scatter(X_test, y_pred2, color = 'gray', alpha=0.5, label = 'KRR predictions w\o CV')
plt.scatter(X_test, y_pred3, color = 'red', alpha=0.5, label = 'PGLS predictions')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.tight_layout()
plt.savefig(
    './../data/imgs/phyloKRR_vs_PGLS.png', dpi = 200,
    bbox_inches = 'tight', pad_inches = 0.1
    )



from phylokrr.utils import k_fold_cv_random

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

best_params = k_fold_cv_random(X_train, y_train, vcv_train,
                                model, 
                                params,
                                folds = 3, 
                                sample = 70)

model.set_params(**best_params)
model.fit(X_train, y_train, vcv = vcv_train)
y_pred_cv = model.predict(X_test)


# control the size
fs = 18
plt.rcParams['figure.figsize'] = [7, 6]
plt.scatter(X_test, y_test, color = 'blue' ,
            alpha=0.5, 
            label = 'Testing (unseen) data',)
plt.scatter(X_test, y_pred_cv, 
            color = 'green', alpha=0.5, 
            label = 'phyloKRR predictions',)
plt.scatter(X_test, y_pred3, 
            color = 'red', alpha=0.5, 
            label = 'PGLS predictions',)
plt.xlabel('x', fontsize = fs)
plt.ylabel('y', fontsize = fs)
plt.legend(fontsize = fs)
plt.tight_layout()
plt.savefig(
    '/Users/ulises/Desktop/ABL/phyloKRR_vs_PGLS_cv.png', dpi = 200,
    bbox_inches = 'tight', pad_inches = 0.1
    )


r2_cv = model.score(X_test, y_test, vcv_test, metric='r2')

model.score(X_test, y_test, vcv_test, metric='rmse')
