"""Exact certificate producer/checker and floating-point trace reference.

Optional, deliberately small development code. Python execution is tested, not
formally extracted. Generated image certificates can additionally be replayed by
Lean's kernel using VibeQCStage2.image_certificate_sound/all_ties.
"""
from dataclasses import dataclass
from fractions import Fraction as Q
from itertools import product
import math

# Conservative forward allowances derived from u0=2^-53, eta0=2^-1074:
# |y-x| <= u0|y|+eta0/2 implies <= 2*u0|x|+eta0.
U = Q(1, 2**52)
ETA = Q(1, 2**1074)


def rational(value):
    if isinstance(value, (float, int)) and not math.isfinite(value):
        raise ValueError('finite represented inputs required')
    try:
        return Q(value)
    except (TypeError, ValueError, OverflowError) as exc:
        raise ValueError('finite real or rational input required') from exc


def dot_budget(pairs, u=U, eta=ETA):
    """Exact evaluation of Lean dotBudget, with the list processed tail first."""
    budget, mass = Q(0), Q(0)
    for a, b in reversed(list(pairs)):
        term = abs(rational(a)*rational(b))
        budget = (1+u)*(u*term+eta+budget)+u*(term+mass)+eta
        mass += term
    return budget


def dot_trace(pairs, u=U, eta=ETA):
    """Check every actual separate binary64 multiply/add against exact operands.

    Invalid/nonfinite operations fail. Does not assume NumPy/BLAS uses this order.
    """
    pairs = [(float(a), float(b)) for a, b in pairs]
    total = 0.0
    operations = []
    for a, b in reversed(pairs):
        exact_product = rational(a)*rational(b)
        p = a*b
        if abs(rational(p)-exact_product) > u*abs(exact_product)+eta:
            raise ValueError('product violates rounding allowance')
        exact_sum = rational(p)+rational(total)
        result = p+total
        if abs(rational(result)-exact_sum) > u*abs(exact_sum)+eta:
            raise ValueError('sum violates rounding allowance')
        operations.append((a, b, p, total, result))
        total = result
    exact = sum((rational(a)*rational(b) for a, b in pairs), Q(0))
    bound = dot_budget(pairs, u, eta)
    if abs(rational(total)-exact) > bound:
        raise AssertionError('trace bound violated')
    return total, exact, bound, operations


def df_bounds(B, D, rho_hat, J_hat):
    """Check native stage errors and the composed absolute bound in exact rationals.

    Bounds are calculated before comparing each observed output. The production
    association is not asserted to be a DotTrace; its two local error premises
    are checked directly for this input/output, without empirical tolerances.
    """
    p, n = len(B), len(D)
    if n == 0 or p == 0 or any(len(row) != n for row in D):
        raise ValueError('nonempty square density required')
    if len(rho_hat) != p or len(J_hat) != n or any(len(row) != n for row in J_hat):
        raise ValueError('output shape mismatch')
    if any(len(bp) != n or any(len(row) != n for row in bp) for bp in B):
        raise ValueError('factor shape mismatch')
    B = [[[rational(x) for x in row] for row in bp] for bp in B]
    D = [[rational(x) for x in row] for row in D]
    rh = list(map(rational, rho_hat))
    rho, inner = [], []
    for a in range(p):
        pairs = [(B[a][i][j], D[i][j]) for i in range(n) for j in range(n)]
        exact = sum((b*d for b, d in pairs), Q(0))
        budget = dot_budget(pairs)
        if abs(rh[a]-exact) > budget:
            raise ValueError('inner DF error exceeds derived allowance')
        rho.append(exact); inner.append(budget)
    rows = []
    for i in range(n):
        for j in range(n):
            exact = sum((B[a][i][j]*rho[a] for a in range(p)), Q(0))
            rounded_inner_exact = sum((B[a][i][j]*rh[a] for a in range(p)), Q(0))
            outer = dot_budget([(B[a][i][j], rh[a]) for a in range(p)])
            observed = rational(J_hat[i][j])
            if abs(observed-rounded_inner_exact) > outer:
                raise ValueError('outer DF error exceeds derived allowance')
            total = outer + sum((abs(B[a][i][j])*inner[a] for a in range(p)), Q(0))
            error = abs(observed-exact)
            if error > total:
                raise AssertionError('composed DF bound violated')
            rows.append({'i': i, 'j': j, 'error': error, 'bound': total,
                         'outer': outer, 'propagated': total-outer})
    return rows


class SearchBudgetExceeded(ValueError):
    pass


def matrix3(A):
    if len(A) != 3 or any(len(row) != 3 for row in A):
        raise ValueError('3 by 3 column-basis matrix required')
    return tuple(tuple(rational(x) for x in row) for row in A)


def inverse3(A):
    a = [list(row)+[Q(i == j) for j in range(3)] for i, row in enumerate(A)]
    for j in range(3):
        pivot = next((i for i in range(j, 3) if a[i][j]), None)
        if pivot is None:
            raise ValueError('singular lattice')
        a[j], a[pivot] = a[pivot], a[j]
        divisor = a[j][j]; a[j] = [v/divisor for v in a[j]]
        for i in range(3):
            if i != j:
                factor = a[i][j]; a[i] = [v-factor*w for v, w in zip(a[i], a[j])]
    return tuple(tuple(row[3:]) for row in a)


def matvec(A, x):
    return tuple(sum((a*b for a, b in zip(row, x)), Q(0)) for row in A)


def cost(A, x, n):
    v = matvec(A, n)
    return sum(((a-b)**2 for a, b in zip(x, v)), Q(0))


def ceil_sqrt(q):
    if q < 0:
        raise ValueError('negative square')
    r = math.isqrt(q.numerator//q.denominator)
    return r if Q(r*r) == q else r+1


def box_points(lo, hi, budget):
    if type(budget) is not int or budget < 1:
        raise ValueError('positive integer work budget required')
    size = math.prod(max(0, h-l+1) for l, h in zip(lo, hi))
    if size > budget:
        raise SearchBudgetExceeded(f'complete box needs {size} candidates, budget {budget}')
    return product(*(range(l, h+1) for l, h in zip(lo, hi)))


@dataclass(frozen=True)
class ImageCertificate:
    A: tuple
    H: tuple
    x: tuple
    R: Q
    lo: tuple
    hi: tuple
    best: tuple
    ties: tuple


def certify_image(A, x, max_candidates=100000, candidate=None):
    """Exact complete 3D triclinic search; no Minkowski or shell hypothesis."""
    A = matrix3(A)
    if len(x) != 3:
        raise ValueError('three displacement coordinates required')
    x = tuple(map(rational, x)); H = inverse3(A); f = matvec(H, x)
    incumbent = tuple(round(v) for v in f) if candidate is None else tuple(candidate)
    if len(incumbent) != 3 or any(type(v) is not int for v in incumbent):
        raise ValueError('integer candidate required')
    R = Q(ceil_sqrt(cost(A, x, incumbent)))
    radii = tuple(sum(map(abs, row))*R for row in H)
    lo = tuple(math.floor(v-r) for v, r in zip(f, radii))
    hi = tuple(math.ceil(v+r) for v, r in zip(f, radii))
    best_cost, ties = None, []
    for n in box_points(lo, hi, max_candidates):
        value = cost(A, x, n)
        if best_cost is None or value < best_cost:
            best_cost, ties = value, [n]
        elif value == best_cost:
            ties.append(n)
    if not ties:
        raise AssertionError('complete box unexpectedly empty')
    certificate = ImageCertificate(A, H, x, R, lo, hi, ties[0], tuple(ties))
    verify_image(certificate, max_candidates)
    return certificate


def verify_image(c, max_candidates=100000):
    """Recheck all ValidImageCert premises and equality of the complete tie set."""
    A, H = matrix3(c.A), matrix3(c.H)
    if len(c.x) != 3 or any(len(v) != 3 for v in (c.lo, c.hi, c.best)):
        raise ValueError('certificate shape mismatch')
    if any(type(v) is not int for row in (c.lo, c.hi, c.best, *c.ties) for v in row):
        raise ValueError('integer images required')
    x, R = tuple(map(rational, c.x)), rational(c.R)
    if any(sum(H[i][a]*A[a][j] for a in range(3)) != Q(i == j)
           for i in range(3) for j in range(3)):
        raise ValueError('inverse witness invalid')
    best = cost(A, x, c.best)
    if R < 0 or best > R*R:
        raise ValueError('incumbent radius invalid')
    f = matvec(H, x); radii = tuple(sum(map(abs, row))*R for row in H)
    if any(l > v-r or v+r > h for l, h, v, r in zip(c.lo, c.hi, f, radii)):
        raise ValueError('box completeness inequalities invalid')
    ties = []
    for n in box_points(c.lo, c.hi, max_candidates):
        value = cost(A, x, n)
        if value < best:
            raise ValueError('selected image is not nearest')
        if value == best:
            ties.append(n)
    if tuple(ties) != c.ties or c.best not in ties:
        raise ValueError('tie set incomplete, duplicated or misordered')
    return True


def lean_q(q):
    q = rational(q)
    return f'({q.numerator} / {q.denominator} : ℚ)'


def lean_vec(xs, rational_values=True):
    return '!['+', '.join(lean_q(x) if rational_values else str(x) for x in xs)+']'


def lean_image_case(label, c):
    """Untrusted producer; kernel checks the full certificate and exact tie set."""
    if not label.isidentifier() or not label.isascii():
        raise ValueError('ASCII Lean case identifier required')
    verify_image(c)
    matrix = lambda A: '!['+', '.join(lean_vec(row) for row in A)+']'
    s = f'def {label}A : Fin 3 → Fin 3 → ℚ := {matrix(c.A)}\n'
    s += f'def {label}H : Fin 3 → Fin 3 → ℚ := {matrix(c.H)}\n'
    s += f'def {label}x : Fin 3 → ℚ := {lean_vec(c.x)}\n'
    for name, xs in [('lo', c.lo), ('hi', c.hi), ('best', c.best)]:
        s += f'def {label}{name} : Fin 3 → ℤ := {lean_vec(xs, False)}\n'
    args = f'{label}A {label}H {label}x {lean_q(c.R)} {label}lo {label}hi {label}best'
    s += f'theorem {label}_valid : ValidImageCert {args} := by decide +kernel\n'
    s += f'theorem {label}_nearest : ∀ n, imageCost {label}A {label}x {label}best ≤ imageCost {label}A {label}x n :=\n'
    s += f'  image_certificate_sound {args} {label}_valid\n'
    expected = '{'+', '.join(lean_vec(n, False) for n in c.ties)+'}'
    s += f'theorem {label}_ties : certifiedTies {label}A {label}x {label}lo {label}hi {label}best = {expected} := by decide +kernel\n'
    return s


def lean_real(q):
    q = rational(q)
    return f'({q.numerator} : ℝ)' if q.denominator == 1 else f'({q.numerator} / {q.denominator} : ℝ)'


def lean_trace_case(label, pairs):
    """Replay the finite trace's local rounding premises with kernel-checked arithmetic."""
    if not label.isidentifier() or not label.isascii():
        raise ValueError('ASCII Lean case identifier required')
    total, exact, bound, operations = dot_trace(pairs)
    list_text = '['+', '.join('('+lean_real(a)+', '+lean_real(b)+')' for a, b in pairs)+']'
    u, eta = lean_real(U), lean_real(ETA)
    s = f'noncomputable def {label}Inputs : List (ℝ × ℝ) := {list_text}\n'
    s += f'theorem {label}_trace : DotTrace {u} {eta} {label}Inputs {lean_real(total)} := by\n'
    s += f'  unfold {label}Inputs\n'
    # Each new head wraps the previously processed tail in the inductive trace.
    def proof(ops, indent):
        pad = ' '*indent
        if not ops:
            return pad+'exact DotTrace.nil\n'
        a,b,p,t,z = ops[-1]
        text = pad+'refine DotTrace.cons '+ ' '.join(map(lean_real,(a,b,p,t,z)))+' _ ?_ ?_ ?_\n'
        text += pad+'·\n'+proof(ops[:-1],indent+2)
        text += pad+'· norm_num\n'+pad+'· norm_num\n'
        return text
    s += proof(operations, 2)
    s += f'theorem {label}_bound : |{lean_real(total)} - dotExact {label}Inputs| ≤ dotBudget {u} {eta} {label}Inputs :=\n'
    s += f'  dot_trace_bound {u} {eta} (by norm_num) {label}Inputs {lean_real(total)} {label}_trace\n'
    return s


def lean_error_case(label, observed, exact, bound):
    if not label.isidentifier() or not label.isascii():
        raise ValueError('ASCII Lean case identifier required')
    return f'theorem {label}_enclosure : |{lean_q(observed)} - {lean_q(exact)}| ≤ {lean_q(bound)} := by decide +kernel\n'