import numpy as np
from mpi4py import MPI
#from tools.realization import realization
from tools.generation.config import DotheLoop, get_config
import os
import sys
from tools.Prealization import realization
from utilities.conditional_decorator import *
from memory_profiler import profile
import multiprocessing as mp

CONFIG_FILE_PATH = 'config.ini' if 'CONFIG_FILE_PATH' not in os.environ else os.environ['CONFIG_FILE_PATH']
IS_TEST = False if 'TEST' not in os.environ else True

def main():
	comm = MPI.COMM_WORLD
	rank = comm.Get_rank()
	pn = comm.Get_size()

	if pn==1:
		sequential()
		return
	if rank==0:
		manager()
	else:
		worker()
	return

@conditional_decorator(profile, IS_TEST)
def sequential():
	comm = MPI.COMM_WORLD
	conffile = CONFIG_FILE_PATH
	parser,iterables = get_config(conffile)
	njobs = DotheLoop(-1,parser,iterables)
	start_job=0
	for job in range(start_job,njobs):
		p = mp.Process(target=realization, args=(job,))
		p.start()
		p.join()
		#realization(job)
	return

def manager():
	comm = MPI.COMM_WORLD
	conffile = CONFIG_FILE_PATH
	parser,iterables = get_config(conffile)
	njobs = DotheLoop(-1,parser,iterables)
	start_job=0
	for job in range(start_job,njobs):
		dest=comm.recv(source=MPI.ANY_SOURCE)
		comm.send(job,dest=dest)
	for i in range(comm.Get_size()-1):
		dest=comm.recv(source=MPI.ANY_SOURCE)
		comm.send(-1,dest=dest)

	return

@conditional_decorator(profile, IS_TEST)
def worker():
	comm = MPI.COMM_WORLD
	rank = comm.Get_rank()
	job=1
	while job!=-1:
		comm.send(rank,dest=0)
		job =  comm.recv(source=0)
		p = mp.Process(target=realization, args=(job,))
		p.start()
		p.join()
		#realization(job)
	return


main()