import numpy as np
import pytest
import torch

from torch_geometric.data import HeteroData
from torch_geometric.loader import HGTLoader
from torch_geometric.nn import GraphConv, to_hetero
from torch_geometric.testing import get_random_edge_index, withPackage
from torch_geometric.typing import SparseTensor
from torch_geometric.utils import k_hop_subgraph


def is_subset(subedge_index, edge_index, src_idx, dst_idx):
    num_nodes = int(edge_index.max()) + 1
    idx = num_nodes * edge_index[0] + edge_index[1]
    subidx = num_nodes * src_idx[subedge_index[0]] + dst_idx[subedge_index[1]]
    mask = torch.from_numpy(np.isin(subidx, idx))
    return int(mask.sum()) == mask.numel()


@withPackage('torch_sparse')
def test_hgt_loader():
    torch.manual_seed(12345)

    data = HeteroData()

    data['paper'].x = torch.arange(100)
    data['author'].x = torch.arange(100, 300)

    data['paper', 'paper'].edge_index = get_random_edge_index(100, 100, 500)
    data['paper', 'paper'].edge_attr = torch.arange(500)
    data['paper', 'author'].edge_index = get_random_edge_index(100, 200, 1000)
    data['paper', 'author'].edge_attr = torch.arange(500, 1500)
    data['author', 'paper'].edge_index = get_random_edge_index(200, 100, 1000)
    data['author', 'paper'].edge_attr = torch.arange(1500, 2500)

    r1, c1 = data['paper', 'paper'].edge_index
    r2, c2 = data['paper', 'author'].edge_index + torch.tensor([[0], [100]])
    r3, c3 = data['author', 'paper'].edge_index + torch.tensor([[100], [0]])
    full_adj = SparseTensor(
        row=torch.cat([r1, r2, r3]),
        col=torch.cat([c1, c2, c3]),
        value=torch.arange(2500),
    )

    batch_size = 20
    loader = HGTLoader(data, num_samples=[5] * 4, batch_size=batch_size,
                       input_nodes='paper')
    assert str(loader) == 'HGTLoader()'
    assert len(loader) == (100 + batch_size - 1) // batch_size

    for batch in loader:
        assert isinstance(batch, HeteroData)

        # Test node and types:
        assert set(batch.node_types) == {'paper', 'author'}
        assert set(batch.edge_types) == set(data.edge_types)

        assert len(batch['paper']) == 4
        assert batch['paper'].n_id.size() == (batch['paper'].num_nodes, )
        assert batch['paper'].x.size() == (40, )  # 20 + 4 * 5
        assert batch['paper'].input_id.numel() == batch_size
        assert batch['paper'].batch_size == batch_size
        assert batch['paper'].x.min() >= 0 and batch['paper'].x.max() < 100

        assert len(batch['author']) == 2
        assert batch['author'].n_id.size() == (batch['author'].num_nodes, )
        assert batch['author'].x.size() == (20, )  # 4 * 5
        assert batch['author'].x.min() >= 100 and batch['author'].x.max() < 300

        # Test edge type selection:
        assert set(batch.edge_types) == {('paper', 'to', 'paper'),
                                         ('paper', 'to', 'author'),
                                         ('author', 'to', 'paper')}

        assert len(batch['paper', 'paper']) == 3
        num_edges = batch['paper', 'paper'].num_edges
        assert batch['paper', 'paper'].e_id.size() == (num_edges, )
        row, col = batch['paper', 'paper'].edge_index
        value = batch['paper', 'paper'].edge_attr
        adj = full_adj[batch['paper'].x, batch['paper'].x]
        assert row.min() >= 0 and row.max() < 40
        assert col.min() >= 0 and col.max() < 40
        assert value.min() >= 0 and value.max() < 500
        assert adj.nnz() == row.size(0)
        assert torch.allclose(row.unique(), adj.storage.row().unique())
        assert torch.allclose(col.unique(), adj.storage.col().unique())
        assert torch.allclose(value.unique(), adj.storage.value().unique())

        assert is_subset(batch['paper', 'paper'].edge_index,
                         data['paper', 'paper'].edge_index, batch['paper'].x,
                         batch['paper'].x)

        assert len(batch['paper', 'author']) == 3
        num_edges = batch['paper', 'author'].num_edges
        assert batch['paper', 'author'].e_id.size() == (num_edges, )
        row, col = batch['paper', 'author'].edge_index
        value = batch['paper', 'author'].edge_attr
        adj = full_adj[batch['paper'].x, batch['author'].x]
        assert row.min() >= 0 and row.max() < 40
        assert col.min() >= 0 and col.max() < 20
        assert value.min() >= 500 and value.max() < 1500
        assert adj.nnz() == row.size(0)
        assert torch.allclose(row.unique(), adj.storage.row().unique())
        assert torch.allclose(col.unique(), adj.storage.col().unique())
        assert torch.allclose(value.unique(), adj.storage.value().unique())

        assert is_subset(batch['paper', 'author'].edge_index,
                         data['paper', 'author'].edge_index, batch['paper'].x,
                         batch['author'].x - 100)

        assert len(batch['author', 'paper']) == 3
        num_edges = batch['author', 'paper'].num_edges
        assert batch['author', 'paper'].e_id.size() == (num_edges, )
        row, col = batch['author', 'paper'].edge_index
        value = batch['author', 'paper'].edge_attr
        adj = full_adj[batch['author'].x, batch['paper'].x]
        assert row.min() >= 0 and row.max() < 20
        assert col.min() >= 0 and col.max() < 40
        assert value.min() >= 1500 and value.max() < 2500
        assert adj.nnz() == row.size(0)
        assert torch.allclose(row.unique(), adj.storage.row().unique())
        assert torch.allclose(col.unique(), adj.storage.col().unique())
        assert torch.allclose(value.unique(), adj.storage.value().unique())

        assert is_subset(batch['author', 'paper'].edge_index,
                         data['author', 'paper'].edge_index,
                         batch['author'].x - 100, batch['paper'].x)

        # Test for isolated nodes (there shouldn't exist any):
        n_id = torch.cat([batch['paper'].x, batch['author'].x])
        row, col, _ = full_adj[n_id, n_id].coo()
        assert torch.cat([row, col]).unique().numel() >= 59


@pytest.mark.dataset
@withPackage('torch_sparse')
def test_hgt_loader_on_cora(get_dataset):
    dataset = get_dataset(name='Cora')
    data = dataset[0]
    data.edge_weight = torch.rand(data.num_edges)

    hetero_data = HeteroData()
    hetero_data['paper'].x = data.x
    hetero_data['paper'].n_id = torch.arange(data.num_nodes)
    hetero_data['paper', 'paper'].edge_index = data.edge_index
    hetero_data['paper', 'paper'].edge_weight = data.edge_weight

    split_idx = torch.arange(5, 8)

    # Sample the complete two-hop neighborhood:
    loader = HGTLoader(hetero_data, num_samples=[data.num_nodes] * 2,
                       batch_size=split_idx.numel(),
                       input_nodes=('paper', split_idx))
    assert len(loader) == 1

    hetero_batch = next(iter(loader))
    batch_size = hetero_batch['paper'].batch_size

    n_id, _, _, e_mask = k_hop_subgraph(split_idx, num_hops=2,
                                        edge_index=data.edge_index,
                                        num_nodes=data.num_nodes)

    n_id = n_id.sort()[0]
    assert n_id.tolist() == hetero_batch['paper'].n_id.sort()[0].tolist()
    assert hetero_batch['paper', 'paper'].num_edges == int(e_mask.sum())

    class GNN(torch.nn.Module):
        def __init__(self, in_channels, hidden_channels, out_channels):
            super().__init__()
            self.conv1 = GraphConv(in_channels, hidden_channels)
            self.conv2 = GraphConv(hidden_channels, out_channels)

        def forward(self, x, edge_index, edge_weight):
            x = self.conv1(x, edge_index, edge_weight).relu()
            x = self.conv2(x, edge_index, edge_weight).relu()
            return x

    model = GNN(dataset.num_features, 16, dataset.num_classes)
    hetero_model = to_hetero(model, hetero_data.metadata())

    out1 = model(data.x, data.edge_index, data.edge_weight)[split_idx]
    out2 = hetero_model(hetero_batch.x_dict, hetero_batch.edge_index_dict,
                        hetero_batch.edge_weight_dict)['paper'][:batch_size]
    assert torch.allclose(out1, out2, atol=1e-6)


@withPackage('torch_sparse')
def test_hgt_loader_disconnected():
    data = HeteroData()

    data['paper'].x = torch.randn(10, 16)
    data['author'].x = torch.randn(10, 16)

    # Paper nodes are disconnected from author nodes:
    data['paper', 'paper'].edge_index = get_random_edge_index(10, 10, 15)
    data['paper', 'paper'].edge_attr = torch.randn(15, 8)
    data['author', 'author'].edge_index = get_random_edge_index(10, 10, 15)
    data['author', 'author'].edge_attr = torch.randn(15, 8)

    loader = HGTLoader(data, num_samples=[2], batch_size=2,
                       input_nodes='paper')

    for batch in loader:
        assert isinstance(batch, HeteroData)

        # Test node and edge types:
        assert set(batch.node_types) == set(data.node_types)
        assert set(batch.edge_types) == set(data.edge_types)

        assert batch['author'].num_nodes == 0
        assert batch['author'].x.size() == (0, 16)
        assert batch['author', 'author'].num_edges == 0
        assert batch['author', 'author'].edge_index.size() == (2, 0)
        assert batch['author', 'author'].edge_attr.size() == (0, 8)
