import numpy as np

from phylokrr.kernels import *


class LKRR:

    def __init__(self, kernel = 'rbf') -> 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}

        # internal
        self.alpha = []
        self.X = []
        self.y = []

    def p1(self, X, beta):
        return 1/(1 + np.exp(X.dot(beta)))

    def p0(self, X, beta):
        EX = np.exp(X.dot(beta))
        return EX/(1 + EX)

    def set_params(self, params):
        self.params = params

    def get_params(self):
        return self.params
    
    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'])

    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
    
    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 rmse(self, K, alpha, Y):

        return np.sqrt( np.mean( (K.dot(alpha) - Y)**2 ) )
    
    def opt_alpha(self, K, y, reg_lam = None):
        # Y = y
        # X = X
        # m = None
        K_Idx = K + reg_lam * np.diag( np.ones( K.shape[0] ) )
        return np.linalg.inv( K_Idx ).dot( y )



# set.seed(666)
# > x1 = rnorm(1000)           # some continuous variables 
# > x2 = rnorm(1000)
# > z = 1 + 2*x1 + 3*x2        # linear combination with a bias
# > pr = 1/(1+exp(-z))         # pass through an inv-logit function
# > y = rbinom(1000,1,pr)      # bernoulli response variable
# > 
# > #now feed it to glm:
# > df = data.frame(y=y,x1=x1,x2=x2)
# > glm( y~x1+x2,data=df,family="binomial")



def P_mat(vcv, chol = False, corr = False):
    
    if corr:
        Kr = np.diag(1/np.sqrt(np.diag(vcv)))
        vcv = Kr @ vcv @ Kr

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

    else:
        Oinv = np.linalg.inv( vcv )
        L,Q  = np.linalg.eig( Oinv )
        P  = Q @ np.diag( np.sqrt( 1/L ) ) @ Q.T

    return P


# ## logistic data

import glob
covs = glob.glob("/Users/ulises/Desktop/Machine Learning/HW3/SantaquiteriasProblem/simulations/rcoal_*.csv")

cov_file = covs[0]
vcv = np.loadtxt(cov_file, delimiter = ",")
nspps = vcv.shape[0]

mean_vector = np.zeros(nspps)

P = P_mat(vcv, chol= False, corr = False)

b = 1
X =  np.random.multivariate_normal(
    mean=mean_vector, 
    cov=vcv,
    # cov=np.diag(np.ones(n_spps)),
    size=2
    ).T

# logistic data
z = (X[:,0] + X[:,1])**b
probs = 1/(1 + np.exp(-z))
y = np.random.binomial(n = 1, p = probs)


sum(y == 1)
sum(y == 0)

np.max(probs)
np.min(probs)


import matplotlib.pyplot as plt
plt.scatter(X[:,0], X[:,1], c = y, alpha=0.6)


# y = rbinom(1000,1,pr) 


# data = np.loadtxt('../data/test_log_data.csv', delimiter=',')

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


# [n,p] = np.shape(data)
# num_train = int(0.5*n)

# X,y = data[:,:-1], data[:,-1]
# P = P_mat(cov,corr=True)

# X = P @ X
# y = P @ y


num_train = round(nspps*0.60)

sample_train = X[0:num_train,:]
sample_test  = X[num_train: ,:]
    
label_train = y[0:num_train]
label_test  = y[num_train: ]
p = 2
    



import warnings

def p1(X, beta):
    return 1/(1 + np.exp(X.dot(beta)))

def p0(X, beta):
    EX = np.exp(X.dot(beta))
    return EX/(1 + EX)

def error_est(X, y,  beta, norm = False):
    if norm:
        return np.linalg.norm(p1(X,beta) - y, ord = 2)
        # pass
    else:
        y_pred = (p1(X, beta) >= 0.5).astype(float)
        return np.mean(y_pred != y)

def NewtonRaphsonStep(X, y, beta):

    Lp = -X.T @ ( y - p1(X, beta) )
    W  = np.diag( p1(X, beta) * p0(X, beta) )

    Lpp = -X.T @ W @ X
    try:
        return beta - np.linalg.inv(Lpp) @ Lp
    except np.linalg.LinAlgError:
        print(Lpp)
        return np.nan
    # return  beta - np.linalg.inv(Lpp) @ Lp


n_updts = 20
beta0 = np.array(np.random.normal(size = p ))
beta = np.copy(beta0)

epsilon = 1e-6
for i in range(n_updts):
    # i = 0
    beta_old = np.copy(beta)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")

        beta = NewtonRaphsonStep(sample_train, label_train, beta)

        if np.any( np.isnan(beta) ):
            print('beta re-redifined')
            beta = np.random.normal(size = p )

        error = error_est(sample_test, label_test, beta, norm=True)
        print("update = %i, error = %s" %(i, round(error, 5))) 

    if error <= epsilon:
        print('error convergence')
        break

    beta_norm = np.linalg.norm(beta_old - beta, ord = np.inf )
    if np.any( beta_norm <= epsilon  ):
        print('coefficient convergence')
        break

y_pred = (p1(sample_test, beta) >= 0.5).astype(float)

from sklearn.metrics import confusion_matrix
confusion_matrix(label_test,y_pred)

