# import torch
# import pytest
# from pina.data.dataset import PinaDatasetFactory, PinaTensorDataset

# input_tensor = torch.rand((100, 10))
# output_tensor = torch.rand((100, 2))

# input_tensor_2 = torch.rand((50, 10))
# output_tensor_2 = torch.rand((50, 2))

# conditions_dict_single = {
#     "data": {
#         "input": input_tensor,
#         "target": output_tensor,
#     }
# }

# conditions_dict_single_multi = {
#     "data_1": {
#         "input": input_tensor,
#         "target": output_tensor,
#     },
#     "data_2": {
#         "input": input_tensor_2,
#         "target": output_tensor_2,
#     },
# }

# max_conditions_lengths_single = {"data": 100}

# max_conditions_lengths_multi = {"data_1": 100, "data_2": 50}


# @pytest.mark.parametrize(
#     "conditions_dict, max_conditions_lengths",
#     [
#         (conditions_dict_single, max_conditions_lengths_single),
#         (conditions_dict_single_multi, max_conditions_lengths_multi),
#     ],
# )
# def test_constructor_tensor(conditions_dict, max_conditions_lengths):
#     dataset = PinaDatasetFactory(
#         conditions_dict,
#         max_conditions_lengths=max_conditions_lengths,
#         automatic_batching=True,
#     )
#     assert isinstance(dataset, PinaTensorDataset)


# def test_getitem_single():
#     dataset = PinaDatasetFactory(
#         conditions_dict_single,
#         max_conditions_lengths=max_conditions_lengths_single,
#         automatic_batching=False,
#     )

#     tensors = dataset.fetch_from_idx_list([i for i in range(70)])
#     assert isinstance(tensors, dict)
#     assert list(tensors.keys()) == ["data"]
#     assert sorted(list(tensors["data"].keys())) == ["input", "target"]
#     assert isinstance(tensors["data"]["input"], torch.Tensor)
#     assert tensors["data"]["input"].shape == torch.Size((70, 10))
#     assert isinstance(tensors["data"]["target"], torch.Tensor)
#     assert tensors["data"]["target"].shape == torch.Size((70, 2))


# def test_getitem_multi():
#     dataset = PinaDatasetFactory(
#         conditions_dict_single_multi,
#         max_conditions_lengths=max_conditions_lengths_multi,
#         automatic_batching=False,
#     )
#     tensors = dataset.fetch_from_idx_list([i for i in range(70)])
#     assert isinstance(tensors, dict)
#     assert list(tensors.keys()) == ["data_1", "data_2"]
#     assert sorted(list(tensors["data_1"].keys())) == ["input", "target"]
#     assert isinstance(tensors["data_1"]["input"], torch.Tensor)
#     assert tensors["data_1"]["input"].shape == torch.Size((70, 10))
#     assert isinstance(tensors["data_1"]["target"], torch.Tensor)
#     assert tensors["data_1"]["target"].shape == torch.Size((70, 2))

#     assert sorted(list(tensors["data_2"].keys())) == ["input", "target"]
#     assert isinstance(tensors["data_2"]["input"], torch.Tensor)
#     assert tensors["data_2"]["input"].shape == torch.Size((50, 10))
#     assert isinstance(tensors["data_2"]["target"], torch.Tensor)
#     assert tensors["data_2"]["target"].shape == torch.Size((50, 2))