import copy


def truncate_list_of_lists(sents_tokenized, num_max, truncate_bottom=True):
    """Return a truncated list, with summ of tokens not exceeding maximum.
    Truncate by lists. If single left list is still too long, truncate it by tokens.
    In our context each element of sents_tokenized is a sentence represented as a list of tokens.

    Args:
        sents_tokenized (List[List[str]]): List, each element is a list.
        num_max (int): maximal allowed number of tokens.
        truncate_bottom (bool): truncate starting from bottom lists.

    Returns:
        sents_tokenized (List[str]): truncated list
    """
    sents_truncated = []
    if truncate_bottom:
        len_truncated = 0
        # Cut by sentences:
        for sent in sents_tokenized:
            len_truncated_maybe = len_truncated + len(sent)
            if len_truncated_maybe > num_max:
                break
            len_truncated = len_truncated_maybe
            sents_truncated.append(sent)
            if len_truncated == num_max:
                break
    else:
        sents_truncated = copy.deepcopy(sents_tokenized)
        len_truncated = sum([len(s) for s in sents_tokenized])
        # Cut by sentences:
        for sent in sents_tokenized:
            if len_truncated <= num_max:
                break
            sents_truncated = sents_truncated[1:]
            len_truncated = len_truncated - len(sent)
    if not sents_truncated:
        sent_use = sents_tokenized[0] if truncate_bottom else sents_tokenized[-1]
        sents_truncated = [copy.deepcopy(sent_use)]
        len_truncated = len(sents_truncated[0])
        # Cut by tokens - always from the top:
        if len_truncated > num_max:
            len_remove = len_truncated - num_max
            sents_truncated[0] = sents_truncated[0][len_remove:]
    assert sum([len(s) for s in sents_truncated]) <= num_max
    return sents_truncated