tests/test_variable_traversal.py

from __future__ import annotations

from pathlib import Path

import pytest

from gcmd_classifier.classification import (
    TermBranchSeed,
    VariableCandidate,
    build_variable_candidates,
    descend_variables,
    validate_variable_candidate_relationship,
)
from gcmd_classifier.config import ModelSettings
from gcmd_classifier.errors import StructuredModelResponseError, UnknownCandidateIDError
from gcmd_classifier.llm import FakeModelClient
from gcmd_classifier.llm.prompts import PromptCandidate
from gcmd_classifier.models import ArticleRecord, SupportType
from gcmd_classifier.vocabulary import build_vocabulary_index, load_vocabulary

FIXTURE_PATH = Path("tests/fixtures/gcmd_hierarchy_small.json")
FULL_HIERARCHY_PATH = Path("data/gcmd_hierarchy.json")


def _index():
    return load_vocabulary(FIXTURE_PATH)


def _article(abstract: str = "Carbon dioxide profiles and methane are discussed.") -> ArticleRecord:
    return ArticleRecord(
        DOI="10.example/variable-traversal",
        Title="Atmospheric carbon observations",
        Year=2025,
        Abstract=abstract,
    )


def _decision(candidate_id: str, confidence: float | None = 0.82) -> dict:
    return {
        "candidate_id": candidate_id,
        "confidence": confidence,
        "evidence": f"Evidence for {candidate_id}.",
        "support_type": "explicit",
        "reason": f"Reason for {candidate_id}.",
    }


def _term_branch(term_uuid: str = "term-atmospheric-chemistry") -> TermBranchSeed:
    index = _index()
    term = index.get(term_uuid)
    topic = index.get(term.parent_uuid or "")
    return TermBranchSeed(
        branch_id="topic:topic_0001/term:term_0001",
        parent_topic_uuid=topic.UUID,
        parent_topic_name=topic.name,
        term_uuid=term.UUID,
        term_name=term.name,
        term_level="Term",
        term_canonical_path=term.canonical_path,
        evidence="The article supports the Term.",
        support_type="explicit",
        confidence=0.77,
        reason="Term was selected.",
        candidate_id="term_0001",
        parent_branch_id="topic:topic_0001",
        prompt_version="term-test",
        model_provider="fake",
        model_name="fake-model",
        retry_count=0,
    )


def _multi_index():
    data = {
        "level": "Category",
        "name": "EARTH SCIENCE",
        "children": [
            {
                "level": "Topic",
                "name": "ATMOSPHERE",
                "UUID": "topic-atmosphere",
                "children": [
                    {
                        "level": "Term",
                        "name": "ATMOSPHERIC CHEMISTRY",
                        "UUID": "term-atmospheric-chemistry",
                        "children": [
                            {
                                "level": "Variable_Level_1",
                                "name": "CARBON",
                                "UUID": "vl1-carbon",
                                "children": [
                                    {
                                        "level": "Variable_Level_2",
                                        "name": "CARBON DIOXIDE",
                                        "UUID": "vl2-carbon-dioxide",
                                    }
                                ],
                            },
                            {
                                "level": "Variable_Level_1",
                                "name": "NITROGEN",
                                "UUID": "vl1-nitrogen",
                            },
                        ],
                    }
                ],
            }
        ],
    }
    return build_vocabulary_index(data, vocabulary_version="multi-variable-test")


def _multi_term_branch() -> TermBranchSeed:
    index = _multi_index()
    term = index.get("term-atmospheric-chemistry")
    topic = index.get(term.parent_uuid or "")
    return TermBranchSeed(
        branch_id="topic:topic_0001/term:term_0001",
        parent_topic_uuid=topic.UUID,
        parent_topic_name=topic.name,
        term_uuid=term.UUID,
        term_name=term.name,
        term_level="Term",
        term_canonical_path=term.canonical_path,
        evidence="The article supports the Term.",
        support_type="explicit",
        confidence=0.77,
        reason="Term was selected.",
        candidate_id="term_0001",
        parent_branch_id="topic:topic_0001",
        prompt_version="term-test",
        model_provider="fake",
        model_name="fake-model",
        retry_count=0,
    )


def test_variable_candidates_are_built_only_from_current_parent_direct_children() -> None:
    index = _index()
    candidates = build_variable_candidates(index, parent_uuid="term-atmospheric-chemistry")

    assert [candidate.variable_uuid for candidate in candidates] == ["vl1-atmosphere-carbon"]
    assert all(
        index.parent_of(candidate.variable_uuid) == "term-atmospheric-chemistry"
        for candidate in candidates
    )


def test_grandchildren_are_not_included_before_child_parent_is_selected() -> None:
    candidates = build_variable_candidates(_index(), parent_uuid="term-atmospheric-chemistry")

    assert "vl2-carbon-dioxide" not in [candidate.variable_uuid for candidate in candidates]
    assert "vl3-carbon-dioxide-profiles" not in [
        candidate.variable_uuid for candidate in candidates
    ]


def test_sibling_branches_are_not_included() -> None:
    candidates = build_variable_candidates(_index(), parent_uuid="vl1-atmosphere-carbon")

    assert [candidate.variable_uuid for candidate in candidates] == [
        "vl2-carbon-dioxide",
        "vl2-methane",
    ]
    assert "vl1-ocean-carbon" not in [candidate.variable_uuid for candidate in candidates]


def test_variable_candidate_ids_are_unique_within_prompt() -> None:
    candidates = build_variable_candidates(_index(), parent_uuid="vl1-atmosphere-carbon")
    candidate_ids = [candidate.candidate_id for candidate in candidates]

    assert candidate_ids == ["variable_0001", "variable_0002"]
    assert len(candidate_ids) == len(set(candidate_ids))


def test_variable_candidate_ids_map_to_uuid_bearing_variable_records() -> None:
    index = _index()

    for candidate in build_variable_candidates(index, parent_uuid="vl1-atmosphere-carbon"):
        record = index.get(candidate.variable_uuid)
        assert record.UUID
        assert record.level.startswith("Variable_Level_")
        assert candidate.prompt_candidate.candidate_id == candidate.candidate_id


def test_variable_candidate_construction_does_not_hard_code_names_or_uuids() -> None:
    index = _index()
    candidates = build_variable_candidates(index, parent_uuid="vl1-atmosphere-carbon")
    direct_children = [index.get(uuid) for uuid in index.children_of("vl1-atmosphere-carbon")]

    assert [candidate.variable_uuid for candidate in candidates] == [
        child.UUID for child in direct_children
    ]
    assert [candidate.prompt_candidate.name for candidate in candidates] == [
        child.name for child in direct_children
    ]


def test_parent_with_no_direct_children_is_terminal_outcome() -> None:
    result = descend_variables(
        article=_article(),
        term_branch=_term_branch("term-weather-events"),
        vocabulary=_index(),
        model_client=FakeModelClient([]),
        settings=ModelSettings(),
    )

    assert result.terminal_count == 1
    assert result.terminals[0].final_uuid == "term-weather-events"
    assert result.terminals[0].stop_reason == "No direct Variable children are available."


def test_stop_at_term_when_no_variable_child_supported() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": "No Variable child is adequately supported.",
            }
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.terminals[0].final_uuid == "term-atmospheric-chemistry"
    assert result.terminals[0].final_level == "Term"
    assert result.terminals[0].stop_reason == "No Variable child is adequately supported."


def test_stop_at_variable_level_1() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": "Variable_Level_1 is the deepest supported concept.",
            },
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.terminals[0].final_uuid == "vl1-atmosphere-carbon"
    assert result.terminals[0].final_level == "Variable_Level_1"
    assert result.terminals[0].stop_reason == "Variable_Level_1 is the deepest supported concept."


def test_stop_at_variable_level_2() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": "Variable_Level_2 is specific enough.",
            },
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.terminals[0].final_uuid == "vl2-carbon-dioxide"
    assert result.terminals[0].final_level == "Variable_Level_2"
    assert result.terminals[0].stop_reason == "Variable_Level_2 is specific enough."


def test_select_variable_level_3_and_stop_naturally_at_leaf() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(prompt_version_variable="variable-test"),
    )

    terminal = result.terminals[0]
    assert terminal.final_uuid == "vl3-carbon-dioxide-profiles"
    assert terminal.final_level == "Variable_Level_3"
    assert terminal.stop_reason == "Selected Variable_Level_3 is a leaf node."
    assert terminal.prompt_version == "variable-test"
    assert len(client.requests) == 3


def test_multiple_selected_variable_level_1_children_create_independent_branches() -> None:
    index = _multi_index()
    client = FakeModelClient(
        [
            {
                "selected": [_decision("variable_0001"), _decision("variable_0002")],
                "stop_at_parent": False,
            },
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": "Stop at carbon.",
            },
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_multi_term_branch(),
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert {terminal.final_uuid for terminal in result.terminals} == {"vl1-carbon", "vl1-nitrogen"}
    assert len({terminal.branch_id for terminal in result.terminals}) == 2


def test_multiple_selected_descendants_under_one_branch_continue_independently() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {
                "selected": [_decision("variable_0001"), _decision("variable_0002")],
                "stop_at_parent": False,
            },
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert {terminal.final_uuid for terminal in result.terminals} == {
        "vl3-carbon-dioxide-profiles",
        "vl2-methane",
    }


def test_sibling_branch_continuation_after_one_branch_stops() -> None:
    index = _multi_index()
    client = FakeModelClient(
        [
            {
                "selected": [_decision("variable_0001"), _decision("variable_0002")],
                "stop_at_parent": False,
            },
            {"selected": [], "stop_at_parent": True, "stop_reason": "Stop at first sibling."},
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_multi_term_branch(),
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert {terminal.final_uuid for terminal in result.terminals} == {"vl1-carbon", "vl1-nitrogen"}
    assert result.errors == ()


def test_sibling_branch_continuation_after_one_branch_fails() -> None:
    index = _multi_index()
    client = FakeModelClient(
        [
            {
                "selected": [_decision("variable_0001"), _decision("variable_0002")],
                "stop_at_parent": False,
            },
            {"selected": [_decision("unknown-variable")], "stop_at_parent": False},
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_multi_term_branch(),
        vocabulary=index,
        model_client=client,
        settings=ModelSettings(),
    )

    assert [terminal.final_uuid for terminal in result.terminals] == ["vl1-nitrogen"]
    assert result.has_errors is True
    assert result.errors[0].code == "UnknownCandidateIDError"


def test_no_forced_leaf_descent() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {
                "selected": [],
                "stop_at_parent": True,
                "stop_reason": "Do not force leaf descent.",
            },
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.terminals[0].final_uuid == "vl1-atmosphere-carbon"
    assert len(client.requests) == 2


def test_no_skipped_hierarchy_levels_or_non_direct_descendant_evaluation() -> None:
    client = FakeModelClient(
        [{"selected": [_decision("vl3-carbon-dioxide-profiles")], "stop_at_parent": False}]
    )

    with pytest.raises(UnknownCandidateIDError):
        descend_variables(
            article=_article(),
            term_branch=_term_branch(),
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_unknown_variable_candidate_id_is_rejected() -> None:
    client = FakeModelClient(
        [{"selected": [_decision("unknown-variable")], "stop_at_parent": False}]
    )

    with pytest.raises(UnknownCandidateIDError):
        descend_variables(
            article=_article(),
            term_branch=_term_branch(),
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_duplicate_variable_candidate_id_is_rejected() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [_decision("variable_0001"), _decision("variable_0001")],
                "stop_at_parent": False,
            }
        ]
    )

    with pytest.raises(UnknownCandidateIDError):
        descend_variables(
            article=_article(),
            term_branch=_term_branch(),
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_empty_or_malformed_variable_candidate_id_rejected_by_structured_validation() -> None:
    client = FakeModelClient([{"selected": [_decision("")], "stop_at_parent": False}])

    with pytest.raises(StructuredModelResponseError):
        descend_variables(
            article=_article(),
            term_branch=_term_branch(),
            vocabulary=_index(),
            model_client=client,
            settings=ModelSettings(),
        )


def test_selected_candidate_from_sibling_parent_is_rejected() -> None:
    invalid_candidate = VariableCandidate(
        candidate_id="variable_0001",
        parent_uuid="vl1-ocean-carbon",
        variable_uuid="vl2-carbon-dioxide",
        prompt_candidate=PromptCandidate(
            candidate_id="variable_0001",
            name="CARBON DIOXIDE",
            level="Variable_Level_2",
        ),
    )

    with pytest.raises(ValueError):
        validate_variable_candidate_relationship(
            invalid_candidate,
            selected_parent_uuid="vl1-atmosphere-carbon",
            index=_index(),
        )


def test_manually_invalid_candidate_to_record_mapping_is_rejected() -> None:
    invalid_candidate = VariableCandidate(
        candidate_id="variable_0001",
        parent_uuid="vl1-atmosphere-carbon",
        variable_uuid="vl1-ocean-carbon",
        prompt_candidate=PromptCandidate(
            candidate_id="variable_0001",
            name="CARBON",
            level="Variable_Level_1",
        ),
    )

    with pytest.raises(ValueError):
        validate_variable_candidate_relationship(
            invalid_candidate,
            selected_parent_uuid="vl1-atmosphere-carbon",
            index=_index(),
        )


def test_terminal_outcome_includes_provenance_and_required_fields() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001", confidence=0.66)], "stop_at_parent": False},
            {"selected": [], "stop_at_parent": True, "stop_reason": "Stop at VL1."},
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(model_name="unit-test-model"),
    )
    terminal = result.terminals[0]

    assert terminal.branch_id.endswith("/variable:variable_0001")
    assert terminal.parent_branch_id == _term_branch().branch_id
    assert terminal.topic_uuid == "topic-atmosphere"
    assert terminal.topic_name == "ATMOSPHERE"
    assert terminal.term_uuid == "term-atmospheric-chemistry"
    assert terminal.term_name == "ATMOSPHERIC CHEMISTRY"
    assert terminal.final_uuid == "vl1-atmosphere-carbon"
    assert terminal.final_name == "CARBON"
    assert terminal.final_level == "Variable_Level_1"
    assert terminal.final_canonical_path.endswith("CARBON")
    assert terminal.path_components[-1] == "CARBON"
    assert terminal.evidence == "Evidence for variable_0001."
    assert terminal.support_type is SupportType.EXPLICIT
    assert terminal.confidence == 0.66
    assert terminal.reason == "Reason for variable_0001."
    assert terminal.stop_reason == "Stop at VL1."
    assert terminal.candidate_id == "variable_0001"
    assert terminal.model_name == "unit-test-model"
    assert len(terminal.evidence_trail) == 2


def test_child_branch_ids_are_distinct_and_preserve_parent_context() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
            {
                "selected": [_decision("variable_0001"), _decision("variable_0002")],
                "stop_at_parent": False,
            },
            {"selected": [_decision("variable_0001")], "stop_at_parent": False},
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    branch_ids = [terminal.branch_id for terminal in result.terminals]
    assert len(branch_ids) == len(set(branch_ids))
    assert all(branch_id.startswith(_term_branch().branch_id) for branch_id in branch_ids)


def test_confidence_is_preserved_only_as_uncalibrated_metadata() -> None:
    client = FakeModelClient(
        [
            {"selected": [_decision("variable_0001", confidence=0.01)], "stop_at_parent": False},
            {"selected": [], "stop_at_parent": True, "stop_reason": "Stop despite low confidence."},
        ]
    )

    result = descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert result.terminals[0].confidence == 0.01
    assert result.terminals[0].final_uuid == "vl1-atmosphere-carbon"


def test_variable_descent_uses_provider_neutral_model_interface() -> None:
    client = FakeModelClient(
        [{"selected": [], "stop_at_parent": True, "stop_reason": "Stop at Term."}]
    )

    descend_variables(
        article=_article(),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(provider="fake-provider", model_name="fake-model"),
    )

    request = client.requests[0]
    assert request.provider == "fake-provider"
    assert request.model_name == "fake-model"
    assert request.response_schema.__name__ == "VariableResponse"


def test_fake_model_receives_article_parent_context_and_direct_variable_candidates() -> None:
    client = FakeModelClient(
        [{"selected": [], "stop_at_parent": True, "stop_reason": "Stop at Term."}]
    )
    article = _article()

    descend_variables(
        article=article,
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    prompt = client.requests[0].prompt
    assert article.DOI in prompt
    assert article.Title in prompt
    assert article.Abstract in prompt
    assert "ATMOSPHERIC CHEMISTRY" in prompt
    assert "candidate_id: variable_0001" in prompt
    assert "CARBON" in prompt
    assert "CARBON DIOXIDE" not in prompt


def test_empty_abstract_remains_valid_and_is_passed_to_variable_prompt() -> None:
    client = FakeModelClient(
        [{"selected": [], "stop_at_parent": True, "stop_reason": "Stop at Term."}]
    )

    descend_variables(
        article=_article(abstract=""),
        term_branch=_term_branch(),
        vocabulary=_index(),
        model_client=client,
        settings=ModelSettings(),
    )

    assert "<ABSTRACT>\n\n</ABSTRACT>" in client.requests[0].prompt


def test_full_data_current_variable_candidate_counts_are_direct_children_only() -> None:
    index = load_vocabulary(FULL_HIERARCHY_PATH)
    counts = {"Variable_Level_1": 0, "Variable_Level_2": 0, "Variable_Level_3": 0}
    variable_uuids: list[str] = []
    parent_uuids = [
        uuid
        for uuid, record in index.records_by_uuid.items()
        if record.level in {"Term", "Variable_Level_1", "Variable_Level_2"}
    ]

    for parent_uuid in parent_uuids:
        for candidate in build_variable_candidates(index, parent_uuid=parent_uuid):
            record = index.get(candidate.variable_uuid)
            counts[record.level] += 1
            variable_uuids.append(record.UUID)
            assert index.parent_of(record.UUID) == parent_uuid
            assert record.UUID

    assert counts["Variable_Level_1"] == 1368
    assert counts["Variable_Level_2"] == 1456
    assert counts["Variable_Level_3"] == 553
    assert len(variable_uuids) == len(set(variable_uuids))