"""Smoke tests against a small fixture derived from the real ttH ROOT sample."""

from __future__ import annotations

import os
from pathlib import Path

import awkward as ak
import pytest
import torch
import uproot

from gnn4colliders.features import build_node_features
from gnn4colliders.graphs import build_dgl_graph

pytest.importorskip("dgl")

pytestmark = [pytest.mark.integration, pytest.mark.real_data]

FEATURE_BRANCHES = [
    ["jet_pt", "ele_pt", "mu_pt", "ph_pt", "MET_met"],
    ["jet_eta", "ele_eta", "mu_eta", "ph_eta", 0],
    ["jet_phi", "ele_phi", "mu_phi", "ph_phi", "MET_phi"],
    "CALC_E",
    ["jet_btag", 0, 0, 0, 0],
    [0, "ele_charge", "mu_charge", 0, 0],
    "NODE_TYPE",
]
OBJECT_TYPES = ["vector", "vector", "vector", "vector", "single"]
FEATURE_SCALES = [0.1, 1, 1, 0.1, 1, 1, 1]


def _fixture_path() -> Path:
    configured = os.environ.get("GNN4COLLIDERS_ROOT_FIXTURE")
    if configured:
        return Path(configured)
    return Path("data/processed/ttH_NLO_64.root")


@pytest.fixture(scope="module")
def root_tree():
    path = _fixture_path()
    if not path.exists():
        pytest.skip(
            f"ROOT sample fixture is absent: {path}; download ttH_NLO.root and "
            "create the reduced fixture first"
        )
    root_file = uproot.open(path)
    tree = root_file["output"]
    try:
        yield tree
    finally:
        root_file.close()


def test_real_root_sample_has_expected_fixture_shape(root_tree):
    assert root_tree.num_entries == 64
    assert {"jet_pt", "MET_met", "weight", "Number"}.issubset(root_tree.keys())


def test_real_root_events_build_features_and_graphs(root_tree):
    arrays = root_tree.arrays(entry_start=0, entry_stop=8, library="ak")
    for event in ak.Array(arrays):
        event_features, lengths = build_node_features(
            event, FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES
        )
        graph = build_dgl_graph(event_features)

        assert event_features.dtype == torch.float32
        assert event_features.shape == (sum(lengths), 7)
        assert graph.number_of_nodes() == event_features.shape[0]
        expected_edges = event_features.shape[0] * max(event_features.shape[0] - 1, 1)
        assert graph.number_of_edges() == expected_edges
        assert graph.edata["features"].shape == (expected_edges, 3)
        assert torch.isfinite(event_features).all()
        assert torch.isfinite(graph.edata["features"]).all()