# -*- 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 os
import numpy as np
from multiprocessing import Pool
from Optimization import divide_grid, grid_search

#os.chdir('folder_name')

#Create the TS, comment out if data_sets already in directory
with open("Get_TS.py") as file:
    exec(file.read())

#Model parameters (fixed)
N=100                   #Network size
Lambda=1/100            #Transition to state 1 parameter
clusters=[0.7,0.3]      #Relative size of clusters
D=0.9                   #Intra cluster connectivity
sigmas=[1,1]            #Variance of information shock
betas=[0.3, 0.7]        #Proportion of beta=1 in each cluster
M=50                    #Number of experiments per run

#Optimization parameters
nb_processors=32        #Number of threads
grid_size=64
min_mu=-200
max_mu=200
max_r=1/2               #min_r is 0
max_gamma=5000          #min_gamma is 0

#Launch the optimization process
subgrids=divide_grid(nb_processors)
data_set=np.load('time_series.npy', allow_pickle=True)
cost_maps=[]
for X in data_set:
    data=X[0]
    title=X[1]
    if __name__ == "__main__":
        pool = Pool(nb_processors)
        results=pool.map(grid_search, subgrids)
        pool.close()
        pool.join()
        #Concatenate the outcome from each thread
        cost_map=results[0][-1]
        for k in range(1, len(results)):
            cost_map=np.concatenate((cost_map, results[k]), axis=0)
        cost_maps.append([cost_map, title])
        np.save('Cost maps', cost_maps)
        