#!/usr/bin/env python3
"""Slow, direct verification of the geometric leaves of both proofs.

This file imports no search code.  It starts with the static cell
lists, the displayed forms, and the displayed rays.  For each proposed cell it
checks the empty-sphere inequality over *all* of Z^5 by exhaustive enumeration
inside a rigorously bounded ellipsoid.  Exact cell volumes sum to one.  It then
checks directly that every displayed ray preserves every cell.

Run with Sage's Python; all arithmetic is rational and every assertion is a
proof check, not a heuristic search.
"""

import argparse
from collections import Counter
from itertools import combinations, product
from math import gcd, lcm

from sage.all import QQ, ZZ, Polyhedron, ceil, floor, matrix, vector

from counterexample_cells import CLASSICAL_CELLS, THETA_CELLS


G = 5
PAIRS = tuple((i, j) for i in range(G) for j in range(i, G))


def symmetric_matrix(coordinates):
    """Upper-triangular coordinates -> symmetric matrix."""
    H = matrix(QQ, G, G, 0)
    for value, (i, j) in zip(coordinates, PAIRS):
        H[i, j] = H[j, i] = value
    return H


def quadratic(H, x):
    x = vector(QQ, x)
    return (x * H * x) / 2


def q_coefficients(x):
    """Linear functional H |-> q_H(x) in upper-triangular coordinates."""
    return vector(QQ, [QQ(x[i] * x[j]) if i != j else QQ(x[i] ** 2) / 2
                       for i, j in PAIRS])


def interpolation_operator(cell):
    """Matrix carrying the coordinates of H to the affine interpolant."""
    basis_rows = []
    for v in cell:
        row = list(v) + [1]
        if matrix(QQ, basis_rows + [row]).rank() > len(basis_rows):
            basis_rows.append(row)
        if len(basis_rows) == G + 1:
            break
    assert len(basis_rows) == G + 1
    interpolation_matrix = matrix(QQ, basis_rows)
    value_matrix = matrix(QQ, [q_coefficients(row[:G]) for row in basis_rows])
    return interpolation_matrix.inverse() * value_matrix


def affine_interpolant(H, cell):
    """Return ell,d with q_H(v)=ell.v+d at every vertex of ``cell``."""
    coordinates = vector(QQ, [H[i, j] for i, j in PAIRS])
    coefficients = interpolation_operator(cell) * coordinates
    ell, constant = vector(QQ, coefficients[:G]), coefficients[G]
    assert all(quadratic(H, v) == ell * vector(QQ, v) + constant
               for v in cell)
    return ell, constant


def slack_row(cell, x):
    """Coefficient row of q_H(x) minus the cell interpolant, as H varies."""
    affine_row = vector(QQ, list(x) + [1]) * interpolation_operator(cell)
    return q_coefficients(x) - affine_row


def primitive_vector(values):
    """Clear denominators and divide by the gcd, preserving orientation."""
    values = vector(QQ, values)
    denominator = 1
    for value in values:
        denominator = lcm(denominator, int(value.denominator()))
    integers = [int(denominator * value) for value in values]
    divisor = gcd(*map(abs, integers))
    assert divisor > 0
    return tuple(value // divisor for value in integers)


def integer_interval(center, square_bound):
    """All integers n satisfying (n-center)^2 <= square_bound, exactly."""
    assert square_bound >= 0
    low = floor(center)
    while (low - center) ** 2 <= square_bound:
        low -= 1
    high = ceil(center)
    while (high - center) ** 2 <= square_bound:
        high += 1
    return range(int(low + 1), int(high))


def nonpositive_points(H, ell, constant):
    """Enumerate every integral x with q_H(x)-ell.x-constant <= 0.

    If z=x-H^-1 ell and F(x)<=0, Cauchy--Schwarz for H gives

        z_i^2 <= (H^-1)_ii z^t H z
              <= 2*h*(H^-1)_ii,

    where h=constant+ell.H^-1.ell/2.  The resulting coordinate box is
    therefore exhaustive, although usually much larger than the ellipsoid.
    """
    assert H.is_positive_definite()
    center = H.solve_right(ell)
    height = constant + (ell * center) / 2
    assert height >= 0
    inverse = H.inverse()
    ranges = [integer_interval(center[i], 2 * height * inverse[i, i])
              for i in range(H.nrows())]
    found = set()
    checked = 0
    for x_tuple in product(*ranges):
        checked += 1
        x = vector(QQ, x_tuple)
        if quadratic(H, x) - ell * x - constant <= 0:
            found.add(tuple(ZZ(a) for a in x))
    return found, checked


def canonical_translate(cell):
    """Canonical key for a bounded lattice polytope modulo translation."""
    return min(tuple(sorted(tuple(a - b for a, b in zip(v, origin))
                            for v in cell))
               for origin in cell)


def verify_delaunay_decomposition(name, B, cells):
    """Check exact cells, one representative per orbit, and volume one."""
    assert B.is_positive_definite()
    assert len({canonical_translate(cell) for cell in cells}) == len(cells)
    total_volume = QQ(0)
    volume_types = Counter()
    enumerated = 0
    for index, cell in enumerate(cells):
        ell, constant = affine_interpolant(B, cell)
        equality_or_below, count = nonpositive_points(B, ell, constant)
        enumerated += count
        assert equality_or_below == set(cell), (name, index)
        polytope = Polyhedron(vertices=cell, base_ring=QQ)
        assert polytope.dim() == G
        volume = polytope.volume()
        total_volume += volume
        volume_types[(len(cell), volume)] += 1
    assert total_volume == 1
    print(f"{name}: {len(cells)} exact Delaunay-cell orbits; volume sum 1")
    print(f"  {enumerated} lattice points checked in rigorous ellipsoid boxes")
    return volume_types


def verify_positive_definite_ray(H, cells):
    """Check q_H lies above every cell's affine interpolant on all Z^5."""
    assert H.is_positive_definite()
    enumerated = 0
    for cell in cells:
        ell, constant = affine_interpolant(H, cell)
        points, count = nonpositive_points(H, ell, constant)
        enumerated += count
        assert all(quadratic(H, x) - ell * vector(QQ, x) - constant == 0
                   for x in points)
    return enumerated


def verify_singular_ray(H, cells, kernel, complement):
    """Reduce a PSD ray modulo its displayed saturated integral kernel."""
    kernel = matrix(ZZ, kernel).transpose()
    complement = matrix(ZZ, complement).transpose()
    change = kernel.augment(complement)
    assert abs(change.det()) == 1
    assert H.is_positive_semidefinite() and H.rank() == complement.ncols()
    assert H * kernel == 0
    enumerated = 0
    for cell in cells:
        ell, constant = affine_interpolant(H, cell)
        assert ell * kernel == 0
        reduced_H = complement.transpose() * H * complement
        reduced_ell = ell * complement
        points, count = nonpositive_points(reduced_H, reduced_ell, constant)
        enumerated += count
        assert all(quadratic(reduced_H, y)
                   - reduced_ell * vector(QQ, y) - constant == 0
                   for y in points)
    return enumerated


def verify_rank_one_ray(H, cells, linear_form):
    """Use (t-k)(t-k-1)>=0 instead of any lattice enumeration."""
    a = vector(ZZ, linear_form)
    gram = 2 * matrix(ZZ, G, 1, list(a)) * matrix(ZZ, 1, G, list(a))
    assert H == gram and gcd(*map(abs, a)) == 1
    for cell in cells:
        values = sorted({a * vector(ZZ, v) for v in cell})
        assert len(values) == 2 and values[1] == values[0] + 1
        k = values[0]
        ell, constant = affine_interpolant(H, cell)
        assert ell == (2 * k + 1) * a
        assert constant == -k * (k + 1)
        # For integral x, the slack is (a.x-k)(a.x-k-1), always nonnegative.


def verify_exact_secondary_cone(name, ray_coordinates, cells, witnesses):
    """Prove that the displayed ray cone is the whole closed DV cone.

    The cell co-sphericity equations cut the 15-dimensional form space down to
    the span of the rays.  On that span, we write down genuine lattice-point
    inequalities and recover their extreme rays by trying every possible set
    of d-1 active inequalities.  The opposite containment was proved by the
    global ray checks above.
    """
    rays = matrix(QQ, ray_coordinates)
    ray_rank = rays.rank()

    equality_rows = []
    for cell in cells:
        for vertex in cell:
            row = slack_row(cell, vertex)
            if row:
                equality_rows.append(row)
    equalities = matrix(QQ, equality_rows)
    assert equalities.rank() == len(PAIRS) - ray_rank
    assert equalities * rays.transpose() == 0

    # Choose ray rows forming a basis of their span and use them as coordinates.
    basis_rows = []
    for ray in rays.rows():
        if matrix(QQ, basis_rows + [ray]).rank() > len(basis_rows):
            basis_rows.append(ray)
    assert len(basis_rows) == ray_rank
    basis = matrix(QQ, basis_rows)

    def in_basis(row):
        coordinates = basis.transpose().solve_right(vector(QQ, row))
        assert coordinates * basis == vector(QQ, row)
        return coordinates

    coordinate_rays = matrix(QQ, [in_basis(ray) for ray in rays.rows()])
    witness_rows = []
    for cell_index, lattice_point in witnesses:
        row = slack_row(cells[cell_index], lattice_point)
        witness_rows.append(row * basis.transpose())
    facets = matrix(QQ, witness_rows)
    profiles = facets * coordinate_rays.transpose()
    assert min(profiles.list()) >= 0

    # Each recorded inequality really is a facet: its zero rays span d-1.
    for profile in profiles.rows():
        zero_rays = [coordinate_rays.row(i)
                     for i, value in enumerate(profile) if value == 0]
        assert matrix(QQ, zero_rays).rank() == ray_rank - 1
        assert any(value > 0 for value in profile)

    # The inequalities define a full-dimensional pointed cone.  A positive
    # sum of the proposed rays is strictly inside it, and their normals span.
    interior = sum(coordinate_rays.rows(), vector(QQ, ray_rank))
    assert all(value > 0 for value in facets * interior)
    assert facets.rank() == ray_rank

    # In a pointed full-dimensional d-cone, every extreme ray is the common
    # kernel of d-1 independent active facet rows.  Exhaust those tiny sets.
    extreme_rays = set()
    for active in combinations(range(facets.nrows()), ray_rank - 1):
        active_rows = matrix(QQ, [facets.row(i) for i in active])
        if active_rows.rank() != ray_rank - 1:
            continue
        direction = vector(QQ, active_rows.right_kernel().basis()[0])
        values = facets * direction
        if all(value <= 0 for value in values):
            direction = -direction
            values = -values
        if all(value >= 0 for value in values):
            assert any(value > 0 for value in values)
            extreme_rays.add(primitive_vector(direction))

    proposed_rays = {primitive_vector(ray) for ray in coordinate_rays.rows()}
    assert len(proposed_rays) == len(ray_coordinates)
    assert extreme_rays == proposed_rays
    print(f"  {name} ray cone is exact: {equalities.rank()} equations and "
          f"{facets.nrows()} witnessed facets")


THETA_FORM = matrix(ZZ, (
    (14, 9, 7, 7, 14),
    (9, 16, 7, 7, 14),
    (7, 7, 16, 7, 14),
    (7, 7, 7, 16, 14),
    (14, 14, 14, 14, 24),
))
THETA_K = (2, 1, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 2)
THETA_RAYS = (
    (0,0,0,0,0,0,0,0,0,0,0,0,1,0,0),
    (0,0,0,0,0,0,0,0,0,1,0,0,0,0,0),
    (0,0,0,0,0,0,0,0,0,1,2,2,1,2,1),
    (0,0,0,0,0,1,0,0,0,0,0,0,0,0,0),
    (1,1,1,1,2,1,1,1,2,1,1,2,1,2,2),
    (1,2,0,2,2,1,0,2,2,0,0,0,1,2,1),
    (1,2,2,0,2,1,2,0,2,1,0,2,0,0,1),
    (4,4,4,4,8,4,4,4,8,4,4,8,4,8,7),
)
THETA_RANK_ONE_FORMS = {
    0: (0, 0, 0, 1, 0),
    1: (0, 0, 1, 0, 0),
    2: (0, 0, 1, 1, 1),
    3: (0, 1, 0, 0, 0),
    5: (1, 1, 0, 1, 1),
    6: (1, 1, 1, 0, 1),
}
THETA_CELL = (
    (0,0,0,0,0), (1,0,0,0,0), (0,1,0,0,0),
    (0,0,1,0,0), (0,0,0,1,0), (1,1,1,1,-2),
)
THETA_HOLE = vector(ZZ, (1, 1, 1, 1, -1))
THETA_BETA = vector(ZZ, (1, 1, 1, 1, 3, 2, 2, 10))
# Each pair is (cell-list index, lattice point). Its empty-sphere slack is a
# supporting inequality for one facet of the eight-ray cone.
THETA_FACET_WITNESSES = (
    (0, (0,1,0,0,-1)),
    (2, (1,1,1,0,-2)),
    (2, (0,0,1,0,-1)),
    (2, (1,1,1,1,-3)),
    (15, (1,1,0,0,-1)),
    (5, (1,1,0,0,-1)),
    (15, (1,0,-1,0,0)),
    (15, (1,0,0,-1,0)),
)


def verify_theta():
    assert all(THETA_FORM[i, i] % 2 == 0 for i in range(G))
    volume_types = verify_delaunay_decomposition(
        "theta", THETA_FORM, THETA_CELLS)
    assert volume_types == Counter({
        (6, QQ(1)/120): 20, (6, QQ(1)/60): 2,
        (7, QQ(1)/60): 32, (9, QQ(1)/15): 4,
    })
    ambient_rays = [tuple(k * x for k, x in zip(THETA_K, ray))
                    for ray in THETA_RAYS]
    ray_forms = [symmetric_matrix(ray) for ray in ambient_rays]
    assert matrix(ZZ, THETA_RAYS).rank() == 8
    assert all(gcd(*map(abs, ray)) == 1 for ray in THETA_RAYS)
    assert sum(ray_forms, matrix(ZZ, G, G, 0)) == THETA_FORM
    enumerated = 0
    for i, H in enumerate(ray_forms):
        assert H.is_positive_semidefinite()
        if i in THETA_RANK_ONE_FORMS:
            verify_rank_one_ray(H, THETA_CELLS, THETA_RANK_ONE_FORMS[i])
        else:
            enumerated += verify_positive_definite_ray(H, THETA_CELLS)
    verify_exact_secondary_cone(
        "theta", ambient_rays, THETA_CELLS, THETA_FACET_WITNESSES)
    assert frozenset(THETA_CELL) in [frozenset(cell) for cell in THETA_CELLS]
    target = THETA_CELL
    support = vector(QQ, [
        affine_interpolant(H, target)[0] * THETA_HOLE
        + affine_interpolant(H, target)[1]
        for H in ray_forms
    ])
    assert support == THETA_BETA
    nonzero_vertices = [vector(ZZ, v) for v in THETA_CELL if any(v)]
    assert sum(nonzero_vertices, vector(ZZ, G)) == 2 * THETA_HOLE
    assert abs(matrix(ZZ, nonzero_vertices).transpose().det()) == 2
    print("  8 theta rays preserve every cell (6 by n(n-1)>=0;")
    print(f"  2 by {enumerated} exhaustive positive-definite checks)")
    print("  determinant-two hole has support values", tuple(support))


CLASSICAL_FORM = matrix(ZZ, (
    (4, 0, 0, 0, -2),
    (0, 4, 0, 0, -2),
    (0, 0, 4, 0, -2),
    (0, 0, 0, 4, -2),
    (-2, -2, -2, -2, 5),
))
CLASSICAL_C = matrix(ZZ, (
    (-1,-1,-1,0,2), (-1,-1,0,-1,2), (-1,0,-1,-1,2),
    (0,-1,-1,-1,2), (-1,-1,-1,-1,2), (-1,-1,-1,-1,3),
    (-2,-1,-1,-1,3), (-1,-2,-1,-1,3),
    (-1,-1,-2,-1,3), (-1,-1,-1,-2,3),
))
CLASSICAL_BETA = vector(ZZ, (1,1,1,1,1,2,2,2,2,2))
CLASSICAL_PROJECTION = matrix(ZZ, (
    (0,0,0,0,0,0,0,0,0,0,0,0,-2,1,0),
    (0,0,0,0,0,0,0,0,0,-2,0,1,0,0,0),
    (0,0,0,0,0,-2,0,0,1,0,0,0,0,0,0),
    (-2,0,0,0,1,0,0,0,0,0,0,0,0,0,0),
    (0,0,0,0,0,0,0,0,0,0,0,0,0,0,1),
))
CLASSICAL_LIFT = matrix(ZZ, (
    (0,0,0,0,0), (0,0,0,0,0), (0,0,0,0,0),
    (0,0,0,0,0), (0,0,0,1,0), (0,0,0,0,0),
    (0,0,0,0,0), (0,0,0,0,0), (0,0,1,0,0),
    (0,0,0,0,0), (0,0,0,0,0), (0,1,0,0,0),
    (0,0,0,0,0), (1,0,0,0,0), (0,0,0,0,1),
))


def standard_basis(i):
    return tuple(1 if i == j else 0 for j in range(G))


CLASSICAL_SINGULAR_DATA = (
    (standard_basis(0), tuple(standard_basis(i) for i in (1,2,3,4))),
    (standard_basis(1), tuple(standard_basis(i) for i in (0,2,3,4))),
    (standard_basis(2), tuple(standard_basis(i) for i in (0,1,3,4))),
    (standard_basis(3), tuple(standard_basis(i) for i in (0,1,2,4))),
    ((1,1,1,1,2), tuple(standard_basis(i) for i in (1,2,3,4))),
)
# One genuine empty-sphere inequality for each facet, in the same
# (cell-list index, lattice point) format as above.
CLASSICAL_FACET_WITNESSES = (
    (4, (0,0,0,1,1)),
    (1, (0,0,1,0,1)),
    (9, (1,1,1,1,2)),
    (0, (0,1,0,0,1)),
    (3, (0,0,-1,0,0)),
    (2, (0,0,0,1,0)),
    (2, (1,0,0,0,1)),
    (0, (0,0,1,0,0)),
    (7, (0,-1,0,0,0)),
    (1, (0,0,0,1,0)),
)


def verify_classical():
    volume_types = verify_delaunay_decomposition(
        "classical", CLASSICAL_FORM, CLASSICAL_CELLS)
    assert volume_types == Counter({(8, QQ(1)/30): 30})
    assert CLASSICAL_PROJECTION * CLASSICAL_LIFT == matrix.identity(ZZ, G)
    ambient_rays = CLASSICAL_C * CLASSICAL_PROJECTION
    ray_forms = [symmetric_matrix(row) for row in ambient_rays.rows()]
    assert ambient_rays.rank() == G and CLASSICAL_C.rank() == G
    assert all(gcd(*map(abs, row)) == 1 for row in CLASSICAL_C.rows())
    assert all(gcd(*map(abs, row)) == 1 for row in ambient_rays.rows())
    assert sum(ray_forms, matrix(ZZ, G, G, 0)) == 5 * CLASSICAL_FORM
    enumerated = 0
    for i, H in enumerate(ray_forms):
        if i < 5:
            kernel, complement = CLASSICAL_SINGULAR_DATA[i]
            enumerated += verify_singular_ray(
                H, CLASSICAL_CELLS, (kernel,), complement)
        else:
            enumerated += verify_positive_definite_ray(H, CLASSICAL_CELLS)
    verify_exact_secondary_cone(
        "classical", ambient_rays.rows(), CLASSICAL_CELLS,
        CLASSICAL_FACET_WITNESSES)
    # Translate the first cell so that it contains both 0 and c_cl.
    origin = vector(ZZ, (1,1,1,1,2))
    target = tuple(tuple(vector(ZZ, v) - origin) for v in CLASSICAL_CELLS[0])
    c = vector(ZZ, (0,0,0,0,-1))
    assert tuple(c) in target and (0,0,0,0,0) in target
    support = vector(QQ, [
        affine_interpolant(H, target)[0] * c
        + affine_interpolant(H, target)[1]
        for H in ray_forms
    ])
    assert support == vector(QQ, (1,1,1,1,1,QQ(3)/2,QQ(3)/2,
                                  QQ(3)/2,QQ(3)/2,QQ(3)/2))
    assert vector(ZZ, [ceil(x) for x in support]) == CLASSICAL_BETA
    print(f"  10 classical rays preserve every cell ({enumerated} checks)")
    print("  c_cl support values:", tuple(support))
    print("  their ceilings give:", tuple(CLASSICAL_BETA))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("example", choices=("classical", "theta", "both"),
                        nargs="?", default="both")
    args = parser.parse_args()
    if args.example in ("classical", "both"):
        verify_classical()
    if args.example in ("theta", "both"):
        verify_theta()
    print("direct geometric certificate: verified")


if __name__ == "__main__":
    main()