# verify2_C22a.py -- second, separately written verifier (Arb balls + KD-tree).
# Author: David Niedbala Giraudin, ORCID 0009-0009-1526-1178
# Technical assistance: Claude (Anthropic)
#
# Certifies, for every certificate C22a_T*_meta.csv (+ _states, _final, ...) next to this
# script (or for the prefixes given as arguments), that the final polygon F is a (p,q) torus
# knot (P0 recomputed from (p,q), every linear step of the path certified by Lemma 2 of the
# note) and that F smoothed by unit-radius arcs has thickness >= 1, and prints rigorous Arb
# upper bounds for its length L and for L / Cr(T(p,q))**(3/4).  Needs numpy, scipy, python-flint.
import sys, time
import math
import numpy as np
from scipy.spatial import cKDTree
from flint import arb, ctx
from fractions import Fraction as Fr
ctx.prec = 120
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from p0 import p0_polygon          # only the definition of P0 is shared

def segdist_ball(A, B, C, D):
    """Arb enclosure of the distance between segments AB, CD (exact inputs):
    minimum of the convex quadratic over the unit square [0,1]x[0,1], computed over the interior
    critical point and the four edges, each in exact rationals, then a ball sqrt."""
    A = [Fr(int(x)) for x in A]; B = [Fr(int(x)) for x in B]
    C = [Fr(int(x)) for x in C]; D = [Fr(int(x)) for x in D]
    u = [B[k]-A[k] for k in range(3)]; v = [D[k]-C[k] for k in range(3)]; w = [A[k]-C[k] for k in range(3)]
    dt = lambda x, y: sum(x[k]*y[k] for k in range(3))
    a, b, c, d, e = dt(u,u), dt(u,v), dt(v,v), dt(u,w), dt(v,w)
    def f(s, t):
        return a*s*s - 2*b*s*t + c*t*t + 2*d*s - 2*e*t + dt(w,w)
    vals = []
    den = a*c - b*b
    if den != 0:
        s = (b*e - c*d)/den; t = (a*e - b*d)/den
        if 0 <= s <= 1 and 0 <= t <= 1: vals.append(f(s, t))
    for s in (Fr(0), Fr(1)):                      # edges s = const: t* = (e + b s)/c
        t = min(max((e + b*s)/c, Fr(0)), Fr(1)); vals.append(f(s, t))
    for t in (Fr(0), Fr(1)):                      # edges t = const: s* = (b t - d)/a
        s = min(max((b*t - d)/a, Fr(0)), Fr(1)); vals.append(f(s, t))
    m = min(vals)
    return arb(m.numerator).sqrt() / arb(m.denominator).sqrt(), m

def embedded_step(P, Q, U):
    """every polygon (1-t)P + tQ embedded: segments move by <= delta; non-adjacent
    pairs need distance > 2 delta (KD-tree on midpoints for candidates, exact
    distance for close ones); adjacent pairs need no folding."""
    N = len(P)
    delta = float(np.sqrt(((Q - P) ** 2).sum(1)).max()) * (1 + 1e-12) + 1e-9
    E = np.roll(P, -1, 0) - P; le = np.sqrt((E ** 2).sum(1).astype(float))
    E0 = np.roll(E, 1, 0); l0 = np.roll(le, 1)
    cth = (E0 * E).sum(1) / (l0 * le)
    if not ((1 + cth) >= 0.1).all():                 # keeps the float evaluation well conditioned
        return False
    if not (2 * delta < np.minimum(l0, le) * np.sqrt((1 + cth) / 2) * (1 - 1e-8)).all():
        return False
    M = (P + np.roll(P, -1, 0)) / 2.0
    rad = 2 * delta + le.max() + 1.0
    pairs = cKDTree(M).query_pairs(rad, output_type='ndarray')
    if len(pairs) == 0:
        return True
    i, j = pairs[:, 0], pairs[:, 1]
    g = np.abs(i - j); g = np.minimum(g, N - g); sel = g >= 2; i, j = i[sel], j[sel]
    lb = np.sqrt(((M[i] - M[j]) ** 2).sum(1)) * (1 - 1e-12) - (le[i] + le[j]) / 2 * (1 + 1e-12)
    for k in np.nonzero(lb <= 2 * delta + 1e-6)[0]:
        a, c = int(i[k]), int(j[k])
        d, _ = segdist_ball(P[a], P[(a+1) % N], P[c], P[(c+1) % N])
        if not d > 2 * delta:
            return False
    return True

def path_ok(states, U):
    n = 0
    for k in range(len(states) - 1):
        A, B = states[k], states[k + 1]
        disp = float(np.sqrt(((B - A) ** 2).sum(1)).max())
        m = max(1, int(np.ceil(disp / (0.05 * U))))
        while True:
            pts = [A + ((B - A) * j) // m for j in range(m)] + [B]
            if all(embedded_step(pts[j], pts[j+1], U) for j in range(m)):
                break
            m *= 2
            if m > 8192:
                return False
        n += m
    print("  path: %d steps, %d sub-steps" % (len(states) - 1, n), flush=True)
    return True

def thickness_ok(F, U):
    N = len(F)
    Fa = [[arb(int(c)) / U for c in row] for row in F]
    E = [[Fa[(i+1) % N][k] - Fa[i][k] for k in range(3)] for i in range(N)]
    nr = lambda x: (x[0]*x[0] + x[1]*x[1] + x[2]*x[2]).sqrt()
    L = [nr(e) for e in E]
    tanh = []
    for v in range(N):
        e1, e2 = E[v-1], E[v]
        dot = e1[0]*e2[0] + e1[1]*e2[1] + e1[2]*e2[2]
        cr = [e1[1]*e2[2]-e1[2]*e2[1], e1[2]*e2[0]-e1[0]*e2[2], e1[0]*e2[1]-e1[1]*e2[0]]
        th = nr(cr) / (L[v-1]*L[v] + dot)               # tan(theta/2)
        if not (2*th < L[v-1] and 2*th < L[v]):
            print("  minrad < 1 at", v); return None
        tanh.append(th)
    up = lambda x: math.nextafter(float(x.upper()), math.inf)     # float >= the ball
    tmax = max(up(t) for t in tanh)
    if not max(up(x) for x in L) < 3.14:                         # edges shorter than pi
        print("  edge too long"); return None
    tau = [t * t / 2 for t in tanh]
    taue = [max(up(tau[i]), up(tau[(i+1) % N])) for i in range(N)]
    cum = [arb(0)]
    for x in L: cum.append(cum[-1] + x)
    Ltot = cum[-1]
    thr = arb.pi() - 2 * tmax
    Ff = F.astype(float) / U
    M = (Ff + np.roll(Ff, -1, 0)) / 2
    lmax = max(float(x.upper()) for x in L)
    pairs = cKDTree(M).query_pairs(2 + 2 * max(taue) + lmax + 0.5, output_type='ndarray')
    nex = 0
    for i, j in pairs:
        i, j = int(min(i, j)), int(max(i, j))
        fwd = cum[j+1] - cum[i]; back = Ltot - (cum[j] - cum[i+1])
        sh = fwd if fwd.upper() < back.upper() else back
        if sh < thr:                         # every point pair at arc distance < pi - 2 tmax
            continue
        need = 2 + taue[i] + taue[j]
        lbf = np.linalg.norm(M[i] - M[j]) - (up(L[i]) + up(L[j])) / 2
        if lbf > need + 1e-6:
            continue
        nex += 1
        d, _ = segdist_ball(F[i], F[(i+1) % N], F[j], F[(j+1) % N])
        d = d / U
        if not d > need:
            print("  pair too close", i, j, d, need); return None
    print("  thickness: %d exact pairs, max tan %.4g" % (nex, tmax), flush=True)
    return Ltot


def load_cert(prefix):
    """Read a certificate stored as CSV files <prefix>_<array>.csv (exact integers);
    the first line of each file gives the shape of the array.  For a sequence of polygons
    marked "delta", the file holds the first polygon and then the differences between
    consecutive polygons, and the polygons are recovered by cumulative sums."""
    import glob as _glob
    import re as _re
    out = {}
    groups = {}                        # a large array may be split into parts <key>_p01, _p02, ...
    for fn in sorted(_glob.glob(prefix + "_*.csv")):
        key = fn[len(prefix) + 1:-4]
        mm = _re.fullmatch(r"(.+)_p(\d\d)", key)
        groups.setdefault(mm.group(1) if mm else key, []).append(fn)
    for key, fns in groups.items():
        with open(fns[0]) as fh:
            head = fh.readline()
        delta = head.startswith("# delta shape ")
        assert delta or head.startswith("# shape "), fns[0]
        shape = tuple(int(x) for x in head.split("shape ")[1].split(","))
        arr = np.concatenate([np.loadtxt(f, dtype=np.int64, delimiter=",", comments="#", ndmin=2)
                              for f in fns]).reshape(shape)
        if delta:                    # first polygon, then differences between consecutive polygons
            arr = np.cumsum(arr, axis=0)
        out[key] = arr
    return out

def main(path):
    z = load_cert(path)
    p, q, U, Rr, rr, kk = [int(v) for v in z["meta"]]
    S = z["states"].astype(np.int64); F = z["final"].astype(np.int64)
    if not (2 <= p < q and math.gcd(p, q) == 1 and Rr == 2 * q + 6 and rr == q + 2 and kk == 16):
        print("REJECTED: parameters"); return
    P0, _ = p0_polygon(p, q, Rr, rr, kk)
    P0i = np.array([[int(c * U) for c in v] for v in P0], dtype=np.int64)
    if not all((c * U).denominator == 1 for v in P0 for c in v):
        print("REJECTED: unit"); return
    if not np.array_equal(P0i, S[0]):
        print("REJECTED: first state"); return
    if "states_ref" in z:
        if not path_ok(list(S), U):
            print("REJECTED: path"); return
        Xl = S[-1]; D0 = np.zeros((2 * len(Xl), 3), dtype=np.int64)
        for v in range(len(Xl)):                     # midpoint refinement, written as a loop
            D0[2 * v] = 2 * Xl[v]; D0[2 * v + 1] = Xl[v] + Xl[(v + 1) % len(Xl)]
        U = 2 * U
        if not path_ok([D0] + list(z["states_ref"].astype(np.int64)) + [F], U):
            print("REJECTED: path"); return
    elif not path_ok(list(S) + [F], U):
        print("REJECTED: path"); return
    L = thickness_ok(F, U)
    if L is None:
        print("REJECTED: thickness"); return
    Cr = min(p * (q - 1), q * (p - 1))
    alpha = L / arb(Cr) ** (arb(3) / 4)
    print("CERTIFIED (verifier 2): L <=", L.upper(), " alpha <=", alpha.upper())

if __name__ == "__main__":
    import glob
    here = os.path.dirname(os.path.abspath(__file__))
    if len(sys.argv) > 1:
        prefixes = sys.argv[1:]
    else:                                   # default: every certificate next to this script
        prefixes = sorted(f[:-len("_meta.csv")] for f in glob.glob(os.path.join(here, "C22a_T*_meta.csv")))
    for pr in prefixes:
        print("=" * 70); print(os.path.basename(pr), flush=True)
        main(pr)