import numpy as np


def split_selection_distributed(region_data, params):
    """
    Splits selected particles between cores. Equal number of distributed files
    per core

    Parameters
    ----------
    region_data : dict
        Global particles indexes to load generated by select_region()
    params : _SwiftSnapshotParams object
        pyread_swift parameters

    Returns
    -------
    index_data : dict
        Local particle indexes for this rank to load
    """

    index_data = {}
    index_data["lefts"] = []
    index_data["rights"] = []
    index_data["files"] = []
    index_data["num_to_load"] = []
    index_data["total_num_to_load"] = 0

    # Loop over each file and assign it to a rank.
    for i, fileno in enumerate(np.unique(region_data["files"])):
        if i % params.comm_size != params.comm_rank:
            continue
        mask = np.where(region_data["files"] == fileno)

        index_data["files"].append(fileno)
        index_data["lefts"].append(region_data["lefts"][mask])
        index_data["rights"].append(region_data["rights"][mask])
        index_data["num_to_load"].append(region_data["num_to_load"][mask])
        index_data["total_num_to_load"] += np.sum(index_data["num_to_load"][-1])

    return index_data


def split_selection_collective(region_data, params):
    """
    Splits selected particles between cores. Equal number of particles per
    core.

    Each file is read collectivley by all ranks, dividing the top level cells
    between the ranks.

    If this will load less than 100 particles per rank on average, it reverts
    to 1 core doing all the reading.

    Parameters
    ----------
    region_data : dict
        Global particles indexes to load generated by select_region()
    params : _SwiftSnapshotParams object
        pyread_swift parameters

    Returns
    -------
    index_data : dict
        Local particle indexes for this rank to load
    """

    index_data = {}
    index_data["lefts"] = []
    index_data["rights"] = []
    index_data["files"] = []
    index_data["num_to_load"] = []
    index_data["total_num_to_load"] = 0

    if (
        params.comm_size > 1
        and region_data["total_num_to_load"] > params.min_in_tot_collective
    ):
        # Loop over each file.
        for fileno in np.unique(region_data["files"]):
            file_mask = np.where(region_data["files"] == fileno)

            # How many particles from this file will each core load?
            num_to_load_from_this_file = np.sum(region_data["num_to_load"][file_mask])
            num_per_core = np.zeros(params.comm_size, dtype="i8")
            num_per_core[:] = num_to_load_from_this_file // params.comm_size
            num_per_core[-1] += num_to_load_from_this_file % params.comm_size

            # The case where this file has few particles compared to the number of ranks.
            # In this case one rank will read all the particles from this file.
            if num_per_core[0] < params.min_in_file_collective:
                if params.comm_rank == params.comm_size - 1:
                    params.message(
                        "Will load %i particles file %i"
                        % (num_to_load_from_this_file, fileno)
                    )

                    index_data["lefts"].append(region_data["lefts"][file_mask])
                    index_data["rights"].append(region_data["rights"][file_mask])
                    index_data["files"].append(fileno)
                    index_data["num_to_load"].append(
                        region_data["num_to_load"][file_mask]
                    )
                    index_data["total_num_to_load"] += num_to_load_from_this_file

            # The case where this file has many particles per rank.
            # In this case we do a collective read.
            else:
                # What particles will each core load?
                tmp_my_lefts = []
                tmp_my_rights = []
                tmp_my_files = []
                my_count = np.zeros(
                    params.comm_size, dtype="i8"
                )  # How many ps have I loaded?

                for l, r, chunk_no in zip(
                    region_data["lefts"][file_mask],
                    region_data["rights"][file_mask],
                    region_data["num_to_load"][file_mask],
                ):

                    chunk_bucket = chunk_no
                    tmp_offset = 0

                    # Allocate this chunk over the cores.
                    while chunk_bucket > 0:
                        for j in range(params.comm_size):
                            # No room left on this core.
                            if my_count[j] == num_per_core[j]:
                                continue

                            if my_count[j] + chunk_bucket <= num_per_core[j]:
                                # Can all fit on this core.
                                my_count[j] += chunk_bucket
                                if params.comm_rank == j:
                                    tmp_my_lefts.append(l + tmp_offset)
                                    tmp_my_rights.append(l + tmp_offset + chunk_bucket)
                                    tmp_my_files.append(fileno)
                                chunk_bucket = 0
                            else:
                                # Only a bit can fit on this core.
                                diff = num_per_core[j] - my_count[j]
                                my_count[j] += diff
                                chunk_bucket -= diff
                                if params.comm_rank == j:
                                    tmp_my_lefts.append(l + tmp_offset)
                                    tmp_my_rights.append(l + tmp_offset + diff)
                                    tmp_my_files.append(fileno)
                                tmp_offset += diff

                            # All allocated.
                            if chunk_bucket == 0:
                                break
                assert np.sum(my_count) == num_to_load_from_this_file

                # Make sure we got them all.
                chk_sum = params.comm.allreduce(
                    np.sum(np.array(tmp_my_rights) - np.array(tmp_my_lefts))
                )
                assert (
                    chk_sum == num_to_load_from_this_file
                ), f"Did not divide ps correctly {chk_sum} ne {num_to_load_from_this_file}"

                params.message(
                    "Will load %i particles l=%s r=%s from file %i"
                    % (
                        num_per_core[params.comm_rank],
                        tmp_my_lefts,
                        tmp_my_rights,
                        fileno,
                    )
                )

                index_data["lefts"].append(tmp_my_lefts)
                index_data["rights"].append(tmp_my_rights)
                index_data["files"].append(fileno)
                index_data["num_to_load"].append(
                    np.array(index_data["rights"][-1])
                    - np.array(index_data["lefts"][-1])
                )
                index_data["total_num_to_load"] += np.sum(index_data["num_to_load"][-1])
    else:
        # Single core case.
        if params.comm_rank == 0:

            # Loop over each file part.
            for fileno in np.unique(region_data["files"]):
                mask = np.where(region_data["files"] == fileno)

                index_data["files"].append(fileno)
                index_data["lefts"].append(region_data["lefts"][mask])
                index_data["rights"].append(region_data["rights"][mask])
                index_data["num_to_load"].append(region_data["num_to_load"][mask])
                index_data["total_num_to_load"] += np.sum(index_data["num_to_load"][-1])

                params.message(
                    f"Will load {index_data['num_to_load'][-1][0]} particles from file {fileno}"
                )
    return index_data
