import torch

from inteprelens.utils.ds import MaskedSequenceDataset
from inteprelens.utils.tensor_dict import TensorDict


def test_masked_sequence_dataset_collate_truncates_right_padding():
    pad_id = 0
    dataset = MaskedSequenceDataset(
        pad_token_id=pad_id,
        input_ids=torch.tensor([[1, 2, 0], [3, 4, 0]]),
        loss_masks=torch.tensor([[True, True, False], [True, True, False]]),
    )

    batch = [
        TensorDict(
            input_ids=dataset["input_ids"][0],
            loss_masks=dataset["loss_masks"][0],
        ),
        TensorDict(
            input_ids=dataset["input_ids"][1],
            loss_masks=dataset["loss_masks"][1],
        ),
    ]

    batch_td = dataset.collate_fn(batch)

    assert batch_td["input_ids"].shape == (2, 2)
    assert torch.equal(
        batch_td["input_ids"], torch.tensor([[1, 2], [3, 4]])
    )
    assert torch.equal(
        batch_td["loss_masks"], torch.tensor([[True, True], [True, True]])
    )


def test_masked_sequence_dataset_collate_truncates_left_padding():
    pad_id = 0
    dataset = MaskedSequenceDataset(
        pad_token_id=pad_id,
        pad_right=False,
        input_ids=torch.tensor([[0, 1, 2], [0, 3, 4]]),
        loss_masks=torch.tensor([[False, True, True], [False, True, True]]),
    )

    batch = [
        TensorDict(
            input_ids=dataset["input_ids"][0],
            loss_masks=dataset["loss_masks"][0],
        ),
        TensorDict(
            input_ids=dataset["input_ids"][1],
            loss_masks=dataset["loss_masks"][1],
        ),
    ]

    batch_td = dataset.collate_fn(batch)

    assert batch_td["input_ids"].shape == (2, 2)
    assert torch.equal(
        batch_td["input_ids"], torch.tensor([[1, 2], [3, 4]])
    )
    assert torch.equal(
        batch_td["loss_masks"], torch.tensor([[True, True], [True, True]])
    )
