from collections import namedtuple
from dataclasses import dataclass

import pytest

from .torch_parsing import list_from_torch

pytest.importorskip("torch")

from torch.utils.data import DataLoader as TorchDataLoader  # noqa: E402
from torch.utils.data import Dataset as TorchDataset  # noqa: E402


def test_list_from_torch_dict_dataset(data: list[dict]):
    class PytorchDictDataset(TorchDataset):
        def __init__(self):
            self.data = data

        def __getitem__(self, i):
            return self.data[i]

        def __len__(self):
            return len(self.data)

    dataset = PytorchDictDataset()
    data_list = list_from_torch(dataset)

    assert isinstance(data_list, list)
    assert len(data_list) == len(dataset)
    assert set(list(data_list[0].keys())) == {"value", "label", "key", "score", "source_id", "partition_id"}


def test_list_from_torch_dataloader(data: list[dict]):
    class PytorchDictDataset(TorchDataset):
        def __init__(self):
            self.data = data

        def __getitem__(self, i):
            return self.data[i]

        def __len__(self):
            return len(self.data)

    dataset = PytorchDictDataset()

    def collate_fn(x: list[dict]):
        return {"value": [item["value"] for item in x], "label": [item["label"] for item in x]}

    dataloader = TorchDataLoader(dataset, batch_size=3, collate_fn=collate_fn)
    data_list = list_from_torch(dataloader)

    assert isinstance(data_list, list)
    assert len(data_list) == len(dataset)
    assert list(data_list[0].keys()) == ["value", "label"]


def test_list_from_torch_tuple_dataset(data: list[dict]):
    class PytorchTupleDataset(TorchDataset):
        def __init__(self):
            self.data = data

        def __getitem__(self, i):
            return self.data[i]["value"], self.data[i]["label"]

        def __len__(self):
            return len(self.data)

    dataset = PytorchTupleDataset()

    # raises error if no column names are passed in
    with pytest.raises(ValueError):
        list_from_torch(dataset)

    # raises error if not enough column names are passed in
    with pytest.raises(ValueError):
        list_from_torch(dataset, column_names=["value"])

    # creates list if correct number of column names are passed in
    data_list = list_from_torch(dataset, column_names=["value", "label"])
    assert isinstance(data_list, list)
    assert len(data_list) == len(dataset)
    assert list(data_list[0].keys()) == ["value", "label"]


def test_list_from_torch_named_tuple_dataset(data: list[dict]):
    # Given a Pytorch dataset that returns a namedtuple for each item
    DatasetTuple = namedtuple("DatasetTuple", ["value", "label"])

    class PytorchNamedTupleDataset(TorchDataset):
        def __init__(self):
            self.data = data

        def __getitem__(self, i):
            return DatasetTuple(self.data[i]["value"], self.data[i]["label"])

        def __len__(self):
            return len(self.data)

    dataset = PytorchNamedTupleDataset()
    data_list = list_from_torch(dataset)
    assert isinstance(data_list, list)
    assert len(data_list) == len(dataset)
    assert list(data_list[0].keys()) == ["value", "label"]


def test_list_from_torch_dataclass_dataset(data: list[dict]):
    @dataclass
    class DatasetItem:
        text: str
        label: int

    class PytorchDataclassDataset(TorchDataset):
        def __init__(self):
            self.data = data

        def __getitem__(self, i):
            return DatasetItem(text=self.data[i]["value"], label=self.data[i]["label"])

        def __len__(self):
            return len(self.data)

    dataset = PytorchDataclassDataset()
    data_list = list_from_torch(dataset)
    assert isinstance(data_list, list)
    assert len(data_list) == len(dataset)
    assert list(data_list[0].keys()) == ["text", "label"]


def test_list_from_torch_invalid_dataset(data: list[dict]):
    class PytorchInvalidDataset(TorchDataset):
        def __init__(self):
            self.data = data

        def __getitem__(self, i):
            return [self.data[i]["value"], self.data[i]["label"]]

        def __len__(self):
            return len(self.data)

    dataset = PytorchInvalidDataset()
    with pytest.raises(ValueError):
        list_from_torch(dataset)
