"""Exact independent controls for the componentwise comparison exponential."""

from decimal import Decimal, localcontext
from fractions import Fraction as Q
from math import factorial

import pytest

from tnfr.mathematics._comparison_flow import comparison_flow_upper, scalar_affine_upper


def _diagonal(values):
    return tuple(
        tuple(value if i == j else Q(0) for j in range(len(values)))
        for i, value in enumerate(values)
    )


def _negative_exp_reference(value):
    """Independent alternating degree 96/97 bounds on exp(-value), 0<=value<=1."""
    assert 0 <= value <= 1
    upper = sum(((-value) ** k / factorial(k) for k in range(97)), Q(0))
    lower = upper - value**97 / factorial(97)
    return lower, upper


def test_zero_duration_keeps_exact_input_radii_with_outward_materialization():
    radii = (Q(1, 3), Q(0), Q(2, 7), Q(1, 1 << 200))
    result = comparison_flow_upper(_diagonal((-100, 0, 1, 8)), radii, Q(0))
    assert all(
        value <= bound <= value + Q(1, 1 << 128) for value, bound in zip(radii, result)
    )
    assert result[1] == 0 and result[3] > 0


def test_zero_uncertainty_is_preserved_exactly():
    matrix = tuple(
        tuple(Q(-1, 2) if i == j else Q(1, 8) for j in range(4)) for i in range(4)
    )
    assert comparison_flow_upper(matrix, (0, 0, 0, 0), Q(1)) == (0, 0, 0, 0)


@pytest.mark.parametrize("duration", (Q(0), Q(1, 3)))
def test_fixed_width_fraction_components_preserve_nilpotent_flow(duration):
    np = pytest.importorskip("numpy")

    def payload(value):
        return Q(np.int64(value.numerator), np.int64(value.denominator))

    matrix = ((Q(0), Q(1, 4)), (Q(0), Q(0)))
    radii = (Q(1, 7), Q(2**62, 3))
    result = comparison_flow_upper(
        tuple(tuple(payload(value) for value in row) for row in matrix),
        tuple(payload(value) for value in radii),
        payload(duration),
    )
    # M**2=0, hence exp(h*M)r = r+h*M*r, independently of the series kernel.
    expected = (radii[0] + duration * radii[1] / 4, radii[1])
    assert all(
        value <= upper < value + Q(1, 10**15) for value, upper in zip(expected, result)
    )
    assert result == comparison_flow_upper(matrix, radii, duration)
    assert all(
        type(value.numerator) is type(value.denominator) is int for value in result
    )


def test_negative_diagonal_retains_contraction_instead_of_absolute_growth():
    radii = (Q(1), Q(2), Q(3), Q(4))
    result = comparison_flow_upper(_diagonal((-1,) * 4), radii, Q(1))
    _, exp_upper = _negative_exp_reference(Q(1))
    for radius, bound in zip(radii, result):
        assert radius * exp_upper <= bound < radius
        assert bound - radius * exp_upper < Q(1, 10**35)


def test_shifted_nilpotent_chain_matches_finite_polynomial_times_decay():
    matrix = tuple(
        tuple(Q(-1) if i == j else Q(int(j == i + 1)) for j in range(4))
        for i in range(4)
    )
    radii, h = (Q(1), Q(2), Q(0), Q(4)), Q(2, 3)
    result = comparison_flow_upper(matrix, radii, h)
    _, decay_upper = _negative_exp_reference(h)
    for i, bound in enumerate(result):
        exact_polynomial = sum(
            (h**k * radii[i + k] / factorial(k) for k in range(4 - i)), Q(0)
        )
        reference = exact_polynomial * decay_upper
        assert reference <= bound < reference + Q(1, 10**34)


def test_positive_diagonal_at_unit_work_limit_encloses_independent_exp_series():
    result = comparison_flow_upper(_diagonal((1,) * 4), (1, 2, 3, 4), Q(1))
    reference = sum((Q(1, factorial(k)) for k in range(101)), Q(0))
    reference_upper = reference + Q(1, factorial(101)) / (1 - Q(1, 102))
    for radius, bound in zip((1, 2, 3, 4), result):
        assert radius * reference_upper <= bound
        assert bound - radius * reference_upper < Q(1, 10**34)


def test_dense_metzler_matrix_matches_independent_projector_decomposition():
    # P=ones/4 is an exact projector. exp(h*(P-2Id)) is therefore
    # exp(-2h)*Id + [exp(-h)-exp(-2h)]*P, including nonuniform input.
    matrix = tuple(
        tuple(Q(1, 4) - (2 if i == j else 0) for j in range(4)) for i in range(4)
    )
    radii = (Q(0), Q(1), Q(2), Q(3))
    result = comparison_flow_upper(matrix, radii, Q(1, 2))
    low_one, high_one = _negative_exp_reference(Q(1))
    _, high_half = _negative_exp_reference(Q(1, 2))
    mean = sum(radii) / 4
    for radius, bound in zip(radii, result):
        coefficient = radius - mean
        reference_upper = mean * high_half + coefficient * (
            high_one if coefficient >= 0 else low_one
        )
        assert reference_upper <= bound < reference_upper + Q(1, 10**34)


def test_comparison_encloses_exact_nonlinear_flow_separation_on_convex_tube():
    # xdot=-x*x has exact solution x/(1+h*x). At h=1/4 both solutions
    # from x in [1,2] and center 3/2 remain in [1/2,2], where J=-2*x<=-1.
    # This supplies an independent dynamical use of the signed comparison.
    h, center = Q(1, 4), Q(3, 2)
    bound = comparison_flow_upper(_diagonal((-1,) * 4), (Q(1, 2),) * 4, h)[0]
    center_endpoint = center / (1 + h * center)
    for initial in (Q(1), Q(2)):
        assert abs(initial / (1 + h * initial) - center_endpoint) <= bound
    assert bound < Q(1, 2)


def test_eight_coordinate_diagonal_matches_independent_scalar_exponentials():
    diagonal = tuple(Q(-index, 8) for index in range(1, 9))
    radii = tuple(Q(index, 3) for index in range(8))
    h = Q(1, 2)
    result = comparison_flow_upper(_diagonal(diagonal), radii, h)
    for rate, radius, bound in zip(diagonal, radii, result):
        _, decay_upper = _negative_exp_reference(-rate * h)
        reference = radius * decay_upper
        assert reference <= bound < reference + Q(1, 10**34)


def test_eight_coordinate_blocks_match_independent_nilpotent_exponentials():
    # Each 2x2 block is -a*Id+N, N^2=0, so its exact response is
    # exp(-a*h)*(Id+h*N). Blocks have distinct decay and coupling rates.
    decay = tuple(Q(index, 4) for index in range(1, 5))
    coupling = tuple(Q(index, 16) for index in range(1, 5))
    matrix = tuple(
        tuple(
            (
                -decay[i // 2]
                if i == j
                else coupling[i // 2] if i % 2 == 0 and j == i + 1 else Q(0)
            )
            for j in range(8)
        )
        for i in range(8)
    )
    radii, h = (Q(0), Q(2), Q(1), Q(3), Q(4), Q(0), Q(2), Q(5)), Q(1, 2)
    result = comparison_flow_upper(matrix, radii, h)
    for index, bound in enumerate(result):
        _, decay_upper = _negative_exp_reference(decay[index // 2] * h)
        polynomial = radii[index]
        if index % 2 == 0:
            polynomial += h * coupling[index // 2] * radii[index + 1]
        reference = polynomial * decay_upper
        assert reference <= bound < reference + Q(1, 10**34)


@pytest.mark.parametrize("dimension", (1, 3, 16, 23, 24))
def test_admitted_dimension_boundaries_retain_contraction(dimension):
    result = comparison_flow_upper(
        _diagonal((-1,) * dimension), (1,) * dimension, Q(1, 2)
    )
    _, decay_upper = _negative_exp_reference(Q(1, 2))
    assert len(result) == dimension
    assert all(decay_upper <= bound < 1 for bound in result)


@pytest.mark.parametrize(
    "matrix,radii,duration,message",
    (
        ((), (), Q(1), "dimension.*1 and 24"),
        (_diagonal((0,) * 25), (0,) * 25, Q(1), "dimension.*1 and 24"),
        (((0,) * 4,) * 3, (0,) * 4, Q(1), "square"),
        (((0,) * 3,) * 4, (0,) * 4, Q(1), "square"),
        (_diagonal((0,) * 4), (0,) * 3, Q(1), "match the matrix dimension"),
        (
            ((0, -1, 0, 0), (0,) * 4, (0,) * 4, (0,) * 4),
            (1,) * 4,
            Q(1),
            "off-diagonals",
        ),
        (_diagonal((0,) * 4), (-1, 0, 0, 0), Q(1), "nonnegative"),
        (_diagonal((0,) * 4), (1,) * 4, Q(-1), "nonnegative"),
        (_diagonal((2,) * 4), (1,) * 4, Q(1), "norm"),
        (_diagonal((-2,) * 4), (1,) * 4, Q(1), "alpha"),
    ),
)
def test_invalid_comparison_or_unresolved_numerical_domain_rejects(
    matrix, radii, duration, message
):
    with pytest.raises(ValueError, match=message):
        comparison_flow_upper(matrix, radii, duration)


@pytest.mark.parametrize("invalid", (False, 0.0, "0"))
def test_raw_nonrational_coordinates_are_not_silently_materialized(invalid):
    with pytest.raises(TypeError, match="exact Fraction"):
        comparison_flow_upper(_diagonal((0,) * 4), (0, 0, 0, invalid), Q(1))
    with pytest.raises(TypeError, match="exact Fraction"):
        comparison_flow_upper(_diagonal((invalid, 0, 0, 0)), (0,) * 4, Q(1))
    with pytest.raises(TypeError, match="exact Fraction"):
        comparison_flow_upper(_diagonal((0,) * 4), (0,) * 4, invalid)


def _decimal(value):
    return Decimal(value.numerator) / Decimal(value.denominator)


@pytest.mark.parametrize("rate", [Q(-2), Q(-1, 3), Q(0), Q(1, 3), Q(2)])
def test_scalar_growth_contains_independent_high_precision_solution(rate):
    radius, duration, defect = Q(1, 7), Q(1, 8), Q(11, 10**6)
    upper = scalar_affine_upper(radius, rate, duration, defect)
    with localcontext() as context:
        context.prec = 100
        exponential = (_decimal(rate) * _decimal(duration)).exp()
        integral = (exponential - 1) / _decimal(rate) if rate else _decimal(duration)
        expected = exponential * _decimal(radius) + integral * _decimal(defect)
        assert _decimal(upper) >= expected
        assert _decimal(upper) - expected < Decimal("1e-30")


@pytest.mark.parametrize("sign", (-1, 1))
@pytest.mark.parametrize("defect", (Q(0), Q(11, 10**6)))
@pytest.mark.parametrize(
    "magnitude,duration",
    (
        (Q(1, 10**100), Q(1)),
        (Q(1, 1 << 129), Q(1)),
        (Q(1, 10**100), Q(10**100)),
    ),
)
def test_scalar_subgrid_nonzero_rates_enclose_independent_decimal_solution(
    sign, magnitude, duration, defect
):
    rate, radius = sign * magnitude, Q(1, 7)
    upper = scalar_affine_upper(radius, rate, duration, defect)
    with localcontext() as context:
        # Resolve exp(1e-100)-1 independently, even before dividing by rate.
        context.prec = 240
        exponential = (_decimal(rate) * _decimal(duration)).exp()
        integral = (exponential - 1) / _decimal(rate)
        expected = exponential * _decimal(radius) + integral * _decimal(defect)
        assert _decimal(upper) >= expected
        assert _decimal(upper) - expected < max(Decimal(1), expected) * Decimal("1e-30")


@pytest.mark.parametrize("sign", (-1, 1))
def test_scalar_subgrid_rate_contains_independent_affine_taylor_bound(sign):
    rate = sign * Q(1, 10**100)
    upper = scalar_affine_upper(0, rate, 1, 1)
    # The integral of exp(rate*s) on [0,1] has this quadratic Taylor upper
    # bound for decay. For growth, use exp(z) <= 1+z+z*z on [0,1/2].
    reference_upper = 1 + rate / 2 + rate**2 / (6 if sign < 0 else 3)
    assert reference_upper <= upper < reference_upper + Q(1, 10**30)


@pytest.mark.parametrize(
    "arguments",
    [
        (True, Q(0), Q(1), Q(0)),
        (Q(-1), Q(0), Q(1), Q(0)),
        (Q(0), Q(0), Q(-1), Q(0)),
        (Q(0), Q(0), Q(1), Q(-1)),
        (Q(0), Q(2), Q(1), Q(0)),
        (Q(0), float("nan"), Q(1), Q(0)),
        (Q(0), Q(0), float("inf"), Q(0)),
        (Q(0), True, Q(0), Q(0)),
        (Q(0), float("nan"), Q(0), Q(0)),
        (Q(0), Q(1, 10**100), Q(0), float("nan")),
        (Q(0), Q(1, 10**100), Q(0), Q(-1)),
    ],
)
def test_scalar_growth_rejects_invalid_primitive_domains(arguments):
    with pytest.raises((TypeError, ValueError)):
        scalar_affine_upper(*arguments)


def test_zero_duration_retains_initial_uncertainty():
    upper = scalar_affine_upper(Q(1, 3), Q(-1), Q(0), Q(1))
    assert Q(1, 3) <= upper < Q(1, 3) + Q(1, 10**30)


@pytest.mark.parametrize("rate", (Q(-1, 10**100), Q(0), Q(1, 10**100)))
@pytest.mark.parametrize("radius", (Q(0), Q(1, 3), Q(1, 10**100)))
def test_zero_duration_subgrid_rate_retains_outward_initial_uncertainty(rate, radius):
    upper = scalar_affine_upper(radius, rate, Q(0), Q(1))
    assert radius <= upper < radius + Q(1, 1 << 128)
    assert (upper == 0) == (radius == 0)