import pytest
import torch

from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples
from gnn4colliders.graphs import build_dgl_graph

dgl = pytest.importorskip("dgl")


def _sample(index: int, node_count: int) -> GraphSample:
    nodes = torch.arange(node_count * 3, dtype=torch.float32).reshape(node_count, 3)
    return GraphSample(
        graph=build_dgl_graph(nodes),
        label=torch.tensor(index),
        global_features=torch.tensor([float(index), -float(index)]),
        metadata=EventMetadata(index, index + 0.5, f"sample_{index}"),
    )


@pytest.mark.parametrize("counts", [(1,), (1, 2), (1, 3, 2)])
def test_graph_batch_preserves_heterogeneous_order_and_metadata(counts):
    samples = [_sample(index, count) for index, count in enumerate(counts)]
    batch = batch_graph_samples(samples)
    assert batch.labels.tolist() == list(range(len(counts)))
    assert batch.metadata.sample_id == tuple(f"sample_{i}" for i in range(len(counts)))
    assert batch.metadata.weight.tolist() == pytest.approx(
        [i + 0.5 for i in range(len(counts))]
    )
    assert batch.global_features[:, 0].tolist() == list(map(float, range(len(counts))))
    assert batch.graph.batch_num_nodes().tolist() == list(counts)
    assert batch.graph.batch_num_edges().tolist() == [
        count * (count - 1) if count > 1 else 1 for count in counts
    ]