import torch

from inteprelens.utils.tensor_dict import TensorDict


def test_tensor_dict_basic_access_and_map():
    td = TensorDict(a=torch.ones(2), b=torch.zeros(2))

    assert td.device == td.a.device
    assert torch.equal(td["a"], td.a)

    td["c"] = [1, 2]
    assert isinstance(td.c, torch.Tensor)
    assert td.c.device == td.device

    td2 = td.map(lambda t: t + 1)
    assert torch.equal(td2.a, td.a + 1)
    assert torch.equal(td2.b, td.b + 1)


def test_tensor_dict_cat_and_stack_with_padding():
    """
    dim=0 operates along the first axis(axis index 0).
    For a 2D tensor shaped(rows, cols):
    * dim=0 rows
    * dim=1 cols

    For 3D tensor shaped(batch, seq, hidden):
    * dim=0 batch
    * dim=1 seq
    * dim=2 hidden

    a = torch.tensor([[1, 2, 3]])   # shape (1, 3)
    b = torch.tensor([[4, 5, 6]])   # shape (1, 3)

    torch.cat([a, b], dim=0)  # shape (2, 3)
    # [[1,2,3],
    #  [4,5,6]]

    torch.cat([a, b], dim=1)  # shape (1, 6)
    # [[1,2,3,4,5,6]]

    """
    td1 = TensorDict(x=torch.tensor([[1, 2]]))
    td2 = TensorDict(x=torch.tensor([[3, 4, 5]]))

    td_cat = TensorDict.cat([td1, td2], dim=0, pad_value=0)
    assert td_cat.x.shape == (2, 3)
    assert torch.equal(td_cat.x[0], torch.tensor([1, 2, 0]))
    assert torch.equal(td_cat.x[1], torch.tensor([3, 4, 5]))

    td_stack = TensorDict.stack([td1, td2], dim=0, pad_value=0)
    assert td_stack.x.shape == (2, 1, 3)
    assert torch.equal(td_stack.x[0, 0], torch.tensor([1, 2, 0]))
    assert torch.equal(td_stack.x[1, 0], torch.tensor([3, 4, 5]))
