import math
import torch


def sliding_window_subbatch(toks, window_size, stride):
    """A sliding-window sub-batching function.

    :param toks:            batched (and encoded) input tokens
    :param window_size:     a sliding window size
    :param stride:          a sliding window stride

    :return:  a tuple (stack sub-batched tokens, the number of sub-batches)
    """
    _, dlen = toks.shape[:2]
    assert dlen > 0
    # Ceiling of the negative number is a negative number too!!!
    # Hence, max(0, ...
    subbatch_qty = math.ceil(max(0, dlen-window_size)/stride) + 1
    assert subbatch_qty > 0, f'Bad sub-batch: {subbatch_qty} dlen {dlen} window_size {window_size} stride {stride}'
    stack = []
    if subbatch_qty == 1:
       return toks, subbatch_qty
    else:
        for s in range(subbatch_qty):
            if s*stride+window_size < dlen:
                stack.append(toks[:, s*stride: s*stride+window_size])
            else:
                nulls = torch.zeros_like(toks[:, :s*stride+window_size - dlen])
                stack.append(torch.cat([toks[:, s*stride:], nulls], dim=1))

        return torch.cat(stack, dim=0), subbatch_qty