"""Reusable interval certificate contracts.

Every failed proof obligation raises :class:`CertificationError`.  None of the
checks in this module rely on ``assert``, so running Python with ``-O`` cannot
silently remove a certificate condition.
"""

from __future__ import annotations

from collections.abc import Callable, Hashable, Iterable, Sequence
from dataclasses import dataclass
from enum import Enum
from typing import Generic, TypeVar

from flint import arb, arb_mat


class CertificationError(RuntimeError):
    """A requested rigorous certificate could not be established."""


class Sign(str, Enum):
    NEGATIVE = "negative"
    ZERO = "zero"
    POSITIVE = "positive"
    UNRESOLVED = "unresolved"


@dataclass(frozen=True)
class LDLProof:
    order: int
    pivots: tuple[arb, ...]


@dataclass(frozen=True)
class CongruenceCertificate:
    order: int
    preconditioner_determinant: arb
    determinant: arb
    ldl: LDLProof


@dataclass(frozen=True)
class NormalizationRecord:
    name: str
    determinant: arb
    divisor: arb
    normalized_determinant: arb


@dataclass(frozen=True)
class MonicExtremalQuadraticFormCertificate:
    order: int
    vector: tuple[arb, ...]
    pivot: arb
    quadratic_form: arb
    normal_residuals: tuple[arb, ...]
    leading_ldl: LDLProof


Key = TypeVar("Key", bound=Hashable)


@dataclass(frozen=True)
class UniqueArgminCertificate(Generic[Key]):
    winner: Key
    value: arb
    comparisons: int


Point = TypeVar("Point")


@dataclass(frozen=True)
class SeriesCrossCheck(Generic[Point]):
    point: Point
    series_value: arb
    scalar_value: arb


@dataclass(frozen=True)
class SeriesCrossCheckRecord(Generic[Point]):
    name: str
    checks: tuple[SeriesCrossCheck[Point], ...]


def classify_sign(value: arb) -> Sign:
    """Classify a real Arb interval without mistaking overlap for zero."""
    if value > 0:
        return Sign.POSITIVE
    if value < 0:
        return Sign.NEGATIVE
    if value == 0:
        return Sign.ZERO
    return Sign.UNRESOLVED


def require_positive(value: arb, label: str) -> arb:
    sign = classify_sign(value)
    if sign is not Sign.POSITIVE:
        raise CertificationError(f"{label} is {sign.value}: {value}")
    return value


def require_nonzero(value: arb, label: str) -> arb:
    sign = classify_sign(value)
    if sign not in (Sign.POSITIVE, Sign.NEGATIVE):
        raise CertificationError(f"{label} is {sign.value}: {value}")
    return value


def _matrix_entries(
    matrix: arb_mat | Sequence[Sequence[arb]], label: str
) -> list[list[arb]]:
    if isinstance(matrix, arb_mat):
        rows = matrix.nrows()
        columns = matrix.ncols()
        entries = [[matrix[row, column] for column in range(columns)] for row in range(rows)]
    else:
        entries = [list(row) for row in matrix]
        rows = len(entries)
        columns = len(entries[0]) if entries else 0
        if any(len(row) != columns for row in entries):
            raise CertificationError(f"{label} has ragged rows")
    if rows == 0 or rows != columns:
        raise CertificationError(f"{label} must be a nonempty square matrix, got {rows}x{columns}")
    return entries


def _validated_symmetric_entries(
    matrix: arb_mat | Sequence[Sequence[arb]], label: str
) -> list[list[arb]]:
    entries = _matrix_entries(matrix, label)
    for row in range(len(entries)):
        for column in range(row + 1, len(entries)):
            left = entries[row][column]
            right = entries[column][row]
            if not left.overlaps(right):
                raise CertificationError(
                    f"{label} is not symmetric at ({row}, {column}): {left} versus {right}"
                )
            common = left.intersection(right)
            entries[row][column] = common
            entries[column][row] = common
    return entries


def certify_positive_ldl(
    matrix: arb_mat | Sequence[Sequence[arb]], *, label: str = "matrix"
) -> LDLProof:
    """Certify positive definiteness using interval ``LDL^T`` pivots."""
    work = _validated_symmetric_entries(matrix, label)
    pivots: list[arb] = []
    for pivot_index in range(len(work)):
        pivot = require_positive(
            work[pivot_index][pivot_index], f"{label} LDL pivot {pivot_index}"
        )
        pivots.append(pivot)
        for row in range(pivot_index + 1, len(work)):
            for column in range(row, len(work)):
                updated = (
                    work[row][column]
                    - work[row][pivot_index] * work[column][pivot_index] / pivot
                )
                work[row][column] = updated
                work[column][row] = updated
    return LDLProof(order=len(work), pivots=tuple(pivots))


def certify_congruence_positive_definite(
    matrix: arb_mat | Sequence[Sequence[arb]],
    preconditioner: arb_mat | Sequence[Sequence[arb]],
    *,
    label: str = "matrix",
) -> CongruenceCertificate:
    """Validate a square symmetric matrix and an invertible congruence proof."""
    matrix_entries = _validated_symmetric_entries(matrix, label)
    preconditioner_entries = _matrix_entries(
        preconditioner, f"{label} preconditioner"
    )
    order = len(matrix_entries)
    if len(preconditioner_entries) != order:
        raise CertificationError(
            f"{label} preconditioner order {len(preconditioner_entries)} does not match {order}"
        )
    raw = arb_mat(matrix_entries)
    congruence = arb_mat(preconditioner_entries)
    preconditioner_determinant = require_nonzero(
        congruence.det(), f"{label} preconditioner determinant"
    )
    transformed = congruence * raw * congruence.transpose()
    ldl = certify_positive_ldl(transformed, label=f"{label} congruence")
    determinant = require_positive(raw.det(), f"{label} determinant")
    return CongruenceCertificate(
        order=order,
        preconditioner_determinant=preconditioner_determinant,
        determinant=determinant,
        ldl=ldl,
    )


def certify_monic_extremal_quadratic_form(
    matrix: arb_mat | Sequence[Sequence[arb]],
    quadratic_matrix: arb_mat | Sequence[Sequence[arb]],
    *,
    label: str = "matrix",
) -> MonicExtremalQuadraticFormCertificate:
    """Certify a monic quadratic extremal and contract a second form.

    For a symmetric order-``n`` matrix ``K``, this certifies the unique
    minimizer of ``q^T K q`` subject to ``q[n-1] = 1`` by proving that the
    leading block is positive definite and rigorously solving its normal
    equations.  ``quadratic_matrix`` is then contracted against the same
    interval vector.
    """
    matrix_entries = _validated_symmetric_entries(matrix, label)
    quadratic_entries = _validated_symmetric_entries(
        quadratic_matrix, f"{label} quadratic matrix"
    )
    order = len(matrix_entries)
    if len(quadratic_entries) != order:
        raise CertificationError(
            f"{label} quadratic matrix order {len(quadratic_entries)} "
            f"does not match {order}"
        )

    if order == 1:
        vector = (arb(1),)
        leading_ldl = LDLProof(order=0, pivots=())
        residuals: tuple[arb, ...] = ()
    else:
        leading_entries = [
            row[: order - 1] for row in matrix_entries[: order - 1]
        ]
        leading_ldl = certify_positive_ldl(
            leading_entries, label=f"{label} leading block"
        )
        leading = arb_mat(leading_entries)
        boundary = arb_mat(
            [[matrix_entries[row][order - 1]] for row in range(order - 1)]
        )
        try:
            solution = leading.solve(-boundary, algorithm="precond")
        except ZeroDivisionError as error:
            raise CertificationError(
                f"{label} leading block solve failed"
            ) from error
        residual_matrix = leading * solution + boundary
        residuals = tuple(
            residual_matrix[index, 0] for index in range(order - 1)
        )
        for index, residual in enumerate(residuals):
            if not residual.contains(arb(0)):
                raise CertificationError(
                    f"{label} normal residual {index} excludes zero: {residual}"
                )
        vector = tuple(
            solution[index, 0] for index in range(order - 1)
        ) + (arb(1),)

    column = arb_mat([[value] for value in vector])
    pivot = (
        column.transpose() * arb_mat(matrix_entries) * column
    )[0, 0]
    quadratic_form = (
        column.transpose() * arb_mat(quadratic_entries) * column
    )[0, 0]
    return MonicExtremalQuadraticFormCertificate(
        order=order,
        vector=vector,
        pivot=pivot,
        quadratic_form=quadratic_form,
        normal_residuals=residuals,
        leading_ldl=leading_ldl,
    )


def _normalization_record(name: str, determinant: arb, divisor: arb) -> NormalizationRecord:
    require_positive(determinant, f"{name} normalization determinant")
    require_positive(divisor, f"{name} normalization divisor")
    normalized = require_positive(
        determinant / divisor, f"{name} normalized determinant"
    )
    return NormalizationRecord(name, determinant, divisor, normalized)


def certify_diagonal_normalization(
    matrix: arb_mat | Sequence[Sequence[arb]], determinant: arb
) -> NormalizationRecord:
    entries = _matrix_entries(matrix, "diagonal normalization matrix")
    divisor = arb(1)
    for index in range(len(entries)):
        divisor *= require_positive(entries[index][index], f"matrix diagonal {index}")
    return _normalization_record("diagonal", determinant, divisor)


def certify_vandermonde_normalization(
    determinant: arb,
    nodes: Sequence[arb],
    *,
    diagonal_weights: Sequence[arb] | None = None,
) -> NormalizationRecord:
    if not nodes:
        raise CertificationError("Vandermonde normalization requires at least one node")
    if diagonal_weights is not None and len(diagonal_weights) != len(nodes):
        raise CertificationError("Vandermonde diagonal weights must match the nodes")
    divisor = arb(1)
    if diagonal_weights is not None:
        for index, weight in enumerate(diagonal_weights):
            divisor *= require_positive(weight, f"Vandermonde diagonal weight {index}")
    vandermonde = arb(1)
    for index, node in enumerate(nodes):
        for previous in range(index):
            difference = require_nonzero(
                node - nodes[previous],
                f"Vandermonde node difference ({previous}, {index})",
            )
            vandermonde *= difference
    divisor *= vandermonde**2
    return _normalization_record("vandermonde", determinant, divisor)


def certify_unique_argmin(
    winner: tuple[Key, arb], competitors: Iterable[tuple[Key, arb]]
) -> UniqueArgminCertificate[Key]:
    winner_key, winner_value = winner
    comparisons = 0
    for key, value in competitors:
        if key == winner_key:
            raise CertificationError(f"argmin winner repeated as a competitor: {key!r}")
        if not winner_value < value:
            relation = "tie" if winner_value == value else "overlap or reversed order"
            raise CertificationError(
                f"unique argmin rejected {winner_key!r} versus {key!r}: {relation}; "
                f"{winner_value} versus {value}"
            )
        comparisons += 1
    return UniqueArgminCertificate(winner_key, winner_value, comparisons)


def certify_series_scalar_agreement(
    name: str,
    points: Iterable[Point],
    *,
    series_evaluator: Callable[[Point], arb],
    scalar_evaluator: Callable[[Point], arb],
) -> SeriesCrossCheckRecord[Point]:
    """Cross-check a rigorously evaluated series against independent scalars.

    ``series_evaluator`` must include any truncation remainder enclosure.  The
    contract deliberately accepts callbacks so special functions and series
    composition strategies remain outside the generic certificate layer.
    """
    checks: list[SeriesCrossCheck[Point]] = []
    for point in points:
        series_value = series_evaluator(point)
        scalar_value = scalar_evaluator(point)
        if not series_value.overlaps(scalar_value):
            raise CertificationError(
                f"{name} series/scalar disagreement at {point!r}: "
                f"{series_value} versus {scalar_value}"
            )
        checks.append(SeriesCrossCheck(point, series_value, scalar_value))
    if not checks:
        raise CertificationError(f"{name} series cross-check requires sample points")
    return SeriesCrossCheckRecord(name, tuple(checks))
