arborist/tests/test_canonical_projection.py

"""Tests for the canonical-projection preflight + ``arborist canon`` CLI.

Two surfaces, one substrate. The π* registry already has unit-level
coverage in ``tests/test_pi_star.py``; these tests exercise the ways
those π*'s become reachable from the user — without RAG, without an
LLM, and without a shard.

- ``_canonical_projection_preflight`` — the sniff + π* dispatch
  helper inside ``arborist.qa.query``.
- ``query()`` short-circuit path — math/logic-shaped questions return
  ``status="canonical_projection"`` before retrieval starts.
- ``arborist canon ...`` CLI subcommand — direct π* call.
"""

from __future__ import annotations

import json
import subprocess
import sys
from pathlib import Path

import pytest

from arborist.qa.client import StubClient
from arborist.qa.query import _canonical_projection_preflight, query


# ----- preflight helper ---------------------------------------------------


@pytest.mark.parametrize(
    "question,expected_ref,expected_canonical",
    [
        ("0.1 + 0.2", "arithmetic@v1", b"3/10"),
        ("0.3", "arithmetic@v1", b"3/10"),
        ("1/3 + 1/6", "arithmetic@v1", b"1/2"),
        ("6/4", "arithmetic@v1", b"3/2"),
        ("(2+3)*4", "arithmetic@v1", b"20/1"),
        ("A IMPL B", "logic-kernel@v1", b"(NOT A OR B)"),
        ("(NOT B) IMPL (NOT A)", "logic-kernel@v1", b"(NOT A OR B)"),
        ("A OR NOT A", "logic-kernel@v1", b"TRUE"),
        ("NOT NOT A", "logic-kernel@v1", b"A"),
    ],
)
def test_preflight_matches_math_or_logic(question, expected_ref, expected_canonical):
    result = _canonical_projection_preflight(question)
    assert result is not None, f"expected match, got None for {question!r}"
    ref, canonical = result
    assert ref == expected_ref
    assert canonical == expected_canonical


# ----- Algebra route (ticket #000030) ------------------------------------


@pytest.mark.parametrize(
    "question",
    [
        "(x+1)**2",
        "x**2 + 3*x + 7",
        "a*b - c",
        "x",
        "sin(x)",
    ],
)
def test_preflight_algebra_route_matches(question):
    """Symbolic-algebra inputs route through algebra-symbolic@v1
    (when sympy is available)."""
    pytest.importorskip("sympy")
    result = _canonical_projection_preflight(question)
    assert result is not None, f"expected match, got None for {question!r}"
    ref, canonical = result
    assert ref == "algebra-symbolic@v1"
    # srepr-form bytes — start with a SymPy class name.
    assert canonical[:1].isalpha()


def test_preflight_algebra_collapses_polynomial_identity():
    pytest.importorskip("sympy")
    a = _canonical_projection_preflight("(x+1)**2")
    b = _canonical_projection_preflight("x**2 + 2*x + 1")
    assert a is not None and b is not None
    assert a[1] == b[1]


@pytest.mark.parametrize(
    "question",
    [
        # leading natural-language verb followed by an operand
        "simplify (x+1)**2",
        "factor x**2 - 1",
        "expand (a+b)**2",
        "evaluate sin(x)",
        # contains '?' — fails algebra charset
        "what is x+1?",
        # arithmetic still owns the no-letters route
        "0.1 + 0.2",
    ],
)
def test_preflight_algebra_falls_through_for_natural_language(question):
    """Natural-language phrasing falls through the algebra route. Pure
    arithmetic with no letters routes to arithmetic@v1, not algebra."""
    pytest.importorskip("sympy")
    r = _canonical_projection_preflight(question)
    if r is not None:
        # Acceptable: arithmetic route catches '0.1 + 0.2' first.
        assert r[0] == "arithmetic@v1"


@pytest.mark.parametrize(
    "question",
    [
        # natural language → no match
        "what is 0.1 + 0.2?",
        "tell me about the capital of france",
        "is A IMPL B the same as B IMPL A",
        # mixed case (lowercase letters) → no logic match
        "a IMPL b",
        # bare uppercase atom without operator → no logic match
        "A",
        "B",
        # empty / whitespace
        "",
        "   ",
    ],
)
def test_preflight_falls_through_for_natural_language(question):
    assert _canonical_projection_preflight(question) is None


def test_preflight_returns_none_on_pi_star_error():
    """Shape-matched but π*-rejected input falls through gracefully.

    ``arithmetic@v1`` rejects identifiers; a string that LOOKS
    arithmetic but contains nothing parseable should not crash query().
    """
    # Empty parens — passes the regex shape but arithmetic@v1 rejects.
    assert _canonical_projection_preflight("()") is None


# ----- query() short-circuit ----------------------------------------------


def test_query_math_short_circuits_with_no_shards(tmp_path: Path):
    """Math-shaped question persists a canonical row + answers without
    needing a shard or hitting the LLM. Status is cache_miss_then_written
    (first call writes the row); audit_mode/verifier_method are the
    canonical-projection tokens; cache_key is real."""
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"  # absent on purpose
    result = query(
        question="0.1 + 0.2",
        qa_db=qa_db,
        chat_client=StubClient(""),
        model_id="test/model",
        shards_dir=no_shards,
    )
    assert result["status"] == "cache_miss_then_written"
    assert result["audit_mode"] == "CANONICAL_PROJECTION"
    assert result["verifier_method"] == "canonical_projection"
    assert result["pi_star_ref"] == "arithmetic@v1"
    assert result["answer_text"] == "3/10"
    assert result["lookup_path"] == "canonical_cache_miss"
    assert result["sources"] == []
    assert result["cache_key"] is not None
    assert result["audit_event_hash"] is not None
    assert result["run_dag_root"] is not None


def test_query_logic_short_circuits(tmp_path: Path):
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"
    result = query(
        question="A IMPL B",
        qa_db=qa_db,
        chat_client=StubClient(""),
        model_id="test/model",
        shards_dir=no_shards,
    )
    assert result["status"] == "cache_miss_then_written"
    assert result["pi_star_ref"] == "logic-kernel@v1"
    assert result["answer_text"] == "(NOT A OR B)"


def test_query_canonical_transient_mode(tmp_path: Path):
    """Operators can disable persistence per-call via
    policy['canonical_projection_preflight_persist']=False; the legacy
    transient render-only behavior comes back (status='canonical_projection',
    cache_key=None, no audit-chain entry)."""
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"
    result = query(
        question="0.1 + 0.2",
        qa_db=qa_db,
        chat_client=StubClient(""),
        model_id="test/model",
        shards_dir=no_shards,
        policy={"canonical_projection_preflight_persist": False},
    )
    assert result["status"] == "canonical_projection"
    assert result["lookup_path"] == "preflight_canonical"
    assert result["cache_key"] is None


def test_query_contrapositive_collapses_to_same_canonical(tmp_path: Path):
    """A IMPL B and (NOT B) IMPL (NOT A) are the same equivalence class."""
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"
    a = query(
        question="A IMPL B",
        qa_db=qa_db,
        chat_client=StubClient(""),
        model_id="m",
        shards_dir=no_shards,
    )
    b = query(
        question="(NOT B) IMPL (NOT A)",
        qa_db=qa_db,
        chat_client=StubClient(""),
        model_id="m",
        shards_dir=no_shards,
    )
    assert a["answer_text"] == b["answer_text"] == "(NOT A OR B)"


def test_query_natural_language_falls_through_to_rag(tmp_path: Path):
    """Natural-language question doesn't trigger the preflight; goes
    to RAG and (with no corpus) returns no_sources."""
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"
    result = query(
        question="what is the capital of france?",
        qa_db=qa_db,
        chat_client=StubClient(""),
        model_id="m",
        shards_dir=no_shards,
    )
    # Falls through — preflight didn't match, RAG runs, finds no shards.
    assert result["status"] != "canonical_projection"
    assert result.get("audit_mode") != "CANONICAL_PROJECTION"


def test_query_canonical_preflight_disable_via_policy(tmp_path: Path):
    """policy['canonical_projection_preflight']=False disables the
    short-circuit even on math-shaped input."""
    qa_db = tmp_path / "qa.db"
    no_shards = tmp_path / "shards"
    result = query(
        question="0.1 + 0.2",
        qa_db=qa_db,
        chat_client=StubClient(""),
        model_id="m",
        shards_dir=no_shards,
        policy={"canonical_projection_preflight": False},
    )
    assert result["status"] != "canonical_projection"


# ----- `arborist canon` CLI subcommand ------------------------------------


def _run_canon(*args: str) -> subprocess.CompletedProcess:
    return subprocess.run(
        [sys.executable, "-m", "arborist.cli", "canon", *args],
        capture_output=True,
        text=True,
        check=False,
    )


def test_cli_canon_arithmetic():
    r = _run_canon("arithmetic@v1", "0.1 + 0.2")
    assert r.returncode == 0, r.stderr
    assert r.stdout.strip() == "3/10"


def test_cli_canon_logic():
    r = _run_canon("logic-kernel@v1", "A IMPL B")
    assert r.returncode == 0, r.stderr
    assert r.stdout.strip() == "(NOT A OR B)"


def test_cli_canon_time_series():
    r = _run_canon(
        "time-series-quantized@v1",
        '{"dt":1,"dv":0.1,"samples":[[0,1.0],[1,2.0]]}',
    )
    assert r.returncode == 0, r.stderr
    assert r.stdout.strip() == "dt=1;dv=0.1;n=2;t0=0:10|20"


def test_cli_canon_list():
    r = _run_canon("--list")
    assert r.returncode == 0, r.stderr
    keys = {line.split("\t")[0] for line in r.stdout.strip().split("\n")}
    assert "arithmetic@v1" in keys
    assert "logic-kernel@v1" in keys
    assert "time-series-quantized@v1" in keys
    assert "code-py-ast@v1" in keys
    assert "claim-lattice@v1" in keys
    assert "wikitext-base@v1" in keys


def test_cli_canon_json_emits_sha256():
    r = _run_canon("--json", "arithmetic@v1", "0.1 + 0.2")
    assert r.returncode == 0, r.stderr
    obj = json.loads(r.stdout)
    assert obj["pi_star_ref"] == "arithmetic@v1"
    assert obj["input"] == "0.1 + 0.2"
    assert obj["canonical"] == "3/10"
    assert len(obj["canonical_sha256"]) == 64


def test_cli_canon_unknown_key_returns_2():
    r = _run_canon("does-not-exist@v1", "anything")
    assert r.returncode == 2
    assert "unknown π*" in r.stderr


def test_cli_canon_pi_star_error_returns_1():
    r = _run_canon("arithmetic@v1", "x + 1")  # identifier rejected
    assert r.returncode == 1
    assert "error" in r.stderr.lower()


def test_cli_canon_missing_args_returns_2():
    r = _run_canon()
    assert r.returncode == 2
    assert "required" in r.stderr.lower()