# -*- coding: utf-8 -*-
"""List-handling utilities."""

from __future__ import division, print_function, absolute_import


def load_balance_list(L, n):
    """Given a list of arbitrary items, split it to n roughly equal-sized parts.

    This is useful for dividing a list of work items in MPI parallelization.
    It is assumed that each work item takes the same amount of time; hence the
    initial distribution is generated by naive integer division.

    If len(L) does not divide evenly with n, the remaining items are distributed
    on an item-by-item basis to the first (len(L) mod n) parts.

    If n > len(L), the items will be distributed on an item-by-item basis to the
    first len(L) parts, and the rest of the parts will get an empty list.
    
    Examples:
        load_balance_list(range(14), 2)
            =>  [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12, 13]]
        In this case, the division is even (no remainder).

        load_balance_list(range(15), 2)
            =>  [[0, 1, 2, 3, 4, 5, 6, 14], [7, 8, 9, 10, 11, 12, 13]]
        In this case, the item "14" is left over after the integer division.
        The leftover item is placed in the first part.

        load_balance_list(range(4), 8)
            =>  [[0], [1], [2], [3], [], [], [], []]
        In this case, n is so large that there are not enough items to place
        even one item in each part. The empty list is generated for those parts
        for which no item is available.

    Parameters:
        L = any Python list.

    Return value:
        List of lists: L_out = [L1, L2, L3, ..., Ln]
        where L1, L2, ... are sublists of L.

        It always holds that len(L_out) == n. Note that the lengths of the
        Lj (individual lists inside L_out) may differ by one item, depending on
        whether the division was even.

        If n == 1, the return value is [L] for compatibility
        of the output format with the n > 1 case.

    """
    if n < 1:
        raise ValueError("n must be >= 1; got %d" % (n))

    # If we're splitting L to one part, this is a no-op.
    # But wrap it - the caller is expecting a list of lists.
    #
    if n == 1:
        return [L]

    nitems = len(L)  # number of items to be distributed

    out = []
    if n <= nitems:
        # Items per process.
        blocklen  = nitems // n  # integer division!

        # Leftover items.
        remainder = nitems - blocklen*n  # this is always < n

        # Distribute the items that divided evenly.
        for m in range(n):
            offs = m*blocklen
            out.append( L[offs:offs+blocklen] )

        # Distribute the leftovers (if any).
        if remainder > 0:
            offs = nitems - remainder
            for m in range(remainder):
                out[m].append( L[offs] )
                offs += 1
    else:
        # Distribute one item to each part as long as items are available.
        for m in range(nitems):
            out.append( [L[m]] )  # wrap the item to make a single-item list.

        # Give an empty list to the rest of the parts.
        nempties = n - nitems  # this many empty lists are needed
        for m in range(nempties):
            out.append( [] )

    assert( len(out) == n )
    return out

