import numpy as np


def distribute_jobs_per_processor(flag_mpi,N_spectra):
    if flag_mpi.upper() == "Y":
        mpi_dict = mpi_variable()
        rank     = mpi_dict["rank"]
        N_proc   = mpi_dict["N_proc"]
    else:
        rank     = 0
        N_proc   = 1
        comm     = None
        mpi_dict = {"rank":rank,"N_proc":N_proc,"comm":comm}

    mpi_job_dict = distribute_common_particle_per_file(N_spectra,N_proc)
    lb_arr       = mpi_job_dict["lb_arr"]
    ub_arr       = mpi_job_dict["ub_arr"]
    indx_lb      = lb_arr[rank]
    indx_ub      = ub_arr[rank]
    N_los_total  = indx_ub - indx_lb
    mpi_dict.update({"indx_lb":indx_lb,"indx_ub":indx_ub,\
                     "N_los_total":N_los_total})

    return mpi_dict

def mpi_variable():
    from mpi4py import MPI
    comm      = MPI.COMM_WORLD
    rank      = comm.Get_rank()
    N_proc    = comm.Get_size()
    status    = MPI.Status()
    proc_name = MPI.Get_processor_name()
    mpi_dict  = {"comm":comm,"rank":rank,"N_proc":N_proc,"status":status,"MPI":MPI,"proc_name":proc_name}
    return mpi_dict

def distribute_common_particle_per_file(N_job,N_proc):
    """
    Let N_job = 1525,
    N_proc = 63
    N_particle_per_file = 25
    part_lb_ub_arr ==>> 25, 25, 25, 25, ...., 25, 25 (63 times)
    N_particle_extra = 25 x 63 - 1525 = 1575 - 1525 = 50
    After the for loop output will be
    part_lb_ub_arr ==>> 24, 24, 24, 24, ...., 25, 25 (24 -->> 50 times, 25 -->> 63-50 = 13 times)


    lb_arr ==>> [0  24 48 ... 1475 1500]
    ub_arr ==>> [24 48 72 ... 1500 1525]

    """
    N_particle_per_file = int(N_job / N_proc) + 1
    part_lb_ub_arr      = np.ones(N_proc,dtype=np.int64)*N_particle_per_file
    N_particle_extra    = int(N_particle_per_file * N_proc - N_job)
    for indx in range(N_particle_extra):
        part_lb_ub_arr[indx] -= 1
    ub_arr      = np.cumsum(part_lb_ub_arr)
    lb_arr      = np.zeros(ub_arr.shape[0],dtype=np.int64)
    lb_arr[1:]  = ub_arr[:-1]
    mpi_job_dict = {"lb_arr":lb_arr,"ub_arr":ub_arr}
    return mpi_job_dict

def mpi_check_number_of_processors(N_Grid,N_proc):

    if (N_Grid % N_proc) != 0:
        print(N_Grid,"should be multiple of N_proc =",N_proc)
        print("Exiting the code ...")
        exit()
    else:
        N_Grid_per_proc = int(N_Grid / N_proc)

    return N_Grid_per_proc



def distribute_common_particle_per_file_test():
    out_dict = distribute_common_particle_per_file(64,1)
    print(out_dict["lb_arr"])
    print(out_dict["ub_arr"])

#distribute_common_particle_per_file_test()
















