tests/unit/data/test_root_dataset.py

from pathlib import Path

import awkward as ak
import numpy as np
import pytest
import uproot

from gnn4colliders.data import EventSample, RootEventDataset
from gnn4colliders.data.root_io import branch_names_from_specs, read_tree
from gnn4colliders.features import build_object_features


def _write_fixture(path: Path) -> None:
    with uproot.recreate(path) as root_file:
        root_file["events"] = {
            "eventNumber": np.array([10, 11, 12], dtype=np.int64),
            "weight": np.array([1.5, -2.0, 0.25], dtype=np.float32),
            "label_branch": np.array([4, 5, 6], dtype=np.int32),
            "jet_pt": ak.Array([[10.0, 20.0], [], [30.0]]),
            "global_x": np.array([0.5, 0.6, 0.7], dtype=np.float32),
        }


def test_branch_selection_and_root_reading(tmp_path: Path):
    path = tmp_path / "events.root"
    _write_fixture(path)
    assert branch_names_from_specs(
        [["jet_pt"], "CALC_E", "NODE_TYPE"],
        tracking_info=["eventNumber"],
        global_features=["global_x"],
    ) == ("jet_pt", "eventNumber", "global_x")
    arrays = read_tree(path, "events", ["weight"])
    assert arrays.fields == ["weight"]
    assert arrays["weight"].to_list() == [1.5, -2.0, 0.25]


def test_dataset_returns_architecture_neutral_samples(tmp_path: Path):
    path = tmp_path / "events.root"
    _write_fixture(path)
    dataset = RootEventDataset(
        path,
        tree_name="events",
        label="label_branch",
        feature_branches=[["jet_pt"]],
        global_features=["global_x"],
        fold_var="eventNumber",
        weight_var="weight",
    )

    assert len(dataset) == 3
    assert dataset.branches == (
        "jet_pt",
        "eventNumber",
        "weight",
        "global_x",
        "label_branch",
    )
    sample = dataset[1]
    assert isinstance(sample, EventSample)
    assert sample.event_index == 1
    assert sample.objects["jet_pt"].to_list() == []
    assert sample.label == 5
    assert sample.tracking.tolist() == [11.0, -2.0]
    assert sample.global_features.tolist() == pytest.approx([0.6])
    assert dataset[-1].event_index == 2


def test_iter_range_preserves_order_across_files(tmp_path: Path):
    first = tmp_path / "first.root"
    second = tmp_path / "second.root"
    _write_fixture(first)
    _write_fixture(second)
    dataset = RootEventDataset(
        [first, second],
        tree_name="events",
        label=[0, 1],
        feature_branches=[["jet_pt"]],
        fold_var="eventNumber",
        weight_var="weight",
    )

    ranged = list(dataset.iter_range(1, 5))

    assert [sample.event_index for sample in ranged] == [1, 2, 3, 4]
    assert [sample.label for sample in ranged] == [0, 0, 1, 1]
    assert [sample.metadata.extra["entry_index"] for sample in ranged] == [1, 2, 0, 1]
    assert [sample.metadata.sample_id for sample in ranged] == [
        f"{first}:events:1",
        f"{first}:events:2",
        f"{second}:events:0",
        f"{second}:events:1",
    ]


def test_multiple_files_keep_input_order_and_feed_feature_builder(tmp_path: Path):
    first = tmp_path / "first.root"
    second = tmp_path / "second.root"
    _write_fixture(first)
    _write_fixture(second)
    dataset = RootEventDataset(
        [first, second],
        tree_name="events",
        label=[0, 1],
        feature_branches=[["jet_pt"], [0], [0], ["CALC_E"]],
        fold_var="eventNumber",
    )

    assert len(dataset) == 6
    assert dataset[3].event_index == 3
    assert dataset[3].label == 1
    features, lengths = build_object_features(
        dataset[0].objects,
        [["jet_pt"], [0], [0], ["CALC_E"]],
        ["vector"],
        [1, 1, 1, 1],
    )
    assert features.shape == (2, 4)
    assert lengths == [2]


def test_root_boundary_reports_missing_tree_and_branch(tmp_path: Path):
    path = tmp_path / "events.root"
    _write_fixture(path)
    with pytest.raises(KeyError, match="tree"):
        read_tree(path, "missing", ["weight"])
    with pytest.raises(KeyError, match="missing requested branches"):
        read_tree(path, "events", ["missing_branch"])


def test_root_boundary_reports_missing_file(tmp_path: Path):
    with pytest.raises((OSError, FileNotFoundError)):
        read_tree(tmp_path / "does-not-exist.root", "events", ["weight"])


def test_sample_identity_is_stable_and_distinguishes_files(tmp_path: Path):
    first = tmp_path / "first.root"
    second = tmp_path / "second.root"
    _write_fixture(first)
    _write_fixture(second)
    kwargs = dict(
        tree_name="events",
        label="label_branch",
        feature_branches=[["jet_pt"]],
        fold_var="eventNumber",
        weight_var="weight",
    )
    left = RootEventDataset(first, **kwargs)
    repeated = RootEventDataset(first, **kwargs)
    both = RootEventDataset([first, second], **kwargs)
    assert left[0].metadata.sample_id == repeated[0].metadata.sample_id
    assert left[0].metadata.sample_id != left[1].metadata.sample_id
    assert both[0].metadata.sample_id != both[3].metadata.sample_id