# verify2_dist.py -- second, separately written verifier for the trefoil distortion certificate:
# path by KD-tree search + floats with margins + exact rational distances; distortion by Arb
# balls with Lipschitz bounds on boxes.  Only the definition of P0 (p0.py) is shared with
# verify_dist.py.  Usage: python3 verify2_dist.py  (no argument).  Needs scipy, python-flint.
# Author: David Niedbala Giraudin, ORCID 0009-0009-1526-1178
# Technical assistance: Claude (Anthropic)
import sys, time, os
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 load_cert(prefix):
    """Read the certificate stored as CSV files <prefix>_<array>.csv (exact integers);
    the first line of each file gives the shape of the array."""
    import glob as _glob
    out = {}
    for fn in sorted(_glob.glob(prefix + "_*.csv")):
        key = fn[len(prefix) + 1:-4]
        with open(fn) as fh:
            head = fh.readline()
        assert head.startswith("# shape "), fn
        shape = tuple(int(x) for x in head[len("# shape "):].split(","))
        arr = np.loadtxt(fn, dtype=np.int64, delimiter=",", comments="#", ndmin=2)
        out[key] = arr.reshape(shape)
    return out

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) >= 1e-4).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 check(F, D, U=1, max_depth=30):
    """F: integer polygon (N,3). D: bound (string or arb). Returns (ok, boxes)."""
    N = len(F); D = arb(D)
    P = [[arb(int(c)) / U for c in v] for v in F]
    E = [[P[(k+1) % N][a] - P[k][a] for a in range(3)] for k in range(N)]
    l = [(e[0]**2 + e[1]**2 + e[2]**2).sqrt() for e in E]
    cum = [arb(0)]
    for x in l: cum.append(cum[-1] + x)
    tot = cum[-1]
    boxes = 0
    def pt(i, s): return [P[i][a] + s * E[i][a] for a in range(3)]
    for i in range(N):
        for j in range(i + 1, N):
            if j == i + 1 or (i == 0 and j == N - 1):
                V = P[j] if j == i + 1 else P[0]
                A = P[i] if j == i + 1 else P[1]
                B = P[(j + 1) % N] if j == i + 1 else P[N - 1]
                u = [A[a] - V[a] for a in range(3)]; v = [B[a] - V[a] for a in range(3)]
                cphi = (u[0]*v[0] + u[1]*v[1] + u[2]*v[2]) / ((u[0]**2+u[1]**2+u[2]**2).sqrt() * (v[0]**2+v[1]**2+v[2]**2).sqrt())
                bound = 1 / ((1 - cphi) / 2).sqrt()          # 1/sin(phi/2)
                if not (bound < D): return False, ("adjacent", i, j)
                continue
            stack = [(0.0, 1.0, 0.0, 1.0, 0)]
            while stack:
                s0, s1, t0, t1, dep = stack.pop(); boxes += 1
                sc = arb((s0 + s1) / 2); tc = arb((t0 + t1) / 2)
                hs = arb((s1 - s0) / 2); ht = arb((t1 - t0) / 2)
                p = pt(i, sc); q = pt(j, tc)
                dc = sum(((q[a] - p[a])**2 for a in range(3)), arb(0)).sqrt()
                dlo = dc - hs * l[i] - ht * l[j]                    # |d| is Lipschitz in (s,t)
                a1 = cum[j] - cum[i] - arb(s0) * l[i] + arb(t1) * l[j]   # max of forward arc on box
                a2 = tot - (cum[j] - cum[i] - arb(s1) * l[i] + arb(t0) * l[j])
                am = a1 if a1 < a2 else (a2 if a2 < a1 else a1)
                if dlo > 0 and am < D * dlo: continue
                if dep >= max_depth: return False, ("depth", i, j, s0, t0)
                sm = (s0 + s1) / 2; tm = (t0 + t1) / 2
                stack += [(s0, sm, t0, tm, dep+1), (sm, s1, t0, tm, dep+1), (s0, sm, tm, t1, dep+1), (sm, s1, tm, t1, dep+1)]
    return True, boxes

def main():
    z = load_cert(os.path.join(os.path.dirname(os.path.abspath(__file__)), "DIST_trefoil"))
    p, q, R, r, k, U = [int(v) for v in z["meta"]]
    if (p, q, R, r, k) != (2, 3, 8, 5, 6):
        print("REJECTED: parameters"); return
    S = list(z["states"].astype(np.int64)); SR = list(z["states_ref"].astype(np.int64))
    P0, _ = p0_polygon(p, q, R, r, k)
    if not all((c * U).denominator == 1 for v in P0 for c in v):
        print("REJECTED: unit"); return
    if not np.array_equal(np.array([[int(c * U) for c in v] for v in P0], dtype=np.int64), S[0]):
        print("REJECTED: first state"); return
    if not path_ok(S, U):
        print("REJECTED: path"); return
    X = S[-1]; D0 = np.zeros((2 * len(X), 3), dtype=np.int64)
    for v in range(len(X)):
        D0[2 * v] = 2 * X[v]; D0[2 * v + 1] = X[v] + X[(v + 1) % len(X)]
    if not path_ok([D0] + SR, 2 * U):
        print("REJECTED: refined path"); return
    ok, info = check(SR[-1], "6.92")
    if not ok:
        print("REJECTED: distortion", info); return
    print("CERTIFIED (verifier 2): trefoil polygon with %d vertices, distortion <= 6.92 (%d boxes)" % (len(SR[-1]), info))

if __name__ == "__main__":
    main()