import torch
import pytest
from triangular_mask import triangular_mask, triangular_mask_slow


@pytest.mark.parametrize(
    "bsz,slen,ranges,pad_mask_idxs",
    [
        (2, 5, {1: [(1, 3)]}, [[0, 1], [0, 2], [1, 4]]),
        (2, 10, {1: [(1, 4), (5, 8)]}, [[1, 3], [1, 0]]),
        (2, 10, {1: [(1, 6)]}, [[1, 3]]),
    ],
)
def test_triangular_mask(bsz, slen, ranges, pad_mask_idxs):
    ranges_lst = [[] for i in range(bsz)]
    for k, v in ranges.items():
        ranges_lst[k] = v
    pad_mask = torch.ones((bsz, slen), dtype=torch.bool)
    for idxs in pad_mask_idxs:
        pad_mask[*idxs] = False
    args = (bsz, slen, ranges_lst, pad_mask, "cpu")
    assert torch.equal(triangular_mask_slow(*args), triangular_mask(*args))