# -*- coding: utf-8 -*-
"""
Part of the GSM-Degroot package.
Copyright (C) 2020-2025, Ivan Conjeaud and Argyris Kalogeratos. This package is a free software distributed under the GPLv3 license.
"""
import numpy as np
import math
import random
import itertools
import os
from multiprocessing import Pool
#os.chdir('folder_name')

#Create the cost maps, comment out if cost_maps already in directory
with open("Cost_maps.py") as file:
    exec(file.read())

#Bootstrap parameters
values_alpha=[0.0001*k for k in range(1,200)]   #Sizes of the BS
K=10

#Optimization parameters: should match with Cost_maps.py
min_mu=-200
max_mu=200
max_r=1/2
max_gamma=5000
grid_size=64

#Set of (normalized) parameters
values_mu=[min_mu+k*(max_mu-min_mu)/(grid_size)+(max_mu-min_mu)/(2*grid_size) for k in range(grid_size)]
values_r=[k*(max_r)/(grid_size)+(max_r)/(2*grid_size) for k in range(grid_size)]
values_gamma=[k*(max_gamma)/(grid_size)+(max_gamma)/(2*grid_size) for k in range(grid_size)]                                            #Number of BS  

normalization_mu=values_mu[-1]-values_mu[0]
normalization_r=values_r[-1]-values_r[0]
normalization_gamma=values_gamma[-1]-values_gamma[0]

normalized_mu=[x/normalization_mu for x in values_mu]
normalized_gamma=[x/normalization_gamma for x in values_gamma]
normalized_r=[x/normalization_r for x in values_r]

def comparison_bootstrap(alpha, K, cost_map):
    #alpha is the proportion of points to be taken, K the number of bootstrap sets to be computed
    N=alpha*len(values_mu)*len(values_r)*len(values_gamma)
    #Compute set of top alpha performing triplets
    best_set=[]
    left_costs=cost_map[1]
    left_costs.sort()
    for k in range(int(N)):
        min_cost=left_costs[k]
        for m in range(grid_size):
            for l in range(grid_size):
                for n in range(grid_size):
                    if min_cost==cost_map[0][m][n][l]:
                        x=[normalized_mu[m],normalized_r[l],normalized_gamma[n]]
        best_set.append(x)
        

    #Compute the bootstrap sets
    bootstraps=[]
    possible_indices=list(itertools.product(range(grid_size),range(grid_size),range(grid_size)))
    for m in range(K):
        bootstraps.append([])
        selected_indices=random.sample(possible_indices,N)
        bootstraps[m].append([[normalized_mu[i[0]], normalized_gamma[i[1]], normalized_r[i[2]]] for i in selected_indices])


##################################################################################################################
#Compute the clustering statistic for top set
    barycenter_top_set=[sum([best_set[k][0] for k in range(len(best_set))])/len(best_set),sum([best_set[k][1] for k in range(len(best_set))])/len(best_set),sum([best_set[k][2] for k in range(len(best_set))])/len(best_set)]
    mean_distance_to_barycenter=sum([math.sqrt((best_set[k][0]-barycenter_top_set[0])**2+(best_set[k][1]-barycenter_top_set[1])**2+(best_set[k][2]-barycenter_top_set[2])**2) for k in range(len(best_set))])/(len(best_set))

#Compute clustering statistics for bootstrap sets
    mean_distances_to_barycenter_bootstrap=[]
    for m in range(K):
        barycenter=[sum([bootstraps[m][k][0] for k in range(len(bootstraps[m]))])/len(bootstraps[m]),sum([bootstraps[m][k][1] for k in range(len(bootstraps[m]))])/len(bootstraps[m]),sum([bootstraps[m][k][2] for k in range(len(bootstraps[m]))])/len(bootstraps[m])]
        mean_distances_to_barycenter=sum([math.sqrt((bootstraps[m][k][0]-barycenter[0])**2+(bootstraps[m][k][1]-barycenter[1])**2+(bootstraps[m][k][2]-barycenter[2])**2) for k in range(len(bootstraps[m]))])/(len(bootstraps[m]))
        mean_distances_to_barycenter_bootstrap.append(mean_distances_to_barycenter)
    mean_mean_distance_to_barycenter_bootstrap=sum(mean_distances_to_barycenter_bootstrap)/(len(mean_distances_to_barycenter_bootstrap))

#Difference between bootstrap average and top set
    Delta=(mean_mean_distance_to_barycenter_bootstrap-mean_distance_to_barycenter)/mean_mean_distance_to_barycenter_bootstrap

    return(Delta)

def get_deltas(parameters):
    K=parameters[0]
    results=parameters[1]
    deltas=[]
    for alpha in values_alpha:
        deltas.append(comparison_bootstrap(alpha, K, results))
    return deltas


#Import cost maps
cost_maps=np.load('cost_maps.npy', allow_pickle=True)

#One bootstrap comparison per TS, each performed by a separate thread
nb_processors=len(cost_maps)
if __name__ == "__main__":
    pool = Pool(nb_processors)
    Parameters=[]
    for k in range(nb_processors):
        Parameters.append([K,cost_maps[k]])
    final_results=pool.map(get_deltas, Parameters)
    np.save('Bootstrap comparison', final_results)
