/tests/unit/test_tda.py

"""Tests for TDA (Task Difficulty Assessment) computation.

Unit tests covering TDAComputer.compute_tdi, _estimate_horizon,
_compute_success_rate (Laplace smoothing), and _compute_evidence_confidence.
"""

from __future__ import annotations

import pytest

from excalibur.planner.models import (
    AttackNode,
    AttackTree,
    EvidenceLevel,
    NodeStatus,
    NodeType,
    TDIScore,
)
from excalibur.planner.tda import TDAComputer


def _build_tree_with_chain(
    depth: int,
    evidence: EvidenceLevel = EvidenceLevel.SPECULATIVE,
) -> tuple[AttackTree, list[AttackNode]]:
    """Build a linear chain of *depth* nodes rooted at index 0.

    Returns:
        Tuple of (tree, ordered_nodes) where ordered_nodes[0] is root.
    """
    nodes: list[AttackNode] = []
    for i in range(depth):
        parent_id = nodes[i - 1].id if i > 0 else None
        node = AttackNode(
            id=f"n{i}",
            node_type=NodeType.ACTION,
            status=NodeStatus.ACTIVE,
            parent_id=parent_id,
            evidence_level=evidence,
        )
        nodes.append(node)

    tree = AttackTree(root_id=nodes[0].id)
    for n in nodes:
        tree.add_node(n)
    return tree, nodes


@pytest.mark.unit
class TestTDAComputerComputeTdi:
    """Tests for TDAComputer.compute_tdi producing valid scores."""

    def test_compute_tdi_returns_tdi_score(self) -> None:
        """compute_tdi returns an instance of TDIScore."""
        tree, nodes = _build_tree_with_chain(1)
        computer = TDAComputer()
        result = computer.compute_tdi(nodes[0], tree)
        assert isinstance(result, TDIScore)

    def test_tdi_value_in_range(self) -> None:
        """Computed TDI value is between 0 and 1."""
        tree, nodes = _build_tree_with_chain(3)
        computer = TDAComputer()
        for node in nodes:
            tdi = computer.compute_tdi(node, tree, context_load=0.5)
            assert 0.0 <= tdi.value <= 1.0

    def test_context_load_propagated(self) -> None:
        """context_load argument is reflected in the TDIScore."""
        tree, nodes = _build_tree_with_chain(1)
        computer = TDAComputer()
        tdi = computer.compute_tdi(nodes[0], tree, context_load=0.75)
        assert tdi.context_load == pytest.approx(0.75)

    def test_custom_weights_used(self) -> None:
        """Custom weights from the constructor are passed through."""
        custom = {
            "horizon": 0.1,
            "evidence": 0.2,
            "context": 0.3,
            "success": 0.4,
        }
        computer = TDAComputer(weights=custom)
        tree, nodes = _build_tree_with_chain(1)
        tdi = computer.compute_tdi(nodes[0], tree)
        assert tdi.weight_horizon == pytest.approx(0.1)
        assert tdi.weight_evidence == pytest.approx(0.2)
        assert tdi.weight_context == pytest.approx(0.3)
        assert tdi.weight_success == pytest.approx(0.4)


@pytest.mark.unit
class TestEstimateHorizon:
    """Tests for TDAComputer._estimate_horizon."""

    def test_single_node_horizon(self) -> None:
        """A single node has horizon == 1.0 (depth 1 / max 1)."""
        tree, nodes = _build_tree_with_chain(1)
        computer = TDAComputer()
        h = computer._estimate_horizon(nodes[0], tree)
        assert h == pytest.approx(1.0)

    def test_deeper_nodes_have_higher_horizon(self) -> None:
        """Deeper nodes receive a higher horizon value."""
        tree, nodes = _build_tree_with_chain(5)
        computer = TDAComputer()
        h_root = computer._estimate_horizon(nodes[0], tree)
        h_leaf = computer._estimate_horizon(nodes[-1], tree)
        assert h_leaf > h_root

    def test_horizon_bounded_zero_to_one(self) -> None:
        """Horizon values always stay in [0, 1]."""
        tree, nodes = _build_tree_with_chain(10)
        computer = TDAComputer()
        for node in nodes:
            h = computer._estimate_horizon(node, tree)
            assert 0.0 <= h <= 1.0


@pytest.mark.unit
class TestComputeSuccessRate:
    """Tests for TDAComputer._compute_success_rate (Laplace smoothing)."""

    def test_zero_visits(self) -> None:
        """With zero visits Laplace smoothing gives (0+1)/(0+2) = 0.5."""
        node = AttackNode(visit_count=0, success_count=0)
        rate = TDAComputer._compute_success_rate(node)
        assert rate == pytest.approx(0.5)

    def test_all_successes(self) -> None:
        """All successes approaches 1 but stays below due to smoothing."""
        node = AttackNode(visit_count=10, success_count=10)
        rate = TDAComputer._compute_success_rate(node)
        # (10+1)/(10+2) = 11/12 ~ 0.9167
        assert rate == pytest.approx(11.0 / 12.0)
        assert rate < 1.0

    def test_all_failures(self) -> None:
        """Zero successes approaches 0 but stays above due to smoothing."""
        node = AttackNode(visit_count=10, success_count=0)
        rate = TDAComputer._compute_success_rate(node)
        # (0+1)/(10+2) = 1/12 ~ 0.0833
        assert rate == pytest.approx(1.0 / 12.0)
        assert rate > 0.0

    def test_mixed_results(self) -> None:
        """Mixed success/visit counts are correctly smoothed."""
        node = AttackNode(visit_count=8, success_count=3)
        rate = TDAComputer._compute_success_rate(node)
        # (3+1)/(8+2) = 4/10 = 0.4
        assert rate == pytest.approx(0.4)


@pytest.mark.unit
class TestComputeEvidenceConfidence:
    """Tests for TDAComputer._compute_evidence_confidence."""

    def test_single_verified_node(self) -> None:
        """Single VERIFIED node produces confidence of 1.0."""
        tree, nodes = _build_tree_with_chain(1, evidence=EvidenceLevel.VERIFIED)
        computer = TDAComputer()
        ec = computer._compute_evidence_confidence(nodes[0], tree)
        assert ec == pytest.approx(1.0)

    def test_path_mean_confidence(self) -> None:
        """Confidence is the mean evidence level along the path."""
        root = AttackNode(
            id="r",
            evidence_level=EvidenceLevel.VERIFIED,
        )
        child = AttackNode(
            id="c",
            parent_id="r",
            evidence_level=EvidenceLevel.SPECULATIVE,
        )
        tree = AttackTree(root_id="r")
        tree.add_node(root)
        tree.add_node(child)

        computer = TDAComputer()
        ec = computer._compute_evidence_confidence(child, tree)
        expected = (EvidenceLevel.SPECULATIVE.value + EvidenceLevel.VERIFIED.value) / 2.0
        assert ec == pytest.approx(expected)

    def test_uniform_speculative(self) -> None:
        """All SPECULATIVE nodes yield confidence of 0.3."""
        tree, nodes = _build_tree_with_chain(4, evidence=EvidenceLevel.SPECULATIVE)
        computer = TDAComputer()
        ec = computer._compute_evidence_confidence(nodes[-1], tree)
        assert ec == pytest.approx(0.3)