# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0


import os
import pandas
import requests
from sklearn.cluster import KMeans
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler
import sys

from functions import *

###############################################################################################
# Script for the experiments of Section 6.2 / Appendix H on the German Credit data set ########
###############################################################################################


### SET PARAMETERS ###############################################################

nr_decimals = 2  # number of decimals for printing
nr_runs_kmeans = 100  # for k-means++, we report average results
dataset = 'predictions'  # 'credit_amount' (size=1000) or 'predictions' (size=500)

k_range = [5, 10, 20, 50]  # numbers of clusters to consider
target = 'uniform'  # either 'uniform' or a list (of length k_range) of lists (with the i-th list being of
# length k_range[i] and its elements summing up to the size of the data set specifying the target cluster sizes

# k_range=[12,12,20,20]
# target=[[25,25,50,50,50,50,50,50,50,50,25,25],[10,15,25,50,50,100,100,50,50,25,15,10],
#        np.tile([10,40],10),np.append(15*np.ones(9),np.append(115*np.ones(2),15*np.ones(9)))]
##################################################################################


# Download the data set if it doesn't exist already
if (not os.path.exists('german.data')):
    print('Data set does not exist in current folder --- have to download it')
    r = requests.get('https://archive.ics.uci.edu/ml/machine-learning-databases/statlog/german/german.data',
                     allow_redirects=True)
    if r.status_code == requests.codes.ok:
        print('Download successful\n')
    else:
        print('Could not download the data set --- please download it manually')
        sys.exit()
    open('german.data', 'wb').write(r.content)

np.random.seed(12345)

if not target == 'uniform':
    assert len(k_range) == len(target), 'The length of k_range must equal the length of target'

if dataset == 'credit_amount':
    df = pandas.read_csv('german.data', sep=' ', header=None)
    data = np.sort(df.iloc[:, 4].values)
    data = data / np.amax(data)
    n = 1000
else:
    n = 500

    # computing the scores if they don't exist already
    if not os.path.exists('1dimexp_NN_scores.csv'):
        print('have to compute the scores')
        df = pandas.read_csv('german.data', sep=' ', header=None)
        label = df.iloc[:, 20].values
        df.drop(20, axis=1, inplace=True)
        df = pandas.get_dummies(df)
        data_matrix = StandardScaler().fit_transform(df.values)
        train_data = data_matrix[:(1000 - n), :]
        train_label = label[:(1000 - n)]
        test_data = data_matrix[(1000 - n):, :]
        test_label = label[(1000 - n):]
        MLP = MLPClassifier(hidden_layer_sizes=(100, 50, 20), random_state=99999).fit(train_data, train_label)
        predictions = MLP.predict(test_data)
        scores = MLP.predict_proba(test_data)[:, 0]
        print('Test accuracy=' + str(np.mean(predictions == test_label)))
        np.savetxt('1dimexp_NN_scores.csv', scores)

    scores = np.loadtxt('1dimexp_NN_scores.csv')
    data = np.sort(scores)

dmat = squareform(pdist(data.reshape((n, 1))))

for ccc, k in enumerate(k_range):

    print('\n######################################################\n')

    if target == 'uniform':
        print('k=' + str(k))
        target_sizes = ((n / k) * np.ones(k)).astype(int)
    else:
        target_sizes = np.array(target[ccc], dtype=int)
        print('k=' + str(k) + ', Target cluster sizes=' + str(target_sizes))

    assert (n == np.sum(target_sizes)), 'Summing over all cluster sizes must yield the number of data points'
    assert (k == target_sizes.size), 'k must equal the length of target sizes'

    ### NAIVE ###
    clustering_with_target_sizes = np.zeros(n)
    clustering_with_target_sizes[0:target_sizes[0]] = 0
    for mmm in np.arange(1, k):
        clustering_with_target_sizes[np.sum(target_sizes[:mmm]):np.sum(target_sizes[:mmm + 1])] = mmm
    _, nr_unfair, violation, _ = measure_ip_stability(dmat, clustering_with_target_sizes)

    print('NAIVE: # Uns=' + str(nr_unfair) +
          ', Max Viol=' + str(round(np.amax(violation), nr_decimals)) +
          ', Mean Viol=' + str(round(np.mean(violation[violation > 1]), nr_decimals)) +
          ', Objective=0' +
          ', Cost=' + str(round(compute_clustering_cost(dmat, clustering_with_target_sizes), nr_decimals)))

    #### DYNAMIC PROGRAMMING APPROACH ###
    clustering_DP, clustering_DP_cluster_sizes = DynamicProgramming_for_1Dimension_PInfinity(data, k, target_sizes)
    _, nr_unfair, violation, _ = measure_ip_stability(dmat, clustering_DP)
    max_dev = np.amax(np.abs(clustering_DP_cluster_sizes - target_sizes))

    print('DP APPROACH: # Uns=' + str(nr_unfair) +
          ', Max Viol=' + str(round(np.amax(violation), nr_decimals)) +
          ', Mean Viol=0' +
          ', Objective=' + str(round(max_dev, nr_decimals)) +
          ', Cost=' + str(round(compute_clustering_cost(dmat, clustering_DP), nr_decimals)))

    ### K-MEANS WITH INITIALIZATION ###
    if target == 'uniform':
        kmeans_init = data[int(target_sizes[0] / 2 - 1)::target_sizes[0]].reshape((k, 1))
    else:
        kmeans_init = data[int(target_sizes[0] / 2 - 1)]
        count_el = target_sizes[0]
        for mmm in np.arange(1, k):
            kmeans_init = np.append(kmeans_init, data[int(count_el + target_sizes[mmm] / 2)])
            count_el = count_el + target_sizes[mmm]
        kmeans_init = kmeans_init.reshape((k, 1))

    k_means_clustering = KMeans(n_clusters=k, init=kmeans_init, n_init=1).fit(data.reshape((n, 1))).labels_
    t1, t2, t3 = np.unique(k_means_clustering, return_index=True, return_counts=True)
    cl_sizes = t3[np.argsort(t2)]
    _, nr_unfair, violation, _ = measure_ip_stability(dmat, k_means_clustering)
    results_k_means_max_dev_from_cluster_sizes = np.amax(np.abs(cl_sizes - target_sizes))

    print('k-MEANS: # Uns=' + str(nr_unfair) +
          ', Max Viol=' + str(round(np.amax(violation), nr_decimals)) +
          ', Mean Viol=' + str(round(np.mean(violation[violation > 1]), nr_decimals)) +
          ', Objective=' + str(results_k_means_max_dev_from_cluster_sizes) +
          ', Cost= ' + str(round(compute_clustering_cost(dmat, k_means_clustering), nr_decimals)))

    ### K-MEANS++ ###
    kmeans_init = 'k-means++'
    results_k_means_nr_unfair = np.zeros(nr_runs_kmeans)
    results_k_means_max_viol = np.zeros(nr_runs_kmeans)
    results_k_means_mean_viol = np.zeros(nr_runs_kmeans)
    results_k_means_max_dev_from_cluster_sizes = np.zeros(nr_runs_kmeans)
    results_k_means_cost = np.zeros(nr_runs_kmeans)

    for mmm in np.arange(nr_runs_kmeans):
        k_means_clustering = KMeans(n_clusters=k, init=kmeans_init).fit(data.reshape((n, 1))).labels_
        t1, t2, t3 = np.unique(k_means_clustering, return_index=True, return_counts=True)
        cl_sizes = t3[np.argsort(t2)]
        _, nr_unfair, violation, _ = measure_ip_stability(dmat, k_means_clustering)

        results_k_means_nr_unfair[mmm] = nr_unfair
        results_k_means_max_viol[mmm] = np.amax(violation)
        if np.sum(violation > 1) > 0:
            results_k_means_mean_viol[mmm] = np.mean(violation[violation > 1])
        results_k_means_max_dev_from_cluster_sizes[mmm] = np.amax(np.abs(cl_sizes - target_sizes))
        results_k_means_cost[mmm] = compute_clustering_cost(dmat, k_means_clustering)

    print('k-MEANS++: # Uns=' + str(np.mean(results_k_means_nr_unfair)) +
          ', Max Viol=' + str(round(np.mean(results_k_means_max_viol), nr_decimals)) +
          ', Mean Viol=' + str(round(np.mean(results_k_means_mean_viol), nr_decimals)) +
          ', Objective=' + str(round(np.mean(results_k_means_max_dev_from_cluster_sizes), nr_decimals)) +
          ', Cost= ' + str(round(np.mean(results_k_means_cost), nr_decimals)))
