# -*- 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.
"""
#Generates a dictionnary containing the model's outcome for:
#-Different network types: BA, ER, grid, WS, SBM
#-Different sizes: big (N=10000) mid (N=1000) and small (N=100)
#-Different values of the parameters gamma, mu and network parameters.

import os
import pickle
import numpy as np
from matplotlib.figure import Figure
from matplotlib.ticker import MaxNLocator
from DGSM_matrix_form_batched import model_non_SBM_batched, model_SBM_batched
os.chdir('/home/conjeaud')


N_big=10000
N_mid=1000
N_small=100
mu_b=-200
gamma_b=2000
m_b=2
r_b=0.1
T_big=50
T_small=75
M=50
sigma=1
beta=0.32
Lambda=10**(-2)
mus=[-500,-300,-200,-100,0]
gammas=[500,1000,2000,4000]
ms=[2,4,6,10]
ps=[0.1,0.2,0.3,0.5]
ks=[2,5,7,10]
rs=[0.01,0.05, 0.1,0.2]
k_b=5
sigmas=[1,1]
betas=[0.2,0.6]
clusters=[0.7,0.3]
D=0.8
p_b=0.2
results_dict={}
#NB: T_big and T_small are different to prevent instability. In large networks, taking a T too large 
#might cause the opinions to grow unbounded in some runs.



################################################################################################
#Barabasi Albert
network_type='BA'
#Small network
#Begin with mu:
results_mu_BA_small=[]
results_gamma_BA_small=[]
results_m_BA_small=[]
print('Generating '+network_type+' small')
for mu in mus:
    model=model_non_SBM_batched(N_small,network_type, [m_b],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_BA_small.append(model)
results_dict['mu_BA_small']=results_mu_BA_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_small,network_type, [m_b],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_BA_small.append(model)
results_dict['gamma_BA_small']=results_gamma_BA_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for m in ms:
    model=model_non_SBM_batched(N_small,network_type, [m],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_m_BA_small.append(model)
results_dict['m_BA_small']=results_m_BA_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
    
#Mid network
#Begin with mu:
results_mu_BA_mid=[]
results_gamma_BA_mid=[]
results_m_BA_mid=[]
print('Generating '+network_type+' mid')
for mu in mus:
    model=model_non_SBM_batched(N_mid,network_type, [m_b],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_BA_mid.append(model)
results_dict['mu_BA_mid']=results_mu_BA_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_mid,network_type, [m_b],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_BA_mid.append(model)
results_dict['gamma_BA_mid']=results_gamma_BA_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for m in ms:
    model=model_non_SBM_batched(N_mid,network_type, [m],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_m_BA_mid.append(model)
results_dict['m_BA_mid']=results_m_BA_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#Big network
#Begin with mu:
results_mu_BA_big=[]
results_gamma_BA_big=[]
results_m_BA_big=[]
print('Generating '+network_type+' big')
for mu in mus:
    model=model_non_SBM_batched(N_big,network_type, [m_b],mu,sigma,beta,gamma_b,Lambda, T_big, M)
    results_mu_BA_big.append(model)
results_dict['mu_BA_big']=results_mu_BA_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_big,network_type, [m_b],mu_b,sigma,beta,gamma,Lambda, T_big, M)
    results_gamma_BA_big.append(model)
results_dict['gamma_BA_big']=results_gamma_BA_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for m in ms:
    model=model_non_SBM_batched(N_big,network_type, [m],mu_b,sigma,beta,gamma_b,Lambda, T_big, M)
    results_m_BA_big.append(model)
results_dict['m_BA_big']=results_m_BA_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#########################################################################################################
#Erdos Renyi
network_type='ER'
#Small network
#Begin with mu:
results_mu_ER_small=[]
results_gamma_ER_small=[]
results_p_ER_small=[]
print('Generating '+network_type+' small')
for mu in mus:
    model=model_non_SBM_batched(N_small,network_type, [p_b],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_ER_small.append(model)
results_dict['mu_ER_small']=results_mu_ER_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_small,network_type, [p_b],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_ER_small.append(model)
results_dict['gamma_ER_small']=results_gamma_ER_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then p:
for p in ps:
    model=model_non_SBM_batched(N_small,network_type, [p],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_p_ER_small.append(model)
results_dict['p_ER_small']=results_gamma_ER_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#Mid network
#Begin with mu:
results_mu_ER_mid=[]
results_gamma_ER_mid=[]
results_m_ER_mid=[]
print('Generating '+network_type+' mid')
for mu in mus:
    model=model_non_SBM_batched(N_mid,network_type, [m_b],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_ER_mid.append(model)
results_dict['mu_ER_mid']=results_mu_ER_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_mid,network_type, [m_b],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_ER_mid.append(model)
results_dict['gamma_ER_mid']=results_gamma_ER_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for m in ms:
    model=model_non_SBM_batched(N_mid,network_type, [m],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_m_ER_mid.append(model)
results_dict['m_ER_mid']=results_m_ER_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)



#Big network
#Begin with mu:
results_mu_ER_big=[]
results_gamma_ER_big=[]
results_p_ER_big=[]
print('Generating '+network_type+' big')
for mu in mus:
    model=model_non_SBM_batched(N_big,network_type, [p_b],mu,sigma,beta,gamma_b,Lambda, T_big, M)
    results_mu_ER_big.append(model)
results_dict['mu_ER_big']=results_mu_ER_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_big,network_type, [p_b],mu_b,sigma,beta,gamma,Lambda, T_big, M)
    results_gamma_ER_big.append(model)
results_dict['gamma_ER_big']=results_gamma_ER_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then p:
for p in ps:
    model=model_non_SBM_batched(N_big,network_type, [p],mu_b,sigma,beta,gamma_b,Lambda, T_big, M)
    results_p_ER_big.append(model)
results_dict['p_ER_big']=results_p_ER_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#######################################################################################################
#Watts-Strogatz
network_type='WS'
#Small network
#Begin with mu:
results_mu_WS_small=[]
results_gamma_WS_small=[]
results_k_WS_small=[]
print('Generating '+network_type+' small')
for mu in mus:
    model=model_non_SBM_batched(N_small,network_type, [k_b,p_b],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_WS_small.append(model)
results_dict['mu_WS_small']=results_mu_WS_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_small,network_type, [k_b,p_b],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_WS_small.append(model)
results_dict['gamma_WS_small']=results_gamma_WS_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for k in ks:
    model=model_non_SBM_batched(N_small,network_type, [k,p_b],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_k_WS_small.append(model)
results_dict['k_WS_small']=results_k_WS_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#Mid network
#Begin with mu:
results_mu_WS_mid=[]
results_gamma_WS_mid=[]
results_m_WS_mid=[]
print('Generating '+network_type+' mid')
for mu in mus:
    model=model_non_SBM_batched(N_mid,network_type, [m_b],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_WS_mid.append(model)
results_dict['mu_WS_mid']=results_mu_WS_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_mid,network_type, [m_b],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_WS_mid.append(model)
results_dict['gamma_WS_mid']=results_gamma_WS_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for m in ms:
    model=model_non_SBM_batched(N_mid,network_type, [m],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_m_WS_mid.append(model)
results_dict['m_WS_mid']=results_m_WS_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)


#Big network
#Begin with mu:
results_mu_WS_big=[]
results_gamma_WS_big=[]
results_k_WS_big=[]
print('Generating '+network_type+' big')
for mu in mus:
    model=model_non_SBM_batched(N_big,network_type, [k_b,p_b],mu,sigma,beta,gamma_b,Lambda, T_big, M)
    results_mu_WS_big.append(model)
results_dict['mu_WS_big']=results_mu_WS_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    model=model_non_SBM_batched(N_big,network_type, [k_b,p_b],mu_b,sigma,beta,gamma,Lambda, T_big, M)
    results_gamma_WS_big.append(model)
results_dict['gamma_WS_big']=results_gamma_WS_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for k in ks:
    model=model_non_SBM_batched(N_big,network_type, [k,p_b],mu_b,sigma,beta,gamma_b,Lambda, T_big, M)
    results_k_WS_big.append(model)
results_dict['k_WS_big']=results_k_WS_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#######################################################################################################
#SBM networks
#Small networks
results_mu_SBM_small=[]
results_gamma_SBM_small=[]
results_r_SBM_small=[]
print('Generating SBM small')
for mu in mus:
    MUS=[mu,mu]
    model=model_SBM_batched(N_small,MUS,sigmas,betas,gamma_b,Lambda, T_small, M, clusters, D, r_b)
    results_mu_SBM_small.append(model)
results_dict['mu_SBM_small']=results_mu_SBM_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
for gamma in gammas:
    MUS=[mu_b,mu_b]
    model=model_SBM_batched(N_small,MUS,sigmas,betas,gamma,Lambda, T_small, M, clusters, D, r_b)
    results_gamma_SBM_small.append(model)
results_dict['gamma_SBM_small']=results_gamma_SBM_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
for r in rs:
    MUS=[mu_b,mu_b]
    model=model_SBM_batched(N_small,MUS,sigmas,betas,gamma_b,Lambda, T_small, M, clusters, D, r)
    results_r_SBM_small.append(model)
results_dict['r_SBM_small']=results_r_SBM_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#Mid network
#Begin with mu:
results_mu_SBM_mid=[]
results_gamma_SBM_mid=[]
results_m_SBM_mid=[]
print('Generating '+network_type+' mid')
for mu in mus:
    MUS=[mu,mu]
    model=model_SBM_batched(N_mid,MUS,sigmas,betas,gamma_b,Lambda, T_big, M, clusters, D, r)
    results_mu_SBM_mid.append(model)
results_dict['mu_SBM_mid']=results_mu_SBM_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    MUS=[mu_b,mu_b]
    model=model_SBM_batched(N_mid,MUS,sigmas,betas,gamma_b,Lambda, T_big, M, clusters, D, r)
    results_gamma_SBM_mid.append(model)
results_dict['gamma_SBM_mid']=results_gamma_SBM_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for m in ms:
    MUS=[mu_b,mu_b]
    model=model_SBM_batched(N_mid,MUS,sigmas,betas,gamma_b,Lambda, T_big, M, clusters, D, r)
    results_m_SBM_mid.append(model)
results_dict['m_SBM_mid']=results_m_SBM_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)


#Big networks
results_mu_SBM_big=[]
results_gamma_SBM_big=[]
results_r_SBM_big=[]
print('Generating SBM big')
for mu in mus:
    MUS=[mu,mu]
    model=model_SBM_batched(N_big,MUS,sigmas,betas,gamma_b,Lambda, T_big, M, clusters, D, r_b)
    results_mu_SBM_big.append(model)
results_dict['mu_SBM_big']=results_mu_SBM_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
for gamma in gammas:
    MUS=[mu_b,mu_b]
    model=model_SBM_batched(N_big,MUS,sigmas,betas,gamma,Lambda, T_big, M, clusters, D, r_b)
    results_gamma_SBM_big.append(model)
results_dict['gamma_SBM_big']=results_gamma_SBM_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
for r in rs:
    MUS=[mu_b,mu_b]
    model=model_SBM_batched(N_big,MUS,sigmas,betas,gamma_b,Lambda, T_big, M, clusters, D, r)
    results_r_SBM_big.append(model)
results_dict['r_SBM_big']=results_r_SBM_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
    
##############################################################################################
#Grid networks
network_type='grid'
grid_b_small=[3,3,3,4]
grid_b_mid=[6,6,6,5]
grid_b_big=[10,10,10,10]
grid_parameters_small=[[10,10],[5,5,4],[3,3,3,4], [2,2,3,3,3]]
grid_parameters_mid=[[32,32],[10,10,10],[6,6,6,5],[5,5,5,4,2]]
grid_parameters_big=[[100,100],[25,20,20],[10,10,10,10],[8,8,8,5,4]]
#Small network
#Begin with mu:
results_mu_grid_small=[]
results_gamma_grid_small=[]
results_dim_grid_small=[]
print('Generating '+network_type+' small')
for mu in mus:
    N=np.product(grid_b_small)
    model=model_non_SBM_batched(N,network_type, [grid_b_small],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_grid_small.append(model)
results_dict['mu_grid_small']=results_mu_grid_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    N=np.product(grid_b_small)
    model=model_non_SBM_batched(N,network_type, [grid_b_small],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_grid_small.append(model)
results_dict['gamma_grid_small']=results_gamma_grid_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for dim in grid_parameters_small:
    N=np.product(dim)
    model=model_non_SBM_batched(N,network_type, [dim],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_dim_grid_small.append(model)
results_dict['dim_grid_small']=results_dim_grid_small
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)

#Mid network
#Begin with mu:
results_mu_grid_mid=[]
results_gamma_grid_mid=[]
results_dim_grid_mid=[]
print('Generating '+network_type+' mid')
for mu in mus:
    N=np.product(grid_b_mid)
    model=model_non_SBM_batched(N,network_type, [grid_b_mid],mu,sigma,beta,gamma_b,Lambda, T_small, M)
    results_mu_grid_mid.append(model)
results_dict['mu_grid_mid']=results_mu_grid_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    N=np.product(grid_b_mid)
    model=model_non_SBM_batched(N,network_type, [grid_b_mid],mu_b,sigma,beta,gamma,Lambda, T_small, M)
    results_gamma_grid_mid.append(model)
results_dict['gamma_grid_mid']=results_gamma_grid_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for dim in grid_parameters_mid:
    N=np.product(dim)
    model=model_non_SBM_batched(N,network_type, [dim],mu_b,sigma,beta,gamma_b,Lambda, T_small, M)
    results_dim_grid_mid.append(model)
results_dict['dim_grid_mid']=results_dim_grid_mid
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)


#Big network
#Begin with mu:
results_mu_grid_big=[]
results_gamma_grid_big=[]
results_dim_grid_big=[]
print('Generating '+network_type+' big')
for mu in mus:
    N=np.product(grid_b_big)
    model=model_non_SBM_batched(N,network_type, [grid_b_big],mu,sigma,beta,gamma_b,Lambda, T_big, M)
    results_mu_grid_big.append(model)
results_dict['mu_grid_big']=results_mu_grid_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then gamma:
for gamma in gammas:
    N=np.product(grid_b_big)
    model=model_non_SBM_batched(N,network_type, [grid_b_big],mu_b,sigma,beta,gamma,Lambda, T_big, M)
    results_gamma_grid_big.append(model)
results_dict['gamma_grid_big']=results_gamma_grid_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
#Then m:
for dim in grid_parameters_big:
    N=np.product(dim)
    model=model_non_SBM_batched(N,network_type, [dim],mu_b,sigma,beta,gamma_b,Lambda, T_big, M)
    results_dim_grid_big.append(model)
results_dict['dim_grid_big']=results_dim_grid_big
with open('results_dict.pkl', 'wb') as fp:
    pickle.dump(results_dict, fp)
    
#########################################################################################################
#Generate and save figures
with open('results_dict.pkl', 'rb') as fp:
    results = pickle.load(fp)
for key in list(results.keys()):
    results_list=results[key]
    fg=Figure()
    ax=fg.gca()
    for k in range(len(results_list)):
        x=results_list[k]
        l,=ax.plot(x)
        if 'mu' in key:
            values=mus
            l.set_label(r"""$\mu$="""+str(values[k]))
        if 'gamma' in key:
            values=gammas
            l.set_label(r"""$\gamma$="""+str(values[k]))
        if 'r_' in key:
            values=rs
            l.set_label(r"""$r$="""+str(values[k]))
        if 'p' in key:
            values=ps
            l.set_label(r"""$p$="""+str(values[k]))
        if 'k' in key:
            values=ks
            l.set_label(r"""$k$="""+str(values[k]))
        if 'm_' in key:
            values=ms
            l.set_label(r"""$m$="""+str(values[k]))
        if 'dim' in key:
            values=[2,3,4,5]
            l.set_label(r"""$dim$="""+str(values[k]))
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.set_xlabel('Time', fontsize=15)
    ax.set_ylabel(r"""$\~S_t$""", rotation=0, fontsize=15, labelpad=10.0)
    ax.legend(loc='upper right')
    
    fg.savefig(key+'.svg')
