import numpy as np


def scatter_array(comm, data, master=0):
    if comm.Get_size() == 1:
        data = data
    else:
        if comm.Get_rank() == master:
            data = np.array_split(data, comm.Get_size())
        else:
            data = None
        data = comm.scatter(data, root=master)

    return data


def bcast_array(comm, data, master=0):
    if comm.Get_size() == 1:
        data = data
    else:
        data = comm.bcast(data, root=master)

    return data


def balance_dataframe(comm, data, master=0):
    import pandas as pd
    data = comm.gather(data, root=master)
    if comm.Get_rank() == master:
        data = pd.concat(data)
        data = np.array_split(data, comm.Get_size())
    else:
        data = None

    data = comm.scatter(data, root=master)

    return data


def get_balanced_distribution(processors, shape):
    while len(shape) < 4:
        shape = (0,) + shape

    fid_dist = {}
    total_rows = shape[2]

    procs_rows = total_rows // processors
    procs_rows_extended = total_rows-(procs_rows*processors)

    rows_sum = 0
    for proc in range(processors):
        if proc < procs_rows_extended:
            aux_rows = procs_rows + 1
        else:
            aux_rows = procs_rows

        total_rows -= aux_rows
        if total_rows < 0:
            rows = total_rows + aux_rows
        else:
            rows = aux_rows

        min_fid = proc * aux_rows * shape[3]
        max_fid = (proc + 1) * aux_rows * shape[3]

        fid_dist[proc] = {
            'y_min': rows_sum,
            'y_max': rows_sum + rows,
            'x_min': 0,
            'x_max': shape[3],
            'fid_min': min_fid,
            'fid_max': max_fid,
            'shape': (shape[0], shape[1], rows, shape[3]),
        }

        rows_sum += rows

    return fid_dist
