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

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):
        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)
        realization(job)
    return


main()