"""Exact dispersion of the kappa = 1 Kuhn lattice along the paired direction, from Appendix F, with checks of Theorem 5.5.

Paper: "Linearized Einstein gravity from Lorentzian Regge calculus on the Kuhn lattice, and the divergence of its one-loop
graviton limit", J. Washburn (2026), https://doi.org/10.5281/zenodo.23764491. Section, theorem and appendix numbers below are
the paper's.

What this program does
  Along the paired covector u = (z+t, z+t, z-t, z-t) the spatial wave number is |k| = 2t and the frequency is omega = 2z/sqrt3
  (Section 5.3); the light cone is z = sqrt3 t. With U = sin^2(t/2) and V = sin^2(z/2) the cross mode is the exact curve
  G_x(U, V) = -6UV + 3U + 4V^2 - V = 0 and the plus mode lies on the degree-six curve G_+(U, V) = 0 printed in Appendix F.
  The program follows, in 40-digit arithmetic, the branch of each curve that leaves the origin along the light cone, and checks
  and prints:
    * the limits of (v_ph - 1)/|k|^2: -59/24 (plus) and 5/24 (cross), and of (v_g - 1)/|k|^2: -59/8 and 5/8 (Theorem 5.5);
    * the windows of Theorem 5.5, |v_ph - 1 - c|k|^2| < (25/16)|k|^3, for 0 < t <= 2e-4 near the cone;
    * where each branch ends: the cross curve turns back where the discriminant of G_x vanishes, U = (sqrt2 - 1)^2/6,
      z = pi/4, |k| = 0.6797; the plus curve reaches omega = 0 where G_+(U, 0) = -2U(4U^2 - 10U + 3) vanishes, |k| = 2.5264.
  With --json FILE.json it writes the samples drawn in the figure "Exact dispersion along the paired direction" on the
  gravity page: the plus branch on a 0.005 grid in |k| and both roots of G_x.

Requires Python 3.9+ and mpmath. Run: python3 paired_dispersion.py   or   python3 paired_dispersion.py --json branches.json
License: MIT. Copyright (c) 2026 Recognition Physics Institute.
"""
import argparse
import json
import sys

from mpmath import asin, diff, mp, mpf, pi, polyroots, polyval, sin, sqrt

mp.dps = 40


def g_cross(U):
    """G_x as a polynomial in V, highest degree first: 4V^2 - (6U + 1)V + 3U."""
    return [mpf(4), -(6 * U + 1), 3 * U]


def g_plus(U):
    """G_+ of Appendix F as a polynomial in V, highest degree first."""
    return [mpf(1024),
            -512 * U - 2816,
            -512 * U**2 + 1792 * U + 2400,
            128 * U**3 + 832 * U**2 - 1584 * U - 648,
            -192 * U**3 - 432 * U**2 + 280 * U + 44,
            80 * U**3 + 88 * U**2 + 20 * U + 2,
            -8 * U**3 + 20 * U**2 - 6 * U]


def group_speed(poly, t, z):
    """Appendix F: v_g = -G_U sin t / (sqrt3 G_V sin z), with U = sin^2(t/2), V = sin^2(z/2)."""
    U, V = sin(t / 2) ** 2, sin(z / 2) ** 2
    G = lambda u, v: polyval(poly(u), v)
    return -diff(lambda u: G(u, V), U) * sin(t) / (sqrt(3) * diff(lambda v: G(U, v), V) * sin(z))


def _real_roots(coeffs):
    roots = polyroots(coeffs, maxsteps=200, extraprec=60)
    tiny = mpf(10) ** -25
    return [r.real for r in roots if abs(r.imag) < tiny and -tiny <= r.real <= 1 + tiny]


def trace(poly, eta1, t_grid):
    """Follow the root V(t) that starts on the light cone as V = 3U(1 + eta1 U); stop where it turns back or ends."""
    out, prev = [], None
    for t in t_grid:
        U = sin(t / 2) ** 2
        guess = 3 * U * (1 + eta1 * U) if prev is None else prev
        roots = _real_roots(poly(U))
        if not roots:
            break
        V = min(roots, key=lambda r: abs(r - guess))
        if prev is not None and abs(V - prev) > mpf('0.05'):
            break
        prev = V
        out.append((2 * t, 4 * asin(sqrt(V)) / sqrt(3)))       # (|k|, omega) = (2t, 2z/sqrt3)
    return out


def cross_closed_form(n=400):
    """Both roots of G_x for U in [0, U_c], U_c = (sqrt2 - 1)^2/6 where the discriminant 36U^2 - 36U + 1 vanishes."""
    Uc = (sqrt(2) - 1) ** 2 / 6
    tc = 2 * asin(sqrt(Uc))
    lower, upper = [], []
    for i in range(n + 1):
        t = tc * i / n
        U = sin(t / 2) ** 2
        D = max(36 * U * U - 36 * U + 1, mpf(0))
        Vm = 6 * U / ((6 * U + 1) + sqrt(D))       # rationalized lower root: no cancellation near the cone
        Vp = ((6 * U + 1) + sqrt(D)) / 8
        lower.append((2 * t, 4 * asin(sqrt(Vm)) / sqrt(3)))
        upper.append((2 * t, 4 * asin(sqrt(Vp)) / sqrt(3)))
    return lower, upper[::-1], (2 * tc, pi / (2 * sqrt(3)))


def plus_static_zero():
    """|k| where the plus branch meets omega = 0: G_+(U, 0) = -2U(4U^2 - 10U + 3), U = (5 - sqrt13)/4."""
    return 4 * asin(sqrt((5 - sqrt(13)) / 4))


def figure_data(step=mpf('0.005')):
    grid = [step * i / 2 for i in range(1, int(mpf('2.53') / step) + 1)]          # t = |k|/2
    k_static = plus_static_zero()
    plus = [(mpf(0), mpf(0))] + [p for p in trace(g_plus, mpf(-238) / 3, grid) if p[0] < k_static] + [(k_static, mpf(0))]
    lower, upper, turn = cross_closed_form()
    return plus, lower, upper, turn


def run_checks():
    ok = True
    small = [mpf(10) ** (-6 + 5 * i / 200) for i in range(200)]       # t from 1e-6 to 1e-1
    for name, poly, eta1, c_ph, c_g in (('plus', g_plus, mpf(-238) / 3, mpf(-59) / 24, mpf(-59) / 8),
                                        ('cross', g_cross, mpf(6), mpf(5) / 24, mpf(5) / 8)):
        branch = trace(poly, eta1, small)
        k0, w0 = branch[0]
        lim_ph = (w0 / k0 - 1) / k0**2
        lim_g = (group_speed(poly, k0 / 2, w0 * sqrt(3) / 2) - 1) / k0**2
        window = max(abs(w / k - 1 - c_ph * k**2) / (mpf(25) / 16 * k**3) for k, w in branch if k / 2 <= mpf('2e-4'))
        good = abs(lim_ph - c_ph) < mpf('1e-6') and abs(lim_g - c_g) < mpf('1e-4') and window < 1
        ok &= good
        print(f"  [{'PASS' if good else 'FAIL'}] {name}: (v_ph - 1)/|k|^2 -> {float(lim_ph):.9f} (Theorem 5.5: {float(c_ph):.9f}); "
              f"(v_g - 1)/|k|^2 -> {float(lim_g):.6f} ({float(c_g):.6f}); window ratio {float(window):.2e} < 1")
    plus, lower, upper, turn = figure_data()
    print(f'  [INFO] cross branch turns back at |k| = {float(turn[0]):.6f}, omega = {float(turn[1]):.6f} = pi/(2 sqrt3), z = pi/4')
    print(f'  [INFO] plus branch reaches omega = 0 at |k| = {float(plus_static_zero()):.6f}; '
          f'its highest frequency on the 0.005 grid is {float(max(w for _, w in plus)):.4f}')
    sample = [(k, w) for k, w in plus if abs(k - mpf('0.4')) < mpf('1e-9')][0]
    print(f'  [INFO] at |k| = 0.4: plus omega = {float(sample[1]):.4f}, phase speed {float(sample[1] / sample[0]):.4f}')
    return ok, (plus, lower, upper, turn)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Exact paired-direction dispersion from Appendix F, checked against Theorem 5.5.')
    parser.add_argument('--json', metavar='FILE.json', help='write the samples drawn in the figure')
    args = parser.parse_args()
    print('Theorem 5.5 and Appendix F: the paired direction')
    good, (plus, lower, upper, turn) = run_checks()
    if args.json:
        data = {'variables': '|k| = 2t, omega = 2z/sqrt3, lattice units',
                'plus': [[float(k), float(w)] for k, w in plus],
                'cross_near_cone': [[float(k), float(w)] for k, w in lower],
                'cross_returning': [[float(k), float(w)] for k, w in upper],
                'cross_turning_point': [float(turn[0]), float(turn[1])]}
        with open(args.json, 'w') as fh:
            json.dump(data, fh)
        print(f'  wrote {args.json}')
    print('ALL CHECKS PASS' if good else 'SOME CHECKS FAILED')
    sys.exit(0 if good else 1)
