# -*- 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 os
from DGSM_matrix_form import model_non_SBM, model_SBM
from multiprocessing import Pool
os.chdir('/home/conjeaud')

#Have each of the M trials of the model generated by a different processor

def model_non_SBM_prime(Parameters):
    N,network_type, network_parameters,mu,sigma,beta,gamma,Lambda, T, M=Parameters[0],Parameters[1],Parameters[2],Parameters[3],Parameters[4],Parameters[5],Parameters[6],Parameters[7],Parameters[8],Parameters[9]
    return np.array(model_non_SBM(N,network_type, network_parameters,mu,sigma,beta,gamma,Lambda, T, M)[0])

def model_SBM_prime(Parameters):
   N,mus,sigmas,betas,gamma,Lambda, T, M, clusters, D, r=Parameters[0],Parameters[1],Parameters[2],Parameters[3],Parameters[4],Parameters[5],Parameters[6],Parameters[7],Parameters[8],Parameters[9], Parameters[10]
   return np.array(model_SBM(N,mus,sigmas,betas,gamma,Lambda, T, M, clusters, D, r)[0])

def model_non_SBM_batched(N,network_type, network_parameters,mu,sigma,beta,gamma,Lambda, T, M):
    if __name__ == "__main__":
        Parameters=[[N,network_type, network_parameters,mu,sigma,beta,gamma,Lambda, T, 1] for _ in range(M)]
        pool = Pool(M)
        results=pool.map(model_non_SBM_prime, Parameters)
        pool.close()
        pool.join()
    return sum(results)/len(results)

def model_SBM_batched(N,mus,sigmas,betas,gamma,Lambda, T, M, clusters, D, r):
    if __name__ == "__main__":
        Parameters=[[N,mus,sigmas,betas,gamma,Lambda, T, 1, clusters, D, r] for _ in range(M)]
        pool = Pool(M)
        results=pool.map(model_SBM_prime, Parameters)
        pool.close()
        pool.join()
    return sum(results)/len(results)
    