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

def make_circles(rho, num_samples, seed, test_num = 0.5):

    np.random.seed(seed)

    # # Step 1: Define correlation structure for observations
    # num_samples = 500
    # 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
    num_features = 2
    X_uncorrelated = np.random.randn(num_samples, num_features)

    # 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 = .01  # Set 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]


    Indx = int(test_num*X.shape[0])
    X_train = X[:Indx]
    X_test = X[Indx:]
    y_train = y[:Indx]
    y_test = y[Indx:]
    V_train = V[:Indx, :Indx]
    V_test = V[Indx:, Indx:]

    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, k_fold_cv_random_vcv, k_fold_cv_random_vcv_v2
import matplotlib.pyplot as plt


(X_train, X_test, 
 y_train, y_test, 
 vcv_train, vcv_test) = make_circles(rho=0.5, 
                                 num_samples=500, 
                                 seed=12038, 
                                 test_num=0.5)


lkrr = Logistic(get_corr=False, tol=1e-4, 
               max_iter=400, kernel='rbf', 
               verbose=True)

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

best_params = k_fold_cv_random_vcv_v2(X_train, y_train, vcv_train, 
                                   lkrr, params,
                                   folds=3, sample=100, 
                                   verbose=True, seed=12038)
print('Best params: ', best_params)
lkrr.set_params(**best_params)
lkrr.fit(X_train, y_train, vcv_train) 
print(lkrr.score(X_test, y_test))


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)

print( 'Loss: ', round(log_reg_loss, 2) )

tmp_axes =  [-4, 4, -4, 4]
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
plt.sca(axes[0])
plot_predictions(log_reg, tmp_axes)
plot_dataset(X_test, y_test, tmp_axes)

plt.sca(axes[1])
plot_predictions(lkrr, tmp_axes)
plot_dataset(X_test, y_test, tmp_axes)
plt.tight_layout()
# plt.savefig('classification_comparisons.png', dpi=300)



