arborist/tests/test_pi_star_calculus_derivative.py

"""calculus-derivative@v1 π* tests (ticket #000030 Phase 2).

Covers JSON-shaped input contract, derivative correctness across
common shapes, n-th order, multi-variable handling, error paths,
and round-trip semantics. Whole file skips when SymPy is absent.
"""

from __future__ import annotations

import json

import pytest

sympy = pytest.importorskip("sympy")  # entire file skipped if absent

from arborist.pi_star import get  # noqa: E402
from arborist.pi_star.protocol import PiStarError  # noqa: E402


@pytest.fixture
def ps():
    return get("calculus-derivative@v1")


@pytest.fixture
def algebra():
    return get("algebra-symbolic@v1")


def _q(**kwargs) -> bytes:
    return json.dumps(kwargs).encode("utf-8")


# --- registry presence ------------------------------------------------


def test_registry_contains_calculus_derivative():
    ps = get("calculus-derivative@v1")
    assert ps.name == "calculus-derivative"
    assert ps.version == "v1"
    assert ps.domain == "calculus"


# --- correctness via cross-check against algebra-symbolic@v1 ---------


@pytest.mark.parametrize(
    "f,x,expected",
    [
        ("x**2", "x", "2*x"),
        ("x**3", "x", "3*x**2"),
        ("x", "x", "1"),
        ("5", "x", "0"),
        ("a*x + b", "x", "a"),
        ("x*y", "x", "y"),  # treats y as constant w.r.t. x
        ("x**2 + 3*x + 7", "x", "2*x + 3"),
        ("(x+1)**2", "x", "2*x + 2"),
    ],
)
def test_first_derivative_matches_expected(ps, algebra, f, x, expected):
    got = ps.canonicalize(_q(f=f, x=x))
    want = algebra.canonicalize(expected.encode())
    assert got == want


@pytest.mark.parametrize(
    "f,x,n,expected",
    [
        ("x**3", "x", 2, "6*x"),
        ("x**4", "x", 3, "24*x"),
        ("x**2", "x", 0, "x**2"),  # n=0 returns the original (expanded)
        ("x**2", "x", 5, "0"),
    ],
)
def test_nth_derivative(ps, algebra, f, x, n, expected):
    got = ps.canonicalize(_q(f=f, x=x, n=n))
    want = algebra.canonicalize(expected.encode())
    assert got == want


def test_trig_derivative(ps, algebra):
    """sin/cos derivatives go through sympy's symbolic engine."""
    out = ps.canonicalize(_q(f="sin(x)", x="x")).decode()
    assert "cos" in out and "Symbol('x')" in out


# --- equivalence classes preserved -----------------------------------


def test_equivalent_inputs_collapse(ps):
    """d/dx of equivalent expressions produces equivalent results."""
    a = ps.canonicalize(_q(f="(x+1)**2", x="x"))
    b = ps.canonicalize(_q(f="x**2 + 2*x + 1", x="x"))
    assert a == b


# --- error paths ------------------------------------------------------


def test_non_bytes_input_raises(ps):
    with pytest.raises(PiStarError, match="expects bytes"):
        ps.canonicalize("not bytes")  # type: ignore[arg-type]


def test_non_json_input_raises(ps):
    with pytest.raises(PiStarError, match="JSON"):
        ps.canonicalize(b"not json at all")


def test_non_object_json_raises(ps):
    with pytest.raises(PiStarError, match="JSON object"):
        ps.canonicalize(b'["x**2", "x"]')
    with pytest.raises(PiStarError, match="JSON object"):
        ps.canonicalize(b'"just a string"')


def test_missing_f_raises(ps):
    with pytest.raises(PiStarError, match="missing field 'f'"):
        ps.canonicalize(_q(x="x"))


def test_missing_x_raises(ps):
    with pytest.raises(PiStarError, match="missing field 'x'"):
        ps.canonicalize(_q(f="x**2"))


def test_empty_field_rejected(ps):
    with pytest.raises(PiStarError, match="non-empty"):
        ps.canonicalize(_q(f="", x="x"))
    with pytest.raises(PiStarError, match="non-empty"):
        ps.canonicalize(_q(f="x", x=""))


def test_non_string_field_rejected(ps):
    with pytest.raises(PiStarError, match="non-empty string"):
        ps.canonicalize(_q(f=42, x="x"))


@pytest.mark.parametrize("bad_n", [-1, 1.5, "1", True, False])
def test_n_must_be_non_negative_int(ps, bad_n):
    with pytest.raises(PiStarError, match="non-negative integer"):
        ps.canonicalize(_q(f="x**2", x="x", n=bad_n))


def test_unparseable_f_raises(ps):
    with pytest.raises(PiStarError, match="cannot parse"):
        ps.canonicalize(_q(f="$$$ % nonsense", x="x"))


def test_relational_f_rejected(ps):
    with pytest.raises(PiStarError, match="boolean/relational"):
        ps.canonicalize(_q(f="x > 0", x="x"))


# --- output is itself an algebra-symbolic@v1 valid input -------------


def test_output_canonicalizes_under_algebra_symbolic(ps, algebra):
    """The derivative output is an expression in srepr form. Feeding
    it back into algebra-symbolic@v1 should be idempotent (it parses
    srepr S-expressions natively)."""
    raw = ps.canonicalize(_q(f="(x+1)**2", x="x"))
    once_through_algebra = algebra.canonicalize(raw)
    twice_through_algebra = algebra.canonicalize(once_through_algebra)
    assert once_through_algebra == twice_through_algebra