"""From the Regge action to the corner couplings of the Kuhn lattice at kappa = 1.

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
  * Builds the Regge action of one Kuhn simplex from Eq. (3): the areas A_h = (1/4)((l1 + l2 - l3)^2 - 4 l1 l2)^(1/2) of its ten
    timelike triangles, the dihedral angle theta_{sigma,h} between the two tetrahedra of the simplex that meet at h (measured in
    the spacelike plane normal to h, from the facet normals), and the number n_h of simplices that share h, which is
    p! r! m! (m + 1) for a triangle whose two chain steps advance p and r counters, m = 4 - p - r.
  * Appendix B. At the flat kappa = 1 lattice a simplex meets its triangles with two consecutive postings at the angles pi/6,
    2pi/3 and pi/6, depending on where they sit in its chain; the six simplices around such a triangle still close up to 2pi.
    With each simplex taking its own flat angle theta0 as its share, S_sigma = sum_h A_h (theta0_{sigma,h} - theta_{sigma,h}),
    the Hessian of one simplex is exactly C/(16 sqrt3), with C the printed integer matrix (checked). With the equal share
    2 pi/n_h the simplex Hessians differ from it by terms (2 pi/n_h - theta0) d^2 A_h on three pairs of edges, which cancel at
    every triangle, so the lattice action and the Bloch matrix B(k) are the same either way (checked at three waves).
  * Propositions 6.3, 6.4 and 7.6. A constant edge profile f changes the squared length of class d by f_d. The first-order shift
    of the zero at the corner pi e_c is delta_c(f) = v_c^dagger (d_f B)(pi e_c) v_c, which is the sum over the 24 simplices of
    a cell of the third derivative of the action along f, w, w, with w the corner mode v_c(d) = 2^(|d|-3) (d_c = 0) carried by the
    Bloch sign of each edge. The program computes it by finite differences in 40-digit arithmetic for the seven backgrounds of
    the gravity page and compares with the paper's exact values; it also recovers Proposition 6.4's coefficients
    (-288, -2592, -2592, 1728)/(16 * 1152 sqrt3) from the four profiles that depend only on |d|, and checks the sum rule
    sum_c delta_c(H) = -396/(1152 sqrt3) tr H of Proposition 6.3 for the strains among them.

Requires Python 3.9+, numpy and mpmath. Run: python3 regge_action.py   [--json shifts.json]
License: MIT. Copyright (c) 2026 Recognition Physics Institute.
"""
import argparse
import itertools
import json
import math
import sys

import mpmath as mp
import numpy as np

mp.mp.dps = 40
PAIRS = [(0, 1), (0, 2), (0, 3), (0, 4), (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
PIDX = {p: i for i, p in enumerate(PAIRS)}
TRIS = list(itertools.combinations(range(5), 3))
# Appendix B: the cell Hessian at the flat kappa = 1 lattice is C/(16 sqrt3), edges in the order of PAIRS.
C = [[2, -2, 1, 0, 0, 3, -2, -6, 2, 2], [-2, -4, 4, -1, 6, -6, 2, 3, -2, 2], [1, 4, 0, -2, -6, 0, 2, -3, 2, -2],
     [0, -1, -2, 2, 2, 2, -2, 2, -1, 0], [0, 6, -6, 2, -6, 6, -3, 0, 3, -6], [3, -6, 0, 2, 6, 0, 0, 6, -6, 3],
     [-2, 2, 2, -2, -3, 0, 0, -6, 4, 1], [-6, 3, -3, 2, 0, 6, -6, -6, 6, 0], [2, -2, 2, -1, 3, -6, 4, 6, -4, -2],
     [2, 2, -2, 0, -6, 3, 1, 0, -2, 2]]
FLAT = [mp.mpf((j - i) * (1 - (j - i))) for i, j in PAIRS]       # squared lengths |d|(1 - |d|) at kappa = 1


def hinge_count(tri):
    i, j, k = tri
    p, r = j - i, k - j
    m = 4 - p - r
    return math.factorial(p) * math.factorial(r) * math.factorial(m) * (m + 1)


NH = [hinge_count(t) for t in TRIS]


def sq(ell, i, j):
    return ell[PIDX[(min(i, j), max(i, j))]]


def angles(ell):
    """Dihedral angles at the ten triangles, from the inward facet normals n_a = e^a, n_0 = -sum_a e^a (dual basis of the edge
    vectors from vertex 0, with the inner product given by the squared lengths)."""
    G = mp.matrix(4, 4)
    for a in range(1, 5):
        for b in range(1, 5):
            G[a - 1, b - 1] = (sq(ell, 0, a) + sq(ell, 0, b) - (0 if a == b else sq(ell, a, b))) / 2
    Gi = G ** -1
    N = mp.matrix(5, 5)
    for i in range(5):
        for j in range(5):
            if i and j:
                N[i, j] = Gi[i - 1, j - 1]
            elif i:
                N[i, j] = -sum(Gi[i - 1, b] for b in range(4))
            elif j:
                N[i, j] = -sum(Gi[a, j - 1] for a in range(4))
            else:
                N[i, j] = sum(Gi[a, b] for a in range(4) for b in range(4))
    out = []
    for tri in TRIS:
        i, j = [v for v in range(5) if v not in tri]
        out.append(mp.acos(-N[i, j] / mp.sqrt(N[i, i] * N[j, j])))
    return out


def areas(ell):
    out = []
    for tri in TRIS:
        l1, l2, l3 = sq(ell, tri[0], tri[1]), sq(ell, tri[1], tri[2]), sq(ell, tri[0], tri[2])
        out.append(mp.sqrt((l1 + l2 - l3) ** 2 - 4 * l1 * l2) / 4)
    return out


THETA0 = angles(FLAT)


def action(ell, share='flat'):
    th, A = angles(ell), areas(ell)
    if share == 'flat':
        return sum(A[t] * (THETA0[t] - th[t]) for t in range(10))
    return sum(A[t] * (2 * mp.pi / NH[t] - th[t]) for t in range(10))


def hessian(share, h=mp.mpf('1e-10')):
    H = [[mp.mpf(0)] * 10 for _ in range(10)]
    S0 = action(FLAT, share)
    for a in range(10):
        for b in range(a, 10):
            def S(da, db):
                e = list(FLAT)
                e[a] += da
                e[b] += db
                return action(e, share)
            if a == b:
                v = (S(h, 0) - 2 * S0 + S(-h, 0)) / h ** 2
            else:
                v = (S(h, h) - S(h, -h) - S(-h, h) + S(-h, -h)) / (4 * h * h)
            H[a][b] = H[b][a] = v
    return H


def class_index(d):
    return int(sum(int(d[a]) << a for a in range(4))) - 1


SIMPLICES = []        # for each of the 24 simplices of a cell: (class index, start vertex) of each edge, in the order of PAIRS
for perm in itertools.permutations(range(4)):
    verts = [np.zeros(4, dtype=int)]
    for a in perm:
        v = verts[-1].copy()
        v[a] = 1
        verts.append(v)
    SIMPLICES.append([(class_index(verts[j] - verts[i]), verts[i].copy()) for i, j in PAIRS])


def bloch(Hs, k):
    """B(k)_{dd'} = sum over the simplices and their edge pairs of H_ff' exp(i k.(b_f' - b_f)) (Appendix B)."""
    M = np.zeros((15, 15), dtype=complex)
    for edges, H in zip(SIMPLICES, Hs):
        for f, (df, bf) in enumerate(edges):
            for g, (dg, bg) in enumerate(edges):
                M[df, dg] += H[f][g] * np.exp(1j * (k @ (bg - bf)))
    return M


def corner_mode(c):
    v = [0.0] * 15
    for m in range(1, 16):
        d = [(m >> a) & 1 for a in range(4)]
        if d[c] == 0:
            v[m - 1] = 2.0 ** (sum(d) - 3)
    return v


def shift(profile, c, h=mp.mpf('1e-8'), tau=mp.mpf('1e-8')):
    """delta_c(f) = sum over simplices of d^3 S / dt du^2 at l0 + t f + u w, by central differences."""
    v = corner_mode(c)
    total = mp.mpf(0)
    for edges in SIMPLICES:
        f = [mp.mpf(profile[d]) for d, _ in edges]
        w = [mp.mpf(v[d]) * (-1) ** int(b[c]) for d, b in edges]

        def S(t, u):
            return action([FLAT[e] + t * f[e] + u * w[e] for e in range(10)])

        def g(t):
            return (S(t, h) - 2 * S(t, 0) + S(t, -h)) / h ** 2

        total += (g(tau) - g(-tau)) / (2 * tau)
    return total


def strain_profile(H):
    out = []
    for m in range(1, 16):
        d = [(m >> a) & 1 for a in range(4)]
        out.append(sum(d[a] * H[a][b] * d[b] for a in range(4) for b in range(4)))
    return out


R3 = mp.sqrt(3)
ONES = [[1] * 4 for _ in range(4)]
G1 = [[(1 if i == j else 0) - 1 for j in range(4)] for i in range(4)]
BACKGROUNDS = [   # name, profile, the paper's shifts at c = 0..3, trace of the strain (None if not a strain), source
    ('the flat metric g1', strain_profile(G1), [0, 0, 0, 0], 0, 'Proposition 6.3(i)'),
    ('the strain 11^T', strain_profile(ONES), [-mp.mpf(11) / (32 * R3)] * 4, 4, 'Proposition 6.3(iii)'),
    ('the probe H_E', strain_profile([[0, 2, -1, -1], [2, 0, -1, -1], [-1, -1, 0, 2], [-1, -1, 2, 0]]), [0, 0, 0, 0], 0,
     'Proposition 6.3(i)'),
    ('the probe H_T2', strain_profile([[-1, 1, 0, 0], [1, -1, 0, 0], [0, 0, 1, -1], [0, 0, -1, 1]]),
     [-R3 / 32, -R3 / 32, R3 / 32, R3 / 32], 0, 'Proposition 6.3(ii)'),
    ('the strain H01 = 1, H23 = -1', strain_profile([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, -1], [0, 0, -1, 0]]),
     [1 / (16 * R3), 1 / (16 * R3), -1 / (16 * R3), -1 / (16 * R3)], 0, 'Proposition 6.3(v)'),
    ('the polarization of a resting mass, (3/4) 11^T + (1/2) g1', strain_profile(
        [[mp.mpf(3) / 4 + (mp.mpf(1) / 2) * G1[i][j] for j in range(4)] for i in range(4)]),
     [-mp.mpf(33) / (128 * R3)] * 4, 3, 'Proposition 7.6'),
    ('the body-diagonal profile (not a strain)', [0] * 14 + [1], [R3 / 32] * 4, None, 'Proposition 6.4'),
]


def report(name, ok, detail):
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}: {detail}")
    return ok


def main():
    parser = argparse.ArgumentParser(description='From the Regge action of a Kuhn simplex to Appendix B and the corner shifts.')
    parser.add_argument('--json', metavar='FILE.json', help='write the computed corner shifts')
    args = parser.parse_args()
    ok = True
    print('Eq. (3): the action of one Kuhn simplex at the flat kappa = 1 lattice')
    for t, tri in enumerate(TRIS):
        if tri in ((0, 1, 2), (1, 2, 3), (2, 3, 4)):
            print(f'  [INFO] triangle {tri}: flat dihedral angle {mp.nstr(THETA0[t] / mp.pi, 6)} pi, shared by n_h = {NH[t]} simplices')
    closes = abs(2 * THETA0[0] + 2 * THETA0[TRIS.index((1, 2, 3))] + 2 * THETA0[TRIS.index((2, 3, 4))] - 2 * mp.pi) < mp.mpf('1e-30')
    ok &= report('the six simplices around a triangle with two consecutive postings close up to 2 pi', closes,
                 '2(pi/6) + 2(2pi/3) + 2(pi/6) = 2 pi')

    print('Appendix B: the cell Hessian')
    Hf = hessian('flat')
    dev = max(abs(Hf[a][b] * 16 * R3 - C[a][b]) for a in range(10) for b in range(10))
    ok &= report('Hessian of one simplex with the flat-angle share = C/(16 sqrt3)', dev < mp.mpf('1e-12'),
                 f'max |16 sqrt3 H - C| = {mp.nstr(dev, 3)}')
    He = hessian('equal')
    diff = [(PAIRS[a], PAIRS[b], He[a][b] * 16 * R3 - C[a][b]) for a in range(10) for b in range(a, 10)
            if abs(He[a][b] * 16 * R3 - C[a][b]) > mp.mpf('1e-10')]
    expected = {((0, 1), (1, 2)): -2 * mp.pi / R3, ((1, 2), (2, 3)): 4 * mp.pi / R3, ((2, 3), (3, 4)): -2 * mp.pi / R3}
    match = len(diff) == 3 and all(abs(x - expected.get((p, q), 0)) < mp.mpf('1e-12') for p, q, x in diff)
    ok &= report('with the equal share 2 pi/n_h, the simplex Hessian differs only by the hinge terms', match,
                 ', '.join(f'{p}-{q}: {mp.nstr(x / (16 * R3) / mp.pi, 6)} pi' for p, q, x in diff))
    Hf_np = [[float(Hf[a][b]) for b in range(10)] for a in range(10)]
    He_np = [[float(He[a][b]) for b in range(10)] for a in range(10)]
    worst = 0.0
    for k in (np.array([0.3, -0.7, 1.1, 0.4]), np.array([math.pi, 0, 0, 0]), np.array([1.0, 2.0, -0.5, 0.2])):
        worst = max(worst, float(np.abs(bloch([He_np] * 24, k) - bloch([Hf_np] * 24, k)).max()))
    ok &= report('the hinge terms cancel: B(k) is the same with either share', worst < 1e-12, f'max difference {worst:.1e} at three waves')

    print('Propositions 6.3, 6.4 and 7.6: first-order shifts of the four corner zeros, from the third variation of the action')
    results = {}
    for name, prof, paper, trace, source in BACKGROUNDS:
        got = [shift(prof, c) for c in range(4)]
        dev = max(abs(g - p) for g, p in zip(got, paper))
        ok &= report(f'{name} ({source})', dev < mp.mpf('1e-12'),
                     f'{", ".join(mp.nstr(g, 8) for g in got)}; paper {", ".join(mp.nstr(mp.mpf(p), 8) for p in paper)}')
        if trace is not None:
            rule = -mp.mpf(396) / (1152 * R3) * trace
            ok &= report('  sum rule, sum_c delta_c = -396/(1152 sqrt3) tr H', abs(sum(got) - rule) < mp.mpf('1e-12'),
                         f'{mp.nstr(sum(got), 8)} = {mp.nstr(rule, 8)}')
        results[name] = [float(g) for g in got]
    coef = []
    for n, paper in zip((1, 2, 3, 4), (-288, -2592, -2592, 1728)):
        prof = [1 if bin(m).count('1') == n else 0 for m in range(1, 16)]
        got = shift(prof, 0) * 16 * 1152 * R3
        coef.append(got)
        ok &= report(f'Proposition 6.4: the profile 1 on |d| = {n}', abs(got - paper) < mp.mpf('1e-6'),
                     f'16 * 1152 sqrt3 * delta = {mp.nstr(got, 9)} (paper {paper})')
    results['proposition_6_4_coefficients'] = [float(x) for x in coef]
    if args.json:
        with open(args.json, 'w') as fh:
            json.dump(results, fh)
        print(f'  wrote {args.json}')
    print('ALL CHECKS PASS' if ok else 'SOME CHECKS FAILED')
    return ok


if __name__ == '__main__':
    sys.exit(0 if main() else 1)
