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


import os
import pandas
import requests
from scipy.cluster.hierarchy import fcluster
from sklearn.cluster import KMeans
from sklearn.preprocessing import scale
from sklearn_extra.cluster import KMedoids
from statistics import mean
import sys

from functions import *


def helper(tup):
    '''
    This is a helper function that takes an output from functions.measure_ip_stability() and summarizes it

    INPUT:
    tup: output from functions.measure_ip_stability (a tuple (is_stable,nr_unstable,violation,closest_cluster))

    OUTPUT:
    nr_unstable: the number of points that are not stable
    max_violation: the maximum violation of any given point
    mean_violation: the average violation among unstable points
    all these are defined at the beginning of Section 6 in the paper
    '''

    violated = list(x for x in tup[2] if x > 1)
    if len(violated) == 0:
        nr_unstable, max_violation, mean_violation = tup[1], np.amax(tup[2]), 0
    else:
        nr_unstable, max_violation, mean_violation = tup[1], np.amax(tup[2]), mean((x for x in tup[2] if x > 1))
    return nr_unstable, max_violation, mean_violation


def run_experiment(data_set, dist_function, k_range, nr_runs):
    '''
    This function runs the various clustering algorithms on a given dataset and for a given distance function and a
    given range for the number of clusters k. It saves the results as a npz-file.

    INPUT:
    data_set: either 'Adult', 'Drug' or 'IndianLiver'
    dist_function: one of the metrics listed on
                    https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.pdist.html
    k_range: an array providing the range for the number of clusters k
    nr_runs: how often to run k-means++ and k-medoids
    '''

    results_folder = 'Results/' + data_set
    if (not os.path.exists(results_folder)):
        if (not os.path.exists('Results')):
            os.mkdir('Results')
        os.mkdir(results_folder)

    # reading the data; the dataset is automatically downloaded if it does't exist
    if data_set == 'Adult':

        if (not os.path.exists('adult.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/adult/adult.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('adult.data', 'wb').write(r.content)

        n = 1000
        df = pandas.read_csv('adult.data', sep=',', header=None)
        df = df[:n]
        ### protected_attribute is only needed when we run the fair k-center algorithm
        protected_attribute = np.array(pandas.get_dummies(df.iloc[:, 9], drop_first=True).values, dtype=int).flatten()
        cont_types = np.where(df.dtypes == 'int')[0]  # =[0,2,4,10,11,12]
        df = df.iloc[:, cont_types]
        data = np.array(df.values, dtype=float)

    if data_set == 'Drug':

        if (not os.path.exists('drug_consumption.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/00373/drug_consumption.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('drug_consumption.data', 'wb').write(r.content)

        df = pandas.read_csv('drug_consumption.data', sep=',', header=None)
        ### protected_attribute is only needed when we run the fair k-center algorithm
        protected_attribute = -np.array(pandas.get_dummies(df.iloc[:, 2], drop_first=True).values,
                                        dtype=int).flatten() + 1
        df.drop(columns=[0], inplace=True)
        cont_types = np.where(df.dtypes == 'float')[0]
        df = df.iloc[:, cont_types]
        data = np.array(df.values, dtype=float)
        n = data.shape[0]  # n=1885

    if data_set == 'IndianLiver':

        if (not os.path.exists('IndianLiver.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/00225/' +
                             'Indian%20Liver%20Patient%20Dataset%20(ILPD).csv', 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('IndianLiver.data', 'wb').write(r.content)

        df = pandas.read_csv('IndianLiver.data', sep=',', header=None)
        df.dropna(inplace=True)
        ### protected_attribute is only needed when we run the fair k-center algorithm
        protected_attribute = np.array(pandas.get_dummies(df.iloc[:, 1], drop_first=True).values, dtype=int).flatten()
        df = pandas.get_dummies(df)
        data = np.array(df.values, dtype=float)
        n = data.shape[0]  # n=579

    # normalize data
    data = scale(data)

    # computing the distance matrix
    distance_matrix = pdist(data, metric=dist_function)

    ### INITIATE THE MEASURES ###
    nr_unf_random = np.zeros((nr_runs, k_range.size))
    nr_unf_kmeans = np.zeros((nr_runs, k_range.size))
    nr_unf_single_linkage = np.zeros(k_range.size)
    nr_unf_average_linkage = np.zeros(k_range.size)
    nr_unf_complete_linkage = np.zeros(k_range.size)
    nr_unf_kmedoids = np.zeros((nr_runs, k_range.size))

    nr_unf_single_linkage_GreedyPruning = np.zeros(k_range.size)
    nr_unf_average_linkage_GreedyPruning = np.zeros(k_range.size)
    nr_unf_complete_linkage_GreedyPruning = np.zeros(k_range.size)

    nr_unf_single_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    nr_unf_average_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    nr_unf_complete_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)

    ###

    maxviol_random = np.zeros((nr_runs, k_range.size))
    maxviol_kmeans = np.zeros((nr_runs, k_range.size))
    maxviol_single_linkage = np.zeros(k_range.size)
    maxviol_average_linkage = np.zeros(k_range.size)
    maxviol_complete_linkage = np.zeros(k_range.size)
    maxviol_kmedoids = np.zeros((nr_runs, k_range.size))

    maxviol_single_linkage_GreedyPruning = np.zeros(k_range.size)
    maxviol_average_linkage_GreedyPruning = np.zeros(k_range.size)
    maxviol_complete_linkage_GreedyPruning = np.zeros(k_range.size)

    maxviol_single_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    maxviol_average_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    maxviol_complete_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)

    ###

    meanviol_random = np.zeros((nr_runs, k_range.size))
    meanviol_kmeans = np.zeros((nr_runs, k_range.size))
    meanviol_single_linkage = np.zeros(k_range.size)
    meanviol_average_linkage = np.zeros(k_range.size)
    meanviol_complete_linkage = np.zeros(k_range.size)
    meanviol_kmedoids = np.zeros((nr_runs, k_range.size))

    meanviol_single_linkage_GreedyPruning = np.zeros(k_range.size)
    meanviol_average_linkage_GreedyPruning = np.zeros(k_range.size)
    meanviol_complete_linkage_GreedyPruning = np.zeros(k_range.size)

    meanviol_single_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    meanviol_average_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    meanviol_complete_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)

    ###
    cost_random = np.zeros((nr_runs, k_range.size))
    cost_kmeans = np.zeros((nr_runs, k_range.size))
    cost_single_linkage = np.zeros(k_range.size)
    cost_average_linkage = np.zeros(k_range.size)
    cost_complete_linkage = np.zeros(k_range.size)
    cost_kmedoids = np.zeros((nr_runs, k_range.size))

    cost_single_linkage_GreedyPruning = np.zeros(k_range.size)
    cost_average_linkage_GreedyPruning = np.zeros(k_range.size)
    cost_complete_linkage_GreedyPruning = np.zeros(k_range.size)

    cost_single_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    cost_average_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)
    cost_complete_linkage_GreedyPruning_criterion2 = np.zeros(k_range.size)

    np.random.seed(12345)

    # Repeat the experiment with different values of k
    for ell, k in enumerate(k_range):

        print('k=' + str(k))

        single_link_clustering = fcluster(linkage(distance_matrix, method='single'), k, criterion='maxclust')
        average_link_clustering = fcluster(linkage(distance_matrix, method='average'), k, criterion='maxclust')
        complete_link_clustering = fcluster(linkage(distance_matrix, method='complete'), k, criterion='maxclust')

        single_link_GreedyPruning_clustering = linkage_clustering_with_greedy_pruning(distance_matrix, k)
        single_link_GreedyPruning_clustering_criterion2 = linkage_clustering_with_greedy_pruning(distance_matrix, k,
                                                                                                 criterion='violation')

        average_link_GreedyPruning_clustering = linkage_clustering_with_greedy_pruning(distance_matrix, k,
                                                                                       method='average')
        average_link_GreedyPruning_clustering_criterion2 = linkage_clustering_with_greedy_pruning(distance_matrix, k,
                                                                                                  method='average',
                                                                                                  criterion='violation')

        complete_link_GreedyPruning_clustering = linkage_clustering_with_greedy_pruning(distance_matrix, k,
                                                                                        method='complete')
        complete_link_GreedyPruning_clustering_criterion2 = linkage_clustering_with_greedy_pruning(distance_matrix, k,
                                                                                                   method='complete',
                                                                                                   criterion='violation')

        ###
        nr_unf_single_linkage[ell], maxviol_single_linkage[ell], meanviol_single_linkage[ell] = helper(
            measure_ip_stability(distance_matrix,
                                 single_link_clustering))
        nr_unf_average_linkage[ell], maxviol_average_linkage[ell], meanviol_average_linkage[ell] = helper(
            measure_ip_stability(distance_matrix,
                                 average_link_clustering))
        nr_unf_complete_linkage[ell], maxviol_complete_linkage[ell], meanviol_complete_linkage[ell] = helper(
            measure_ip_stability(distance_matrix,
                                 complete_link_clustering))

        nr_unf_single_linkage_GreedyPruning[ell], maxviol_single_linkage_GreedyPruning[ell], \
        meanviol_single_linkage_GreedyPruning[ell] = helper(measure_ip_stability(distance_matrix,
                                                                                 single_link_GreedyPruning_clustering))
        nr_unf_average_linkage_GreedyPruning[ell], maxviol_average_linkage_GreedyPruning[ell], \
        meanviol_average_linkage_GreedyPruning[ell] = helper(measure_ip_stability(distance_matrix,
                                                                                  average_link_GreedyPruning_clustering))
        nr_unf_complete_linkage_GreedyPruning[ell], maxviol_complete_linkage_GreedyPruning[ell], \
        meanviol_complete_linkage_GreedyPruning[ell] = helper(measure_ip_stability(distance_matrix,
                                                                                   complete_link_GreedyPruning_clustering))

        nr_unf_single_linkage_GreedyPruning_criterion2[ell], maxviol_single_linkage_GreedyPruning_criterion2[
            ell], meanviol_single_linkage_GreedyPruning_criterion2[
            ell] = helper(measure_ip_stability(distance_matrix,
                                               single_link_GreedyPruning_clustering_criterion2))
        nr_unf_average_linkage_GreedyPruning_criterion2[ell], maxviol_average_linkage_GreedyPruning_criterion2[
            ell], meanviol_average_linkage_GreedyPruning_criterion2[
            ell] = helper(measure_ip_stability(distance_matrix,
                                               average_link_GreedyPruning_clustering_criterion2))
        nr_unf_complete_linkage_GreedyPruning_criterion2[ell], maxviol_complete_linkage_GreedyPruning_criterion2[
            ell], meanviol_complete_linkage_GreedyPruning_criterion2[
            ell] = helper(measure_ip_stability(distance_matrix,
                                               complete_link_GreedyPruning_clustering_criterion2))

        ###
        cost_single_linkage[ell] = compute_clustering_cost(distance_matrix, single_link_clustering)
        cost_average_linkage[ell] = compute_clustering_cost(distance_matrix, average_link_clustering)
        cost_complete_linkage[ell] = compute_clustering_cost(distance_matrix, complete_link_clustering)

        cost_single_linkage_GreedyPruning[ell] = compute_clustering_cost(distance_matrix,
                                                                         single_link_GreedyPruning_clustering)
        cost_average_linkage_GreedyPruning[ell] = compute_clustering_cost(distance_matrix,
                                                                          average_link_GreedyPruning_clustering)
        cost_complete_linkage_GreedyPruning[ell] = compute_clustering_cost(distance_matrix,
                                                                           complete_link_GreedyPruning_clustering)

        cost_single_linkage_GreedyPruning_criterion2[ell] = compute_clustering_cost(distance_matrix,
                                                                        single_link_GreedyPruning_clustering_criterion2)
        cost_average_linkage_GreedyPruning_criterion2[ell] = compute_clustering_cost(distance_matrix,
                                                                        average_link_GreedyPruning_clustering_criterion2)
        cost_complete_linkage_GreedyPruning_criterion2[ell] = compute_clustering_cost(distance_matrix,
                                                                        complete_link_GreedyPruning_clustering_criterion2)

        for mmm in np.arange(nr_runs):
            random_clustering = np.random.randint(k, size=n)
            k_means_clustering = KMeans(n_clusters=k).fit(data).labels_
            k_medoids_clustering = KMedoids(n_clusters=k, metric=dist_function, init='k-medoids++').fit(data).labels_

            ##
            nr_unf_random[mmm, ell], maxviol_random[mmm, ell], meanviol_random[mmm, ell] = \
                helper(measure_ip_stability(distance_matrix, random_clustering))
            nr_unf_kmeans[mmm, ell], maxviol_kmeans[mmm, ell], meanviol_kmeans[mmm, ell] = \
                helper(measure_ip_stability(distance_matrix, k_means_clustering))
            nr_unf_kmedoids[mmm, ell], maxviol_kmedoids[mmm, ell], meanviol_kmedoids[mmm, ell] = \
                helper(measure_ip_stability(distance_matrix, k_medoids_clustering))

            ##
            cost_random[mmm, ell] = compute_clustering_cost(distance_matrix, random_clustering)
            cost_kmeans[mmm, ell] = compute_clustering_cost(distance_matrix, k_means_clustering)
            cost_kmedoids[mmm, ell] = compute_clustering_cost(distance_matrix, k_medoids_clustering)

    # save data
    save_str = {'Adult': 'AdultData', 'Drug': 'DrugData', 'IndianLiver': 'IndianLiver'}[data_set]

    np.savez(results_folder + '/' + save_str + '_NrReplicates=' + str(nr_runs) + '_distance=' + dist_function,
             data_set=data_set,
             nr_runs=nr_runs,
             dist_function=dist_function,
             k_range=k_range,
             n=n,
             protected_attribute=protected_attribute,
             #
             nr_unf_random=nr_unf_random,
             nr_unf_kmeans=nr_unf_kmeans,
             nr_unf_single_linkage=nr_unf_single_linkage,
             nr_unf_average_linkage=nr_unf_average_linkage,
             nr_unf_complete_linkage=nr_unf_complete_linkage,
             nr_unf_kmedoids=nr_unf_kmedoids,
             nr_unf_single_linkage_GreedyPruning=nr_unf_single_linkage_GreedyPruning,
             nr_unf_average_linkage_GreedyPruning=nr_unf_average_linkage_GreedyPruning,
             nr_unf_complete_linkage_GreedyPruning=nr_unf_complete_linkage_GreedyPruning,
             nr_unf_single_linkage_GreedyPruning_criterion2=nr_unf_single_linkage_GreedyPruning_criterion2,
             nr_unf_average_linkage_GreedyPruning_criterion2=nr_unf_average_linkage_GreedyPruning_criterion2,
             nr_unf_complete_linkage_GreedyPruning_criterion2=nr_unf_complete_linkage_GreedyPruning_criterion2,
             #
             maxviol_random=maxviol_random,
             maxviol_kmeans=maxviol_kmeans,
             maxviol_single_linkage=maxviol_single_linkage,
             maxviol_average_linkage=maxviol_average_linkage,
             maxviol_complete_linkage=maxviol_complete_linkage,
             maxviol_kmedoids=maxviol_kmedoids,
             maxviol_single_linkage_GreedyPruning=maxviol_single_linkage_GreedyPruning,
             maxviol_average_linkage_GreedyPruning=maxviol_average_linkage_GreedyPruning,
             maxviol_complete_linkage_GreedyPruning=maxviol_complete_linkage_GreedyPruning,
             maxviol_single_linkage_GreedyPruning_criterion2=maxviol_single_linkage_GreedyPruning_criterion2,
             maxviol_average_linkage_GreedyPruning_criterion2=maxviol_average_linkage_GreedyPruning_criterion2,
             maxviol_complete_linkage_GreedyPruning_criterion2=maxviol_complete_linkage_GreedyPruning_criterion2,
             #
             meanviol_random=meanviol_random,
             meanviol_kmeans=meanviol_kmeans,
             meanviol_single_linkage=meanviol_single_linkage,
             meanviol_average_linkage=meanviol_average_linkage,
             meanviol_complete_linkage=meanviol_complete_linkage,
             meanviol_kmedoids=meanviol_kmedoids,
             meanviol_single_linkage_GreedyPruning=meanviol_single_linkage_GreedyPruning,
             meanviol_average_linkage_GreedyPruning=meanviol_average_linkage_GreedyPruning,
             meanviol_complete_linkage_GreedyPruning=meanviol_complete_linkage_GreedyPruning,
             meanviol_single_linkage_GreedyPruning_criterion2=meanviol_single_linkage_GreedyPruning_criterion2,
             meanviol_average_linkage_GreedyPruning_criterion2=meanviol_average_linkage_GreedyPruning_criterion2,
             meanviol_complete_linkage_GreedyPruning_criterion2=meanviol_complete_linkage_GreedyPruning_criterion2,
             #
             cost_random=cost_random,
             cost_kmeans=cost_kmeans,
             cost_single_linkage=cost_single_linkage,
             cost_average_linkage=cost_average_linkage,
             cost_complete_linkage=cost_complete_linkage,
             cost_kmedoids=cost_kmedoids,
             cost_single_linkage_GreedyPruning=cost_single_linkage_GreedyPruning,
             cost_average_linkage_GreedyPruning=cost_average_linkage_GreedyPruning,
             cost_complete_linkage_GreedyPruning=cost_complete_linkage_GreedyPruning,
             cost_single_linkage_GreedyPruning_criterion2=cost_single_linkage_GreedyPruning_criterion2,
             cost_average_linkage_GreedyPruning_criterion2=cost_average_linkage_GreedyPruning_criterion2,
             cost_complete_linkage_GreedyPruning_criterion2=cost_complete_linkage_GreedyPruning_criterion2)
