# -*- 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
import math
import random
from multiprocessing import Pool
from DGSM_matrix_form import model_SBM_robust as model


#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())

#Optimization takes place over mu, gamma and r, it proceeds in two steps:
#-An initial grid search.
#-A refinement of the grid search using simulated annealing

#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=20        #Number of threads
grid_size=20
min_mu=-200
max_mu=200
max_r=1/2               #min_r is 0
max_gamma=5000          #min_gamma is 0

#Grid
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)]

#Hyperparameters for simulated annealing
K=600                   #Number of iterations
alpha=0.001             #Next point selection hyperparameter
Temp1=1                 #Initial temperature
eta1=0.95               #Rate of temperature decrease              

#Divide the grid for multi-threading
#Returns a list with one subgrid per thread
def divide_grid(nb_processors):
    if grid_size/nb_processors > int(grid_size/nb_processors ):
        print('No way to divide the space, try with a multiple of '+str(grid_size))
        return()
    else:
        number_of_values_mu=int(grid_size/nb_processors)
        set_of_values=[]
        dep_point=0
        for i in range(nb_processors):
            set_of_values.append([])
            list_mu=[values_mu[i] for i in range(dep_point, dep_point+number_of_values_mu)]
            set_of_values[i]=[list_mu, values_r, values_gamma]
            dep_point+=number_of_values_mu
        return(set_of_values)

#Compute the cost for an output of the model and a TS from the data
def cost(TS, data):
    T=len(TS)
    norm_model=sum([TS[i]**2 for i in range(T)])
    norm_data=sum([data[i]**2 for i in range(T)])
    scalar_product=sum([TS[i]*data[i] for i in range(T)])
    best_lambda=scalar_product/norm_data
    scaled_data=[best_lambda*data[i] for i in range(T)]
    L=math.sqrt(sum([(TS[i]-scaled_data[i])**2 for i in range(T)])/norm_model)
    return L


#One grid search will be run per thread, hence the input as a list
def grid_search(Parameters):
    subgrid=Parameters[0] 
    data=Parameters[1]      #TS from the data to fit the model to
    nb_mu=len(subgrid[0])
    nb_r=len(subgrid[1])
    nb_gamma=len(subgrid[2])
    counter=0
    best_L=0
    costs=np.zeros((nb_mu, nb_r, nb_gamma))
    for k in range(nb_mu):
        for n in range(nb_r):
            for l in range(nb_gamma):
                mus=[subgrid[0][k], subgrid[0][k]]
                r=subgrid[1][n]
                gamma=subgrid[2][l]
                model_outcome=model(N,mus,sigmas,betas,gamma,Lambda, len(data), M, clusters, D, r)[0]
                L=cost(model_outcome, data)
                if best_L<=L or counter==0:
                    best_L=L
                    best_param=[mus[0],r,gamma]
                    best_ts=model_outcome
                costs[k][n][l]=L
    return(best_L, best_param, best_ts, costs)

def simulated_annealing(parameters_set, data):
    K=parameters_set[0]                 #Number of iterations
    eta1=parameters_set[1]              #Temperature rate of decrease
    Temp1=parameters_set[2]             #Initial temperature
    alpha=parameters_set[3]             #Next point selection hyperparameter
    min_mu=parameters_set[4]            
    max_mu=parameters_set[5]
    max_r=parameters_set[6]
    max_gamma=parameters_set[7]
    M=parameters_set[8]                 #Number of experiments per run
    mu_0=parameters_set[9]              #Starting point mu
    r_0=parameters_set[10]              #Starting point r
    gamma_0=parameters_set[11]          #Starting point_gamma
    previous_L=parameters_set[12]       #Cost attained by grid search
    
    #Initialize parameters
    mu=mu_0
    mus=[mu,mu]
    r=r_0
    gamma=gamma_0
    Temp=Temp1*1/(eta1)
    
    for n in range(K):
        #Generate the model outcome
        model_outcome=model(N,mus,sigmas,betas,gamma,Lambda, len(data), M, clusters, D, r)
        TS=model_outcome[0]
        #Compute associated cost
        L=cost(TS, data)
        #Save initial parameters and cost as best at n=0
        if n == 0:
            best_L=previous_L
            best_parameters=[mu, r, gamma]
            L_mem=previous_L
            parameters_memory=[mu,r,gamma]
        else:
            #When a better performance is achieved 
            if L<L_mem:
            #Compute the new neighborhood to draw next parameters from
                range_mu=[max(min_mu, mu-((max_mu-min_mu)/2)*alpha**(1/3)), min(max_mu, mu+((max_mu-min_mu)/2)*alpha**(1/3))]
                range_r=[max(0, r-((max_r)/2)*alpha**(1/3)), min(max_r, r+((max_r)/2)*alpha**(1/3))]
                range_gamma=[max(0, gamma-((max_gamma)/2)*alpha**(1/3)), min(max_gamma, gamma+((max_gamma)/2)*alpha**(1/3))]
            #Keep parameters in memory if best so far
                if L<best_L:
                    best_L=L
                    best_parameters=[mu, r, gamma]
            #Keep L and parameters in memory for next time
                L_mem=L
                parameters_memory=[mu,r,gamma]
            #Draw new parameters
                mu=random.uniform(range_mu[0],range_mu[1])
                r=random.uniform(range_r[0], range_r[1])
                gamma=random.uniform(range_gamma[0], range_gamma[1])
            #When poorer performance is achieved
            if L>=L_mem:
                draw=random.random()
                #With some probability, accept the new parameters
                if draw<math.exp(-((L-L_mem)/Temp)):
                    #Compute the neighborhood
                    range_mu=[max(min_mu, mu-((max_mu-min_mu)/2)*alpha**(1/3)), min(max_mu, mu+((max_mu-min_mu)/2)*alpha**(1/3))]
                    range_r=[max(0, r-((max_r)/2)*alpha**(1/3)), min(max_r, r+((max_r)/2)*alpha**(1/3))]
                    range_gamma=[max(0, gamma-((max_gamma)/2)*alpha**(1/3)), min(max_gamma, gamma+((max_gamma)/2)*alpha**(1/3))]
                    #Draw new parameters
                    mu=random.uniform(range_mu[0],range_mu[1])
                    r=random.uniform(range_r[0], range_r[1])
                    gamma=random.uniform(range_gamma[0], range_gamma[1])
                    #Keep L in memory for next time
                    L_mem=L
                #If not accepted, take previous parameters' neighborhood
                else:
                    mu=parameters_memory[0]
                    r=parameters_memory[1]
                    gamma=parameters_memory[2]
                    range_mu=[max(min_mu, mu-((max_mu-min_mu)/2)*alpha**(1/3)), min(max_mu, mu+((max_mu-min_mu)/2)*alpha**(1/3))]
                    range_r=[max(0, r-((max_r)/2)*alpha**(1/3)), min(max_r, r+((max_r)/2)*alpha**(1/3))]
                    range_gamma=[max(0, gamma-((max_gamma)/2)*alpha**(1/3)), min(max_gamma, gamma+((max_gamma)/2)*alpha**(1/3))]
                    #Draw new parameters
                    mu=random.uniform(range_mu[0],range_mu[1])
                    r=random.uniform(range_r[0], range_r[1])
                    gamma=random.uniform(range_gamma[0], range_gamma[1])
            #Once the step is completed decrease temperature
            Temp=Temp*eta1
    return(best_parameters, best_L)

#Launch the optimization process
subgrids=divide_grid(nb_processors)
data_set=np.load('time_series.npy', allow_pickle=True)
for X in data_set:
    data=X[0]
    title=X[1]
    #Launch grid search on each subgrid on different threads
    parameters=[[subgrids[i], data] for i in range(nb_processors)]
    pool = Pool(nb_processors)
    results=pool.map(grid_search, parameters)
    pool.close()
    pool.join()
    #Identify the best performing triplet among the ones found by each thread
    costs=[R[0] for R in results]
    best_triplet=results[costs.index(min(costs))]
    mu_0=best_triplet[0]
    r_0=best_triplet[1]
    gamma_0=best_triplet[2]
    #Launch the simulated annealing starting at the best triplet
    parameters_SA=[K,eta1,Temp1,alpha,min_mu,max_mu,max_r,max_gamma,M, mu_0,r_0,gamma_0, min(costs)]
    SA_results=simulated_annealing(parameters_SA, data)
    #Save the results
    np.save(title, SA_results)