import warnings


def chunk_baselines_by_redundant_groups(reds, max_chunk_size):
    """Chunk list of baselines by redundant group constrained by number of baselines.

    This method calculates chunks of baselines to load simultaneously constrained by the conditions that
    all baselines simultaneously loaded completely fill out redundant groups and the number of baselines
    loaded is less then or equal to max_chunk_size.
    If a redundant group exceeds max_chunk_size it will be loaded in its entirety anyways
    and warning will be raised.

    Parameters
    ----------
    reds : list
        list of lists with redundant groups of antpairpols
    max_chunk_size :
        maximum number of baselines to group in each chunk
                    while filling out redundant groups.

    Returns
    -------

    chunked_baselines : list
        list of lists of antpairpols where each chunk containes entire redundant groups
        and (if possible) includes a number of baselines less then max_chunk_size.
    """
    baseline_chunks = []
    group_index = 0
    # iterate through redundant groups
    for grp in reds:
        if len(grp) > max_chunk_size:
            # if red group is larger then the chunk size.
            # then give a warning and treate the red group as a chunk anyways.
            warnings.warn("Warning: baseline group of length %d encountered with number"
                          " of baselines exceeding max_chunk_size=%d."
                          " First baseline is %s"
                          " Loading group anyways." % (len(reds[group_index]), max_chunk_size, str(reds[group_index][0])))
            baseline_chunks.append(grp)
        else:
            # if baseline_chunks is empty, initialize the first chunk.
            if len(baseline_chunks) == 0:
                baseline_chunks.append([])
            if len(baseline_chunks[-1]) + len(grp) <= max_chunk_size:
                baseline_chunks[-1].extend(grp)
            else:
                baseline_chunks.append(grp)
    return baseline_chunks