
import random
import numpy as np

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



import numpy as np
from phylokrr.metrics import rmse, R2

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



class KRR:

    def __init__(self, kernel = 'rbf', 
                 fit_intercept = True, 
                 check_cov = False,
                 store_gram = 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

        store_gram: bool, default = False
            store the gram matrix used to fit the model
        """

        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.store_gram = store_gram
        
        # internal
        self.intercept = 0
        self.alpha = np.array([])
        self.X = np.array([])
        self.K = 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 Gram_mat(self, X1, X2):
        """
        Gram matrix

        X1: np.array
            data matrix (e.g., testing data)

        X2: np.array
            data matrix (e.g., training data)
        """

        if self.kernel == 'rbf':
            K = RBF_kernel(X1, X2, self.params['gamma'])

        else:
            K = linear_kernel(X1, X2, self.params['c'])

        return K
        

    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)
        K_train = self.Gram_mat(self.X, self.X)

        if self.store_gram:
            self.K = K_train

        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"

        K_test = self.Gram_mat(X_test, self.X)
        
        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)


# 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)
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)) 

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 = seed) # seed defined above



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

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

# set hyperparamters
model.set_params(**params)

import time
start = time.time()
# 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')
end = time.time()
print("Time: ", end - start)


# import matplotlib.pyplot as plt
# from phylokrr.utils import PGLS

# # 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

# 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_pred3, color = 'red', alpha=0.5, label = 'PGLS predictions')
# plt.xlabel('x')
# plt.ylabel('y')
# plt.legend()

from scipy.optimize import minimize
import numpy as np

# )
model = KRR(kernel='rbf',
            fit_intercept = True,
            store_gram = True)
model.fit(X, y, vcv = vcv)

K = model.K
a = model.alpha
l = model.params['lambda']


def objective_function(K, y, l, vcv, lambda_val):
    n,p = X.shape
    # lambda_val = 0.05

    W = np.diag( np.diag(vcv) )
    C = lambda_val * vcv + (1 - lambda_val) * W
    
    nlI = n * l * np.eye(n)
    OpK = np.linalg.solve(C, K)
    Opy = np.linalg.solve(C, y)

    a = np.linalg.solve(OpK+ nlI, Opy)
    e = y - K @ a
    
    dual = (1/n) * e.T @ np.linalg.solve(C, e) + l*( a.T @ K @ a )



    

    C_inv = np.linalg.inv(C)
    beta  = np.linalg.inv(X.T @ C_inv @ X) @ (X.T @ C_inv @ y)
    e = y - X @ beta 
    s = (1/n)*e.T @ C_inv @ e

    objective = n*np.log(s) + np.log(np.linalg.det(C)) + (1/s)*e.T@ C_inv @e
    return objective.astype(np.float64)[0][0]

def optimize_problem(X, y, Sigma):
    n, p = X.shape
    X = np.hstack((np.ones((n, 1)), X))

    # Define the bounds for lambda
    bounds = [(0, 1)]

    # Initial guess for lambda 
    initial_lambda = np.random.uniform(0, 1)

    # Define the optimization problem
    fnc = lambda lambda_val: objective_function(X, y, Sigma, lambda_val)
    result = minimize(fnc, x0 = initial_lambda, bounds=bounds, method='SLSQP')

    return result


