import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from scipy.stats import multivariate_normal
from sklearn.metrics import zero_one_loss
from sklearn.base import clone

def compute_zero_one_loss(true_labels, predicted_labels):
    return zero_one_loss(true_labels, predicted_labels)

# Margin Noise
def generate_data_margin(m_distance, dim = 5, num_samples = 1000, random_seed = 0):
    np.random.seed(random_seed)
    m = m_distance / 2
    mean_0 = np.array([-m]*dim)
    cov_0 = np.identity(dim)

    mean_1 = np.array([m]*dim)
    cov_1 = np.identity(dim)

    data_0 = np.random.multivariate_normal(mean_0, cov_0, num_samples)
    data_1 = np.random.multivariate_normal(mean_1, cov_1, num_samples)

    labels_0 = np.zeros(num_samples)
    labels_1 = np.ones(num_samples)

    data = np.vstack((data_0, data_1))
    labels = np.concatenate((labels_0, labels_1))

    return data, labels


range_dist = [0.5, 3]

m_values = np.linspace(range_dist[0], range_dist[1], 100)

num_iterations = 10
num_dimensions = 7
num_points = 1000

plt.figure(figsize = (5,2.5))

for num_dimensions in [3,5,7,10]:

    data, labels = generate_data_margin(range_dist[1], num_dimensions, num_points, 0)

    model = LogisticRegression()
    model.fit(data, labels)
    predicted_labels = model.predict(data)
    zero_one_loss_value = compute_zero_one_loss(labels, predicted_labels)

    #print("model", zero_one_loss_value)

    # Get the original coefficients
    original_coefficients = model.coef_[0]
    # Add a random vector to the coefficients
    random_vector = np.random.normal(0, 0.2, size=original_coefficients.shape)
    new_coefficients = original_coefficients + random_vector

    # Create a new model with the modified coefficients
    new_model = LogisticRegression()
    new_model.fit(data, labels)
    new_model.coef_ = np.array([new_coefficients])
    new_model.intercept_ = model.intercept_

    predicted_labels = new_model.predict(data)
    zero_one_loss_value = compute_zero_one_loss(labels, predicted_labels)
    #print(model.coef_, new_model.coef_)

    #print("new model", zero_one_loss_value)

    model = new_model

    to_plot = []
    for m in m_values:
        variance_array = []

        for k in range(num_iterations):
            data, labels = generate_data_margin(m, num_dimensions, num_points, k)

            #model = LogisticRegression()
            #model.fit(data, labels)

            predicted_labels = model.predict(data)
            zero_one_loss_value = compute_zero_one_loss(labels, predicted_labels)
            #zero_one_losses.append(zero_one_loss_value)

            #variance = zero_one_loss_value * (1 - zero_one_loss_value)
            variance = np.var((labels == predicted_labels).astype(int))
            variance_array += [variance]

        variance_mean = np.mean(variance_array)
        variance_std = np.std(variance_array)
        to_plot.append((variance_mean, variance_std))

    to_plot_mean = np.array(to_plot)[:, 0]
    to_plot_std = np.array(to_plot)[:, 1]

    plt.plot(m_values, to_plot_mean, label=str(num_dimensions) + ' dim')
    plt.fill_between(m_values, to_plot_mean-to_plot_std, to_plot_mean+to_plot_std, alpha = 0.5)

    plt.xlabel('Margin noise, m', fontsize = 18)
    plt.ylabel('Var of the loss', fontsize = 18)
    

plt.legend(loc = "upper left", fontsize = 16)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16)
plt.gca().invert_xaxis()
plt.savefig("g_margin.png", bbox_inches='tight', dpi = 200)

#Attribute noise
def generate_data_attribute(mean1, mean2, dim = 5, num_samples = 1000, random_seed = 0):
    np.random.seed(random_seed)
    mean_0 = np.array([mean1]*dim)
    cov_0 = np.identity(dim)

    mean_1 = np.array([mean2]*dim)
    cov_1 = np.identity(dim)

    data_0 = np.random.multivariate_normal(mean_0, cov_0, num_samples)
    data_1 = np.random.multivariate_normal(mean_1, cov_1, num_samples)

    labels_0 = np.zeros(num_samples)
    labels_1 = np.ones(num_samples)

    data = np.vstack((data_0, data_1))
    labels = np.concatenate((labels_0, labels_1))

    return data, labels


noise_range = [0, 0.5]
noise_values = np.linspace(noise_range[0], noise_range[1], 100)

mean1 = 0
mean2 = 2

num_iterations = 10
num_dimensions = 7
num_points = 1000

plt.figure(figsize = (5,2.5))

for num_dimensions in [3,5,7,10]:

    data, labels = generate_data_attribute(mean1, mean2, num_dimensions, num_points, 0)

    model = LogisticRegression()
    model.fit(data, labels)
    predicted_labels = model.predict(data)
    zero_one_loss_value = compute_zero_one_loss(labels, predicted_labels)

    #print("model", zero_one_loss_value)

    # Get the original coefficients
    original_coefficients = model.coef_[0]
    # Add a random vector to the coefficients
    random_vector = np.random.normal(0, 0.2, size=original_coefficients.shape)
    new_coefficients = original_coefficients + random_vector

    # Create a new model with the modified coefficients
    new_model = LogisticRegression()
    new_model.fit(data, labels)
    new_model.coef_ = np.array([new_coefficients])
    new_model.intercept_ = model.intercept_

    predicted_labels = new_model.predict(data)
    zero_one_loss_value = compute_zero_one_loss(labels, predicted_labels)
    #print(model.coef_, new_model.coef_)

    #print("new model", zero_one_loss_value)

    model = new_model

    to_plot = []
    for noise in noise_values:
        variance_array = []

        for k in range(num_iterations):
            mean_1 = np.array([0]*num_dimensions)
            cov_1 = noise*np.identity(num_dimensions)

            random = np.random.multivariate_normal(mean_1, cov_1, 2*num_points)
            data_noise = data + random

            predicted_labels = model.predict(data_noise)
            zero_one_loss_value = compute_zero_one_loss(labels, predicted_labels)

            variance = np.var((labels == predicted_labels).astype(int))
            variance_array += [variance]

        variance_mean = np.mean(variance_array)
        variance_std = np.std(variance_array)
        to_plot.append((variance_mean, variance_std))

    to_plot_mean = np.array(to_plot)[:, 0]
    to_plot_std = np.array(to_plot)[:, 1]

    plt.plot(noise_values, to_plot_mean, label=str(num_dimensions) + ' dim')
    plt.fill_between(noise_values, to_plot_mean-to_plot_std, to_plot_mean+to_plot_std, alpha = 0.5)

    plt.xlabel(r'Additive attribute noise, $\sigma$', fontsize = 18)
    plt.ylabel('Var of the loss', fontsize = 18)


plt.legend(loc = "upper left", fontsize = 16)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16)
#plt.gca().invert_xaxis()
plt.savefig("g_additive.png", bbox_inches='tight', dpi = 200)