"""Contract, numerical, and snapshot tests for the shared computation store."""
import json

import numpy as np
import pytest
import sympy as sp

from autodiff_visualizer import ComputationGraph, safe_parse_expr
from backward_rules import ScalarLocalRule, TensorBackwardRule
from diagram_renderer import DIVIDER, render_diagram, step_label
from math_renderer import MathRenderer
from tensor_engine import TensorComputationGraph
from trace_renderer import _contribution_box_html, render_trace, render_trace_html


def scalar(expr, feed):
    graph = ComputationGraph().build(safe_parse_expr(expr))
    graph.evaluate({sp.Symbol(k): v for k, v in feed.items()})
    return graph.to_computation_model()


def tensor_graph(expr, feed, forward_only=False):
    graph = TensorComputationGraph().build_from_ast(expr)
    graph.evaluate(feed, forward_only=forward_only)
    return graph


CASES = [
    ('scalar', '(x*x)*y+(y+2)', {'x': 3., 'y': 4.}),
    ('scalar', 'sin(x) + x**y', {'x': 2., 'y': 3.}),
    ('tensor', 'sum(x/y+b)', {'x': np.array([2., 4.]), 'y': np.array([1., 2.]), 'b': 2.}),
    ('tensor', 'sum(W @ x)', {'W': np.array([[1., 2.], [3., 4.]]), 'x': np.array([2., 3.])}),
    ('tensor', 'sum(x @ W)', {'W': np.array([[1., 2., 3.], [3., 4., 5.]]), 'x': np.array([2., 3.])}),
    ('tensor', 'softmax(x)[-1]', {'x': np.array([1., 2., 3.])}),
    ('tensor', 'sum(x*x)', {'x': np.array([1., 2., 3.])}),
    ('tensor', 'sum(x*x)', {'x': 3.}),
]


@pytest.mark.parametrize('notation', ['code', 'mathbook'])
@pytest.mark.parametrize('engine,formula,feed', CASES)
def test_every_view_uses_identical_local_equations(notation, engine, formula, feed):
    model = scalar(formula, feed) if engine == 'scalar' else tensor_graph(formula, feed).to_computation_model()
    text, html, diagram = render_trace(model, notation), render_trace_html(model, notation), render_diagram(model, notation)
    for step in model.steps:
        renderer = MathRenderer(notation, html=True)
        upper = renderer.backward_rule(step.explanation)
        lower = renderer.application_derivation(step.explanation, layout="auto")
        plain_upper, plain_lower = MathRenderer(notation).explanation(step.explanation)
        assert step_label(step, notation).split(DIVIDER)[0] == f"<div class='backward-rule'>{upper}</div>"
        assert f'>{lower}</div>' in step_label(step, notation)
        box = _contribution_box_html(step.operation, step, notation)
        assert f'class="tv-deriv-rule">{upper}</div>' in box
        assert f'class="tv-deriv-result">{lower}</div>' in box
        assert box in html
        assert upper in diagram and lower in diagram
        assert plain_upper in text and plain_lower in text
        assert 'diff(out,' not in upper and '∂out/' not in upper
        # Semantic identity: the lower compartment never re-states the symbolic rule.
        assert step.rule.application not in step.explanation.application_derivation.stages
        assert step.explanation.application_derivation.lhs == step.explanation.contribution
        # Expressions are data, independent of both display notation and containers.
        assert '<span' not in json.dumps(step.rule.application.to_data())


def test_reference_scalar_example_and_occurrence_identity():
    model = scalar('(x*x)*y+(y+2)', {'x': 3., 'y': 4.})
    x_steps = [s for s in model.steps if s.destination_quantity.name == 'x']
    renderer = MathRenderer()
    assert renderer.explanation(x_steps[0].explanation) == (
        'a = x, b = x\n∂(a · b)/∂a = b = x = 3',
        'c_n1→x, 1 = 4 · 3 = 12')
    assert renderer.explanation(x_steps[1].explanation) == (
        'a = x, b = x\n∂(a · b)/∂b = a = x = 3',
        'c_n1→x, 2 = 4 · 3 = 12')
    assert x_steps[0].rule.operand.key != x_steps[1].rule.operand.key
    assert x_steps[0].destination_quantity is x_steps[1].destination_quantity
    assert '∂(x · x)/∂x' not in render_trace(model)
    assert 'Adjoint:\nx̄ = ∂out/∂x = 24' in render_trace(model)
    for step in x_steps:
        assert renderer.expression(step.explanation.contribution) in renderer.explanation(step.explanation)[1]
        assert MathRenderer(html=True).expression(step.explanation.contribution) in render_trace_html(model)


def test_slot_names_do_not_confuse_graph_quantity_names():
    model = scalar('a*a', {'a': 3.})
    assert MathRenderer().explanation(model.steps[0].explanation)[0] == 'b = a, c = a\n∂(b · c)/∂b = c = a = 3'


def test_contract_depends_on_local_shapes():
    scalar_model = scalar('x*x', {'x': 3.})
    tensor_scalar = tensor_graph('x*x', {'x': 3.}).to_computation_model()
    assert all(isinstance(s.rule, ScalarLocalRule) for s in tensor_scalar.steps)
    for first, second in zip(scalar_model.steps, tensor_scalar.steps):
        assert MathRenderer().explanation(first.explanation) == MathRenderer().explanation(second.explanation)
    dot = tensor_graph('dot(x,x)', {'x': np.array([2., 3.])}).to_computation_model()
    assert all(isinstance(s.rule, TensorBackwardRule) for s in dot.steps)
    assert all(s.contribution.shape == (2,) for s in dot.steps)


@pytest.mark.parametrize('formula,feed', [
    ('sum(W @ x)', {'W': np.array([[1., 2., 3.], [4., 5., 6.]]), 'x': np.array([.2, .3, .4])}),
    ('sum(x @ W)', {'W': np.array([[1., 2., 3.], [4., 5., 6.]]), 'x': np.array([.2, .3])}),
    ('mean(A @ B)', {'A': np.array([[1., 2.], [3., 4.]]), 'B': np.array([[.2, .3, .4], [.5, .6, .7]])}),
    ('x @ y', {'x': np.array([.2, .3]), 'y': np.array([.4, .5])}),
    ('sum(x/b+b/x)', {'x': np.array([2., 3.]), 'b': 2.}),
    ('sum(tanh(x)*sigmoid(x)+exp(x)+log(x))', {'x': np.array([.2, .3])}),
    ('sum(x**-2)', {'x': np.array([2., 3.])}),
    ('softmax(x)[-1]', {'x': np.array([.2, .3, .4])}),
])
def test_refactored_tensor_rules_against_finite_differences(formula, feed):
    graph = tensor_graph(formula, feed)
    for name, original in feed.items():
        value = np.asarray(original)
        for index in np.ndindex(value.shape):
            plus = {k: np.array(v, copy=True) for k, v in feed.items()}
            minus = {k: np.array(v, copy=True) for k, v in feed.items()}
            plus[name][index] += 1e-6
            minus[name][index] -= 1e-6
            numerical = (tensor_graph(formula, plus, True).root.val - tensor_graph(formula, minus, True).root.val) / 2e-6
            assert np.isclose(np.asarray(graph.name_to_input[name].grad)[index], numerical, atol=1e-7)
    model = graph.to_computation_model()
    values = {q.uid: q.primal_value for q in model.quantities}
    for step in model.steps:
        env = {**values, 'adjoint:' + step.source_quantity.uid: step.incoming_adjoint}
        actual = step.rule.application.evaluate(env)
        assert np.shape(actual) == step.destination_quantity.shape
        np.testing.assert_allclose(actual, step.contribution)


def test_snapshot_is_independent_of_feed_and_reevaluation():
    feed = {'x': np.array([2., 3.])}
    graph = tensor_graph('sum(x*x)', feed)
    model = graph.to_computation_model()
    before = render_trace(model)
    feed['x'][0] = 100.
    graph.evaluate({'x': np.array([10., 20.])})
    assert render_trace(model) == before
    x = next(q for q in model.quantities if q.name == 'x')
    np.testing.assert_array_equal(x.primal_value, [2., 3.])
    with pytest.raises(ValueError):
        x.primal_value[0] = 99.


def test_rendering_does_not_evaluate_rules(monkeypatch):
    from math_expression import Expr
    model = tensor_graph('sum(W @ x)', {'W': np.ones((2, 3)), 'x': np.ones(3)}).to_computation_model()
    def forbidden(*args, **kwargs):
        raise AssertionError('Renderers must consume stored explanations')
    monkeypatch.setattr(Expr, 'evaluate', forbidden)
    render_diagram(model)
    render_trace(model)
    render_trace_html(model)


def test_forward_only_html_does_not_construct_backward_examples(monkeypatch):
    from teaching_examples import multiplication_example
    multiplication_example.cache_clear()
    graph = ComputationGraph().build(safe_parse_expr('x*x'), forward_only=True)
    graph.evaluate({sp.Symbol('x'): 3.}, forward_only=True)
    def forbidden(*args, **kwargs):
        raise AssertionError('Forward-only rendering must not differentiate examples')
    monkeypatch.setattr(sp, 'diff', forbidden)
    html = render_trace_html(graph.to_computation_model(forward_only=True), mode='forward')
    body = html.split('</style>')[1]
    assert '∂' not in body and 'diff(' not in body
    assert 'tv-adjoint-step' not in body and 'tv-deriv-box' not in body


def test_repeated_base_and_exponent_keep_different_rules():
    model = scalar('x**x', {'x': 2.})
    first, second = model.steps
    assert first.rule.operand.key != second.rule.operand.key
    assert first.local_value == 4.
    assert np.isclose(second.local_value, 4 * np.log(2))
    assert 'log(' in MathRenderer('code').explanation(second.explanation)[0]
    assert np.isclose(model.accumulations[0].resulting_adjoint, 4 + 4 * np.log(2))


def test_numerical_outer_product_is_not_rendered_as_row_times_column():
    model = tensor_graph('sum(W @ x)', {'W': np.ones((2, 2)), 'x': np.array([2., 3.])}).to_computation_model()
    upper, lower = MathRenderer().explanation(model.steps[0].explanation)
    assert upper == 'c_n1→W = n̄1 xᵀ'
    assert lower == 'c_n1→W = outer([1, 1], [2, 3])\n       = [[2, 3], [2, 3]]'


@pytest.mark.parametrize('value', [2., np.array([0., 2.])])
def test_constant_zero_power_does_not_multiply_infinite_adjoint(value):
    count = np.size(value)
    graph = tensor_graph(f'log(sum(x**0)-{count})', {'x': value})
    np.testing.assert_array_equal(graph.name_to_input['x'].grad, np.zeros_like(value))
    model = graph.to_computation_model()
    step = next(s for s in model.steps if s.operation.kind == 'pow')
    assert not step.rule.application.has_adjoint()
    assert 'nan' not in MathRenderer().explanation(step.explanation)[1]


@pytest.mark.parametrize('engine', ['scalar', 'tensor'])
@pytest.mark.parametrize('formula, expected', [
    ('nograd(x*x)', None),
    ('nograd(x*x)+x', 1.),
    ('x*x', 0.),
])
def test_unreached_adjoints_are_absent_but_active_zeros_remain(engine, formula, expected):
    from diagram_renderer import quantity_label
    from tikz_renderer import render_tikz
    feed = {'x': 0. if formula == 'x*x' else 3.}
    model = (scalar(formula, feed) if engine == 'scalar' else
             tensor_graph(formula, feed).to_computation_model())
    x = next(q for q in model.quantities if q.name == 'x')
    assert x.adjoint_value == expected
    operations = {op.quantity.uid: op for op in model.operations}
    if 'nograd' in formula:
        stopped_operand = next(op.operands[0] for op in model.operations if op.function == 'nograd')
        assert stopped_operand.adjoint_value is None
        stop_result = next(op.quantity for op in model.operations if op.function == 'nograd')
        assert stop_result.adjoint_value == 1
        assert 'adjoint-state' not in quantity_label(stopped_operand, operations[stopped_operand.uid], 'code')
    label = quantity_label(x, None, 'code')
    assert ('adjoint-state' in label) == (expected is not None)
    tikz = render_tikz(model)
    assert (r'\overline{x}=' in tikz) == (expected is not None)
    text = render_trace(model, 'code')
    assert ('diff(out, x) = x_bar =' in text) == (expected is not None)
    html = render_trace_html(model, 'code')
    assert 'None' not in html


def test_semantic_derivation_stages_model_contract():
    """Lower compartment = (concrete application, result); never the symbolic rule."""
    def check(step, concrete_kind):
        d = step.explanation.application_derivation
        assert d.lhs == step.explanation.contribution
        assert len(d.stages) == 2
        assert d.stages[0].kind == concrete_kind
        assert d.stages[1].kind == 'literal'
        # Identity, not substring: the symbolic rule lives only in local_rule.
        assert step.rule.application not in d.stages
        assert not set(step.explanation.local_rule.terms[1:]) & set(d.stages)
        return d.stages[1].value

    # 1. Scalar multiplication: c_out→x = 1 · 3 = 3
    assert check(scalar('x*y', {'x': 2., 'y': 3.}).steps[0], 'mul') == 3.0
    # 2. Repeated scalar operand: x*x
    for s in scalar('x*x', {'x': 3.}).steps:
        assert check(s, 'mul') == 3.0
    # 3. softmax(z)[1]: softmax VJP, then index selection
    tg_sm = tensor_graph('softmax(z)[1]', {'z': np.array([1., 2., 3.])}).to_computation_model()
    np.testing.assert_allclose(check(tg_sm.steps[0], 'elementwise'),
                               [-0.02203304, 0.18483645, -0.1628034], atol=1e-5)
    np.testing.assert_allclose(check(tg_sm.steps[1], 'scale'), [0., 1., 0.])
    # 4. sum(W @ x): outer product to W, transpose-matmul to x
    tg_mv = tensor_graph('sum(W @ x)', {'W': np.ones((2, 3)), 'x': np.ones(3)}).to_computation_model()
    np.testing.assert_array_equal(check(tg_mv.steps[0], 'outer'), np.ones((2, 3)))
    np.testing.assert_array_equal(check(tg_mv.steps[1], 'matmul'), [2., 2., 2.])
    # 5. Broadcasting: sum(x/y + b)
    tg_bc = tensor_graph('sum(x/y + b)', {'x': np.array([2., 4.]), 'y': np.array([1., 2.]), 'b': 2.}).to_computation_model()
    assert check(next(s for s in tg_bc.steps if s.destination_quantity.name == 'b'), 'sum') == 2.0


def test_derivation_stages_forbid_embedded_equality():
    """Negative invariant: No derivation stage may itself be an Equation or encode internal equality."""
    from backward_explanation import Equation
    for engine, formula, feed in CASES:
        model = scalar(formula, feed) if engine == 'scalar' else tensor_graph(formula, feed).to_computation_model()
        for step in model.steps:
            d = step.explanation.application_derivation
            # lhs must be an Expr, not Equation
            assert not isinstance(d.lhs, Equation)
            assert d.lhs.kind != 'equation'
            for stage in d.stages:
                # stages must be Expr, never Equation
                assert not isinstance(stage, Equation)
                assert stage.kind != 'equation'