import numbers
from typing import Dict, Generator, List, Optional, Tuple, Union

import pytest
import torch

from torch_geometric.data import Data, HeteroData
from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
from torch_geometric.transforms import Pad
from torch_geometric.transforms.pad import (
    AttrNamePadding,
    EdgeTypePadding,
    NodeTypePadding,
    Padding,
    UniformPadding,
)
from torch_geometric.typing import EdgeType, NodeType


def fake_data() -> Data:
    return FakeDataset(avg_num_nodes=10, avg_degree=5, edge_dim=2)[0]


def fake_hetero_data(node_types=2, edge_types=5) -> HeteroData:
    return FakeHeteroDataset(num_node_types=node_types,
                             num_edge_types=edge_types, avg_num_nodes=10,
                             edge_dim=2)[0]


def _generate_homodata_node_attrs(data: Data) -> Generator[str, None, None]:
    for attr in data.keys():
        if data.is_node_attr(attr):
            yield attr


def _generate_homodata_edge_attrs(data: Data) -> Generator[str, None, None]:
    for attr in data.keys():
        if data.is_edge_attr(attr):
            yield attr


def _generate_heterodata_nodes(
    data: HeteroData
) -> Generator[Tuple[NodeType, str, torch.Tensor], None, None]:
    for node_type, store in data.node_items():
        for attr in store.keys():
            yield node_type, attr


def _generate_heterodata_edges(
    data: HeteroData
) -> Generator[Tuple[EdgeType, str, torch.Tensor], None, None]:
    for edge_type, store in data.edge_items():
        for attr in store.keys():
            yield edge_type, attr


def _check_homo_data_nodes(
    original: Data,
    padded: Data,
    max_num_nodes: Union[int, Dict[NodeType, int]],
    node_pad_value: Optional[Padding] = None,
    is_mask_available: bool = False,
    exclude_keys: Optional[List[str]] = None,
):
    assert padded.num_nodes == max_num_nodes

    compare_pad_start_idx = original.num_nodes

    if is_mask_available:
        assert padded.pad_node_mask.numel() == padded.num_nodes
        assert torch.all(padded.pad_node_mask[:compare_pad_start_idx])
        assert not torch.any(padded.pad_node_mask[compare_pad_start_idx:])

    for attr in _generate_homodata_node_attrs(original):
        if attr in exclude_keys:
            assert attr not in padded.keys()
            continue

        assert attr in padded.keys()

        if not isinstance(padded[attr], torch.Tensor):
            continue

        assert padded[attr].shape[0] == max_num_nodes

        # Check values in padded area.
        pad_value = node_pad_value.get_value(
            None, attr) if node_pad_value is not None else 0.0
        assert all(
            i == pad_value
            for i in torch.flatten(padded[attr][compare_pad_start_idx:]))

        # Check values in non-padded area.
        assert torch.equal(original[attr],
                           padded[attr][:compare_pad_start_idx])


def _check_homo_data_edges(
    original: Data,
    padded: Data,
    max_num_edges: Optional[int] = None,
    edge_pad_value: Optional[Padding] = None,
    is_mask_available: bool = False,
    exclude_keys: Optional[List[str]] = None,
):
    # Check edge index attribute.
    if max_num_edges is None:
        max_num_edges = padded.num_nodes**2
    assert padded.num_edges == max_num_edges
    assert padded.edge_index.shape[1] == max_num_edges
    assert padded.edge_index.shape[1] == max_num_edges

    compare_pad_start_idx = original.num_edges
    expected_node = original.num_nodes

    # Check values in padded area.
    assert all(
        padded.edge_index[0, i] == padded.edge_index[1, i] == expected_node
        for i in range(compare_pad_start_idx, max_num_edges))
    # Check values in non-padded area.
    assert torch.equal(original.edge_index,
                       padded.edge_index[:, :compare_pad_start_idx])

    if is_mask_available:
        assert padded.pad_edge_mask.numel() == padded.num_edges
        assert torch.all(padded.pad_edge_mask[:compare_pad_start_idx])
        assert not torch.any(padded.pad_edge_mask[compare_pad_start_idx:])

    # Check other attributes.
    for attr in _generate_homodata_edge_attrs(original):
        if attr == 'edge_index':
            continue
        if attr in exclude_keys:
            assert attr not in padded.keys()
            continue

        assert attr in padded.keys()

        if not isinstance(padded[attr], torch.Tensor):
            continue

        assert padded[attr].shape[0] == max_num_edges

        # Check values in padded area.
        pad_value = edge_pad_value.get_value(
            None, attr) if edge_pad_value is not None else 0.0
        assert all(
            i == pad_value
            for i in torch.flatten(padded[attr][compare_pad_start_idx:, :]))

        # Check values in non-padded area.
        assert torch.equal(original[attr],
                           padded[attr][:compare_pad_start_idx, :])


def _check_hetero_data_nodes(
    original: HeteroData,
    padded: HeteroData,
    max_num_nodes: Union[int, Dict[NodeType, int]],
    node_pad_value: Optional[Padding] = None,
    is_mask_available: bool = False,
    exclude_keys: Optional[List[str]] = None,
):
    if is_mask_available:
        for store in padded.node_stores:
            assert 'pad_node_mask' in store

    expected_nodes = max_num_nodes

    for node_type, attr in _generate_heterodata_nodes(original):
        if attr in exclude_keys:
            assert attr not in padded[node_type].keys()
            continue

        assert attr in padded[node_type].keys()

        if not isinstance(padded[node_type][attr], torch.Tensor):
            continue

        compare_pad_start_idx = original[node_type].num_nodes
        padded_tensor = padded[node_type][attr]

        if attr == 'pad_node_mask':
            assert padded_tensor.numel() == padded[node_type].num_nodes
            assert torch.all(padded_tensor[:compare_pad_start_idx])
            assert not torch.any(padded_tensor[compare_pad_start_idx:])
            continue

        original_tensor = original[node_type][attr]

        # Check the number of nodes.
        if isinstance(max_num_nodes, dict):
            expected_nodes = max_num_nodes[node_type]

        assert padded_tensor.shape[0] == expected_nodes

        compare_pad_start_idx = original_tensor.shape[0]
        pad_value = node_pad_value.get_value(
            node_type, attr) if node_pad_value is not None else 0.0
        assert all(
            i == pad_value
            for i in torch.flatten(padded_tensor[compare_pad_start_idx:]))
        # Compare non-padded area with the original.
        assert torch.equal(original_tensor,
                           padded_tensor[:compare_pad_start_idx])


def _check_hetero_data_edges(
    original: HeteroData,
    padded: HeteroData,
    max_num_edges: Optional[Union[int, Dict[EdgeType, int]]] = None,
    edge_pad_value: Optional[Padding] = None,
    is_mask_available: bool = False,
    exclude_keys: Optional[List[str]] = None,
):
    if is_mask_available:
        for store in padded.edge_stores:
            assert 'pad_edge_mask' in store

    for edge_type, attr in _generate_heterodata_edges(padded):
        if attr in exclude_keys:
            assert attr not in padded[edge_type].keys()
            continue

        assert attr in padded[edge_type].keys()

        if not isinstance(padded[edge_type][attr], torch.Tensor):
            continue

        compare_pad_start_idx = original[edge_type].num_edges
        padded_tensor = padded[edge_type][attr]

        if attr == 'pad_edge_mask':
            assert padded_tensor.numel() == padded[edge_type].num_edges
            assert torch.all(padded_tensor[:compare_pad_start_idx])
            assert not torch.any(padded_tensor[compare_pad_start_idx:])
            continue

        original_tensor = original[edge_type][attr]

        if isinstance(max_num_edges, numbers.Number):
            expected_num_edges = max_num_edges
        elif max_num_edges is None or edge_type not in max_num_edges.keys():
            v1, _, v2 = edge_type
            expected_num_edges = padded[v1].num_nodes * padded[v2].num_nodes
        else:
            expected_num_edges = max_num_edges[edge_type]

        if attr == 'edge_index':
            # Check the number of edges.
            assert padded_tensor.shape[1] == expected_num_edges

            # Check padded area values.
            src_nodes = original[edge_type[0]].num_nodes
            assert all(
                i == src_nodes
                for i in torch.flatten(padded_tensor[0,
                                                     compare_pad_start_idx:]))
            dst_nodes = original[edge_type[2]].num_nodes
            assert all(
                i == dst_nodes
                for i in torch.flatten(padded_tensor[1,
                                                     compare_pad_start_idx:]))

            # Compare non-padded area with the original.
            assert torch.equal(original_tensor,
                               padded_tensor[:, :compare_pad_start_idx])
        else:
            # Check padded area size.
            assert padded_tensor.shape[0] == expected_num_edges

            # Check padded area values.
            pad_value = edge_pad_value.get_value(
                edge_type, attr) if edge_pad_value is not None else 0.0
            assert all(i == pad_value for i in torch.flatten(padded_tensor[
                compare_pad_start_idx:, :]))

            # Compare non-padded area with the original.
            assert torch.equal(original_tensor,
                               padded_tensor[:compare_pad_start_idx, :])


def _check_data(
    original: Union[Data, HeteroData],
    padded: Union[Data, HeteroData],
    max_num_nodes: Union[int, Dict[NodeType, int]],
    max_num_edges: Optional[Union[int, Dict[EdgeType, int]]] = None,
    node_pad_value: Optional[Union[Padding, int, float]] = None,
    edge_pad_value: Optional[Union[Padding, int, float]] = None,
    is_mask_available: bool = False,
    exclude_keys: Optional[List[str]] = None,
):

    if not isinstance(node_pad_value, Padding) and node_pad_value is not None:
        node_pad_value = UniformPadding(node_pad_value)
    if not isinstance(edge_pad_value, Padding) and edge_pad_value is not None:
        edge_pad_value = UniformPadding(edge_pad_value)

    if is_mask_available is None:
        is_mask_available = False

    if exclude_keys is None:
        exclude_keys = []

    if isinstance(original, Data):
        _check_homo_data_nodes(original, padded, max_num_nodes, node_pad_value,
                               is_mask_available, exclude_keys)
        _check_homo_data_edges(original, padded, max_num_edges, edge_pad_value,
                               is_mask_available, exclude_keys)
    else:
        _check_hetero_data_nodes(original, padded, max_num_nodes,
                                 node_pad_value, is_mask_available,
                                 exclude_keys)
        _check_hetero_data_edges(original, padded, max_num_edges,
                                 edge_pad_value, is_mask_available,
                                 exclude_keys)


def test_pad_repr():
    pad_str = 'Pad(max_num_nodes=10, max_num_edges=15, ' \
        'node_pad_value=UniformPadding(value=3.0), ' \
        'edge_pad_value=UniformPadding(value=1.5))'
    assert str(eval(pad_str)) == pad_str


@pytest.mark.parametrize('data', [fake_data(), fake_hetero_data()])
@pytest.mark.parametrize('num_nodes', [32, 64])
@pytest.mark.parametrize('add_pad_mask', [True, False])
def test_pad_auto_edges(data, num_nodes, add_pad_mask):
    transform = Pad(max_num_nodes=num_nodes, add_pad_mask=add_pad_mask)

    out = transform(data)
    _check_data(data, out, num_nodes, is_mask_available=add_pad_mask)


@pytest.mark.parametrize('num_nodes', [32, 64])
@pytest.mark.parametrize('num_edges', [300, 411])
@pytest.mark.parametrize('add_pad_mask', [True, False])
def test_pad_data_explicit_edges(num_nodes, num_edges, add_pad_mask):
    data = fake_data()
    transform = Pad(max_num_nodes=num_nodes, max_num_edges=num_edges,
                    add_pad_mask=add_pad_mask)

    out = transform(data)
    _check_data(data, out, num_nodes, num_edges,
                is_mask_available=add_pad_mask)


@pytest.mark.parametrize('num_nodes', [32, {'v0': 64, 'v1': 36}])
@pytest.mark.parametrize('num_edges', [300, {('v0', 'e0', 'v1'): 397}])
@pytest.mark.parametrize('add_pad_mask', [True, False])
def test_pad_heterodata_explicit_edges(num_nodes, num_edges, add_pad_mask):
    data = fake_hetero_data()
    transform = Pad(max_num_nodes=num_nodes, max_num_edges=num_edges,
                    add_pad_mask=add_pad_mask)

    out = transform(data)
    _check_data(data, out, num_nodes, num_edges,
                is_mask_available=add_pad_mask)


@pytest.mark.parametrize('node_pad_value', [10, AttrNamePadding({'x': 3.0})])
@pytest.mark.parametrize('edge_pad_value',
                         [11, AttrNamePadding({'edge_attr': 2.0})])
def test_pad_data_pad_values(node_pad_value, edge_pad_value):
    data = fake_data()
    num_nodes = 32
    transform = Pad(max_num_nodes=num_nodes, node_pad_value=node_pad_value,
                    edge_pad_value=edge_pad_value)
    out = transform(data)
    _check_data(data, out, num_nodes, node_pad_value=node_pad_value,
                edge_pad_value=edge_pad_value)


@pytest.mark.parametrize('node_pad_value', [
    UniformPadding(12),
    AttrNamePadding({'x': 0}),
    NodeTypePadding({
        'v0': UniformPadding(12),
        'v1': AttrNamePadding({'x': 7})
    })
])
@pytest.mark.parametrize('edge_pad_value', [
    UniformPadding(13),
    EdgeTypePadding({
        ('v0', 'e0', 'v1'):
        UniformPadding(13),
        ('v1', 'e0', 'v0'):
        AttrNamePadding({'edge_attr': UniformPadding(-1.0)})
    })
])
def test_pad_heterodata_pad_values(node_pad_value, edge_pad_value):
    data = fake_hetero_data()
    num_nodes = 32
    transform = Pad(max_num_nodes=num_nodes, node_pad_value=node_pad_value,
                    edge_pad_value=edge_pad_value)

    out = transform(data)
    _check_data(data, out, num_nodes, node_pad_value=node_pad_value,
                edge_pad_value=edge_pad_value)


@pytest.mark.parametrize('data', [fake_data(), fake_hetero_data()])
@pytest.mark.parametrize('add_pad_mask', [True, False])
@pytest.mark.parametrize('exclude_keys', [
    ['y'],
    ['edge_attr'],
    ['y', 'edge_attr'],
])
def test_pad_data_exclude_keys(data, add_pad_mask, exclude_keys):
    num_nodes = 32
    transform = Pad(max_num_nodes=num_nodes, add_pad_mask=add_pad_mask,
                    exclude_keys=exclude_keys)

    out = transform(data)
    _check_data(data, out, num_nodes, is_mask_available=add_pad_mask,
                exclude_keys=exclude_keys)


@pytest.mark.parametrize('is_hetero', [False, True])
def test_pad_invalid_max_num_nodes(is_hetero):
    if is_hetero:
        data = fake_hetero_data(node_types=1)
    else:
        data = fake_data()

    transform = Pad(max_num_nodes=data.num_nodes - 1)

    with pytest.raises(AssertionError, match="after padding"):
        transform(data)


@pytest.mark.parametrize('is_hetero', [False, True])
def test_pad_invalid_max_num_edges(is_hetero):
    if is_hetero:
        data = fake_hetero_data(node_types=1, edge_types=1)
    else:
        data = fake_data()

    transform = Pad(max_num_nodes=data.num_nodes + 10,
                    max_num_edges=data.num_edges - 1)

    with pytest.raises(AssertionError, match="after padding"):
        transform(data)


def test_pad_num_nodes_not_complete():
    data = fake_hetero_data(node_types=2, edge_types=1)
    transform = Pad(max_num_nodes={'v0': 100})

    with pytest.raises(KeyError):
        transform(data)


def test_pad_invalid_padding_type():
    with pytest.raises(ValueError, match="to be an integer or float"):
        Pad(max_num_nodes=100, node_pad_value='somestring')
    with pytest.raises(ValueError, match="to be an integer or float"):
        Pad(max_num_nodes=100, edge_pad_value='somestring')


def test_pad_data_non_tensor_attr():
    data = fake_data()
    batch_size = 13
    data.batch_size = batch_size

    transform = Pad(max_num_nodes=100)
    padded = transform(data)
    assert padded.batch_size == batch_size

    exclude_transform = Pad(max_num_nodes=101, exclude_keys=('batch_size', ))
    padded = exclude_transform(data)
    assert 'batch_size' not in padded.keys()


@pytest.mark.parametrize('mask_pad_value', [True, False])
def test_pad_node_additional_attr_mask(mask_pad_value):
    data = fake_data()
    mask = torch.randn(data.num_nodes) > 0
    mask_names = ['train_mask', 'test_mask', 'val_mask']
    for mask_name in mask_names:
        setattr(data, mask_name, mask)
    padding_num = 20

    max_num_nodes = int(data.num_nodes) + padding_num
    max_num_edges = data.num_edges + padding_num

    transform = Pad(max_num_nodes, max_num_edges, node_pad_value=0.1,
                    mask_pad_value=mask_pad_value)
    padded = transform(data)
    padded_masks = [getattr(padded, mask_name) for mask_name in mask_names]

    for padded_mask in padded_masks:
        assert padded_mask.ndim == 1
        assert padded_mask.size()[0] == max_num_nodes
        assert torch.all(padded_mask[-padding_num:] == mask_pad_value)


def test_uniform_padding():
    pad_val = 10.0
    p = UniformPadding(pad_val)
    assert p.get_value() == pad_val
    assert p.get_value("v1", "x") == pad_val

    p = UniformPadding()
    assert p.get_value() == 0.0

    with pytest.raises(ValueError, match="to be an integer or float"):
        UniformPadding('')


def test_attr_name_padding():
    x_val = 10.0
    y_val = 15.0
    default = 3.0
    padding_dict = {'x': x_val, 'y': UniformPadding(y_val)}
    padding = AttrNamePadding(padding_dict, default=default)

    assert padding.get_value(attr_name='x') == x_val
    assert padding.get_value('v1', 'x') == x_val
    assert padding.get_value(attr_name='y') == y_val
    assert padding.get_value('v1', 'y') == y_val
    assert padding.get_value(attr_name='x2') == default

    padding = AttrNamePadding({})
    assert padding.get_value(attr_name='x') == 0.0


def test_attr_name_padding_invalid():
    with pytest.raises(ValueError, match="to be a dictionary"):
        AttrNamePadding(10.0)

    with pytest.raises(ValueError, match="to be a string"):
        AttrNamePadding({10: 10.0})

    with pytest.raises(ValueError, match="to be of type"):
        AttrNamePadding({"x": {}})

    with pytest.raises(ValueError, match="to be of type"):
        AttrNamePadding({"x": {}})

    node_type_padding = NodeTypePadding({"x": 10.0})
    with pytest.raises(ValueError, match="to be of type"):
        AttrNamePadding({'x': node_type_padding})


@pytest.mark.parametrize('store_type', ['node', 'edge'])
def test_node_edge_type_padding(store_type):
    if store_type == "node":
        stores = ['v1', 'v2', 'v3', 'v4']
        padding_cls = NodeTypePadding
    else:
        stores = [('v1', 'e1', 'v1'), ('v1', 'e2', 'v1'), ('v1', 'e1', 'v2'),
                  ('v2', 'e1', 'v1')]
        padding_cls = EdgeTypePadding

    s0_default = 3.0
    s0_padding_dict = {'x': 10.0, 'y': -12.0}
    s0_padding = AttrNamePadding(s0_padding_dict, s0_default)
    s1_default = 0.1
    s1_padding_dict = {'y': 0.0, 'p': 13.0}
    s1_padding = AttrNamePadding(s1_padding_dict, s1_default)

    s2_default = 7.5
    store_default = -11.0
    padding_dict = {
        stores[0]: s0_padding,
        stores[1]: s1_padding,
        stores[2]: s2_default
    }
    padding = padding_cls(padding_dict, store_default)

    assert padding.get_value(stores[0], 'x') == s0_padding_dict['x']
    assert padding.get_value(stores[0], 'y') == s0_padding_dict['y']
    assert padding.get_value(stores[0], 'p') == s0_default
    assert padding.get_value(stores[0], 'z') == s0_default

    assert padding.get_value(stores[1], 'x') == s1_default
    assert padding.get_value(stores[1], 'y') == s1_padding_dict['y']
    assert padding.get_value(stores[1], 'p') == s1_padding_dict['p']
    assert padding.get_value(stores[1], 'z') == s1_default

    assert padding.get_value(stores[2], 'x') == s2_default
    assert padding.get_value(stores[2], 'z') == s2_default

    assert padding.get_value(stores[3], 'x') == store_default


def test_edge_padding_invalid():
    with pytest.raises(ValueError, match="to be a tuple"):
        EdgeTypePadding({'v1': 10.0})

    with pytest.raises(ValueError, match="got 1"):
        EdgeTypePadding({('v1', ): 10.0})

    with pytest.raises(ValueError, match="got 2"):
        EdgeTypePadding({('v1', 'v2'): 10.0})

    with pytest.raises(ValueError, match="got 4"):
        EdgeTypePadding({('v1', 'e2', 'v1', 'v2'): 10.0})