import numpy as np
import matplotlib.pyplot as plt
from numpy.linalg import cholesky


def pre_process_V(V):
    n, _ = V.shape
    V = V/np.max(V)
    detV_n = np.linalg.det(V)**(1/n)
    print('detV_n: ', detV_n)
    V = V/detV_n if detV_n != 0 else V
    return V

def make_circles(rho, num_samples, seed, test_num = 0.5, sigma=1, radius=0.01):

    np.random.seed(seed)
    # # Step 1: Define correlation structure for observations
    # num_samples = 5
    # rho = 0.5  # Strength of correlation between consecutive observations

    # Construct correlation matrix (exponential decay)
    V = np.array([[rho**abs(i-j) for j in range(num_samples)] for i in range(num_samples)])

    # Step 2: Generate independent Gaussian features
    # noise = 0.1  # Add noise to the observations
    num_features = 2
    X_uncorrelated = np.random.randn(num_samples, num_features) * sigma

    # Step 3: Apply Cholesky decomposition to correlate observations
    L = cholesky(V)
    X = L @ X_uncorrelated   # Now observations are correlated

    # Step 4: Define a **circular non-linearity** (donut shape)
    # radius sets the radius for the decision boundary

    decision_boundary = np.sqrt(X[:, 0]**2 + X[:, 1]**2)  # Compute distance from origin
    y = ((decision_boundary > radius) & (decision_boundary < radius + 1)).astype(int)  # Define labels

    # Step 5: Shuffle the data and split into training and testing sets
    all_index = list(range(X.shape[0]))
    np.random.shuffle(all_index)
    
    X = X[all_index]
    y = y[all_index]
    V = V[all_index, :][:, all_index]


    X = X[all_index,:]
    y = y[all_index]
    V = V[all_index, :][:, all_index]

    Indx = int(test_num*X.shape[0])
    test_idx = np.random.choice(all_index, Indx, replace=False)
    train_idx = list(set(all_index) - set(test_idx))

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

    V_train = V[train_idx,:][:,train_idx]
    V_test = V[test_idx,:][:,test_idx]
    
    return X_train, X_test, y_train, y_test, V_train, V_test


def plot_dataset(X, y, axes):
    plt.plot(X[:, 0][y==0], X[:, 1][y==0], "bs")
    plt.plot(X[:, 0][y==1], X[:, 1][y==1], "g^")
    plt.axis(axes)
    plt.grid(True, which='major', linestyle='--', linewidth='0.5')
    plt.xlabel(r"$x_1$", fontsize=20)
    plt.ylabel(r"$x_2$", fontsize=20, rotation=0)

def plot_predictions(clf, axes):
    x0s = np.linspace(axes[0], axes[1], 100)
    x1s = np.linspace(axes[2], axes[3], 100)
    x0, x1 = np.meshgrid(x0s, x1s)
    X = np.c_[x0.ravel(), x1.ravel()]
    y_pred = clf.predict(X).reshape(x0.shape)
    plt.contourf(x0, x1, y_pred, cmap=plt.cm.bwr , alpha=0.3)

from krgee import Logistic
import matplotlib.pyplot as plt
from cross_validation_krgee import (k_fold_cv_random_vcv_v2)

# seed = 1203
seed = 12038
cv_sample = 100
cv_folds = 4
tol = 1e-3
max_iter = 70

(X_train, X_test, 
 y_train, y_test, 
 vcv_train, vcv_test) = make_circles(
                                rho=0.5,
                                num_samples=500,
                                 seed=seed, 
                                 test_num=0.35, 
                                 sigma=1,
                                 radius=.1
                                 )


# plot_dataset(X_train, y_train, np.array([-1, 1, -1, 1])*4)
# plot_dataset(X_test, y_test, np.array([-1, 1, -1, 1])*4)
# vcv_train = pre_process_V(vcv_train)

print('X_train shape: ', X_train.shape)
print('X_test shape: ', X_test.shape)

print('label_ratio train:', np.sum(y_train ==1)/len(y_train))
print('label_ratio test:', np.sum(y_test ==1)/len(y_test))


X_train_offset = X_train.mean(axis=0)
X_train_scale = X_train.std(axis=0)
X_train -= X_train_offset
X_train /= X_train_scale

# X_test -= np.mean(X_test, axis=0)
# X_test /= np.std(X_test, axis=0)

X_test -= X_train_offset
X_test /= X_train_scale


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


I_train = np.eye(X_train.shape[0])
lkrr = Logistic(get_corr=False, tol=tol, 
               max_iter=max_iter, kernel='rbf', 
               verbose=True, seed=seed, 
               fix_phi=True)
# lkrr.phi = 1 


import time
start = time.time()
lkrr_best_params = k_fold_cv_random_vcv_v2(
    X_train, y_train, I_train, 
    lkrr, params,
    folds=cv_folds,
    sample=cv_sample, 
    verbose=True, 
    seed=seed
)
print('lkrr CV Time: ', time.time() - start)


print('lkrr best params: ', lkrr_best_params)
lkrr.set_params(**lkrr_best_params)
lkrr.fit(X_train, y_train, I_train)
lkrr_loss = lkrr.score(X_test, y_test)
print('lkrr loss: ', lkrr_loss)
print('phi: ', lkrr.phi)
print('degrees of freedom: ', lkrr.dof)
print()


lkrgee = Logistic(get_corr=False, tol=tol, 
               max_iter=max_iter, kernel='rbf', 
               verbose=True, seed=seed, 
               fix_phi=True)



import time
start = time.time()
lkrgee_best_params = k_fold_cv_random_vcv_v2(
    X_train, y_train, vcv_train, 
    lkrgee, params,
    folds=cv_folds,
    sample=cv_sample, 
    verbose=True, 
    seed=seed
)
print('lkrgee CV Time: ', time.time() - start)

print('lkrgee best params: ', lkrgee_best_params)
lkrgee.set_params(**lkrgee_best_params)
# lkrgee.dof = emp_df
# lkrgee.verbose = True
# lkrgee._verbose = False
lkrgee.fit(X_train, y_train, vcv_train)
# print("R inv max: ", np.max(lkrgee.R_inv))
# print("R inv min: ", np.min(lkrgee.R_inv))
lkrgee_loss = lkrgee.score(X_test, y_test)
print('lkrgee_loss testing: ', lkrgee_loss)
print('phi: ', lkrgee.phi)
print('degrees of freedom: ', lkrgee.dof)
# print('rho: ', lkrgee.rho)





from sklearn.linear_model import LogisticRegression
log_reg = LogisticRegression()
log_reg.fit(X_train, y_train)
log_reg_loss = 1 -  log_reg.score(X_test, y_test)


# tmp_axes =  [-4, 4, -4, 4]
# fig, axes = plt.subplots(1, 3, figsize=(15, 5))
# plt.sca(axes[0])
# plt.title(f'Logistic Regression (loss: {round(log_reg_loss, 3)})')
# plot_predictions(log_reg, tmp_axes)
# plot_dataset(X_test, y_test, tmp_axes)

# plt.sca(axes[1])
# plt.title(f'Logistic KRR (loss: {round(lkrr_loss, 3)})')
# plot_predictions(lkrr, tmp_axes)
# plot_dataset(X_test, y_test, tmp_axes)

# plt.sca(axes[2])
# plt.title(f'Logistic KRGEE (loss: {round(lkrgee_loss, 3)})')
# plot_predictions(lkrgee, tmp_axes)
# plot_dataset(X_test, y_test, tmp_axes)
# plt.tight_layout()



# # # plt.savefig('classification_model_comparison_v0.png', dpi=300)


# # # # ROC curve
from sklearn.metrics import roc_curve, auc
log_reg_probs = log_reg.predict_proba(X_test)[:, 1]
lkrr_probs = lkrr.predict(X_test)
lkrgee_probs = lkrgee.predict(X_test)

log_reg_fpr, log_reg_tpr, _ = roc_curve(y_test, log_reg_probs)
lkrr_fpr, lkrr_tpr, _ = roc_curve(y_test, lkrr_probs)
lkrgee_fpr, lkrgee_tpr, _ = roc_curve(y_test, lkrgee_probs)

log_reg_auc = auc(log_reg_fpr, log_reg_tpr)
lkrr_auc = auc(lkrr_fpr, lkrr_tpr)
lkrgee_auc = auc(lkrgee_fpr, lkrgee_tpr)




print()
print('log_reg_loss: ', log_reg_loss)
print('lkrr_loss: ', lkrr_loss)
print('lkrgee loss: ', lkrgee_loss)
print()

print('log_reg_auc: ', log_reg_auc)
print('lkrr_auc: ', lkrr_auc)
print('lkrgee_auc: ', lkrgee_auc)


tmp_axes =  [-4, 4, -4, 4]
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
plt.sca(axes[0])
plt.title(f'Logistic Regression (loss: {round(log_reg_loss, 3)})')
plot_predictions(log_reg, tmp_axes)
plot_dataset(X_test, y_test, tmp_axes)

plt.sca(axes[1])
plt.title(f'Logistic KRR (loss: {round(lkrr_loss, 3)})')
plot_predictions(lkrr, tmp_axes)
plot_dataset(X_test, y_test, tmp_axes)

plt.sca(axes[2])
plt.title(f'Logistic KRGEE (loss: {round(lkrgee_loss, 3)})')
plot_predictions(lkrgee, tmp_axes)
plot_dataset(X_test, y_test, tmp_axes)
plt.tight_layout()




# plt.figure(figsize=(8, 6))
# plt.plot(log_reg_fpr, log_reg_tpr, 'b-', label='Logistic Regression (area = %0.4f)' % log_reg_auc)
# plt.plot(lkrr_fpr, lkrr_tpr, 'g-', label='Logistic KRR (area = %0.4f)' % lkrr_auc)
# plt.plot(lkrgee_fpr, lkrgee_tpr, 'r-', label='Logistic KRGEE  (area = %0.4f)' % lkrgee_auc)
# plt.plot([0, 1], [0, 1], 'k--')
# plt.xlabel('False Positive Rate')
# plt.ylabel('True Positive Rate')
# plt.legend()
# plt.grid(True)
# plt.tight_layout()








# # from sklearn.metrics import confusion_matrix
# # log_reg_pred = (log_reg_probs >= 0.5).astype(float)
# # lkrr_pred = (lkrr_probs >= 0.5).astype(float)
# # lkrgee_pred = (lkrgee_probs >= 0.5).astype(float)

# # log_reg_cm = confusion_matrix(y_test, log_reg_pred)
# # lkrr_cm = confusion_matrix(y_test, lkrr_pred)
# # lkrgee_cm = confusion_matrix(y_test, lkrgee_pred)

# # import seaborn as sns

# # fig, axes = plt.subplots(1, 3, figsize=(15, 5))
# # plt.sca(axes[0])
# # plt.title('Logistic Regression')
# # plt.xlabel('Predicted')
# # plt.ylabel('True')
# # # # sns.heatmap(log_reg_cm, annot=True, fmt='g', cmap='Blues', cbar=False)

# # # # plt.sca(axes[1])
# # # # plt.title('Logistic KRR')
# # # # plt.xlabel('Predicted')
# # # # plt.ylabel('True')
# # # # sns.heatmap(lkrr_cm, annot=True, fmt='g', cmap='Blues', cbar=False) 

# # # # plt.sca(axes[2])
# # # # plt.title('Logistic KRGEE')
# # # # plt.xlabel('Predicted')
# # # # plt.ylabel('True')
# # # # sns.heatmap(lkrgee_cm, annot=True, fmt='g', cmap='Blues', cbar=False)
# # # # plt.tight_layout()
# # # # plt.savefig('classification_confusion_matrix.png', dpi=300)





# # # # y_pred = (p1(f_test) >= 0.5).astype(float)
# # # # cm = confusion_matrix(y_test, y_pred).astype(float)
# # # # import seaborn as sns
# # # # cm[0,:] = cm[0,:]/np.sum(cm[0,:])
# # # # cm[1,:] = cm[1,:]/np.sum(cm[1,:])

# # # # cm = np.round(cm, 2)
# # # # # cm[1:,1] = cm[1:,1]/np.sum(cm[1:,1])
# # # # sns.heatmap(cm, annot=True, fmt='g', cmap='Blues', cbar=False)
# # # # plt.xlabel('Predicted')
# # # # plt.ylabel('True')
# # # # plt.title('Confusion Matrix')

