#!/usr/bin/env python3
"""
qg_regge_triangulation_d2.py -- D2 companion: an ACTUAL Regge triangulation
of the (Euclidean) Schwarzschild exterior.

Closes the open engineering task flagged in scripts/qg_strong_theory_derivations.py
(derivation D2) and in the Beltracchi audit: build a genuine simplicial mesh of
a Schwarzschild window, compute metric edge lengths, GEOMETRIC dihedral angles
and deficits, hinge areas, and assemble both actions

    S_Regge = (1/8 pi G) sum_sigma A_sigma delta_sigma
    S_RS    = (1/8 pi G) sum_sigma A_sigma sinh(delta_sigma)

then MEASURE the convergence exponent of the density |S_RS - S_Regge| / Vol
under uniform mesh refinement, instead of asserting h^4 by power counting.

Construction
------------
* Euclidean Schwarzschild, G = c = 1, M = 1, r_s = 2:
      ds^2 = f dtau^2 + f^-1 dr^2 + r^2 dtheta^2 + r^2 sin^2(theta) dphi^2,
      f = 1 - r_s/r  (f > 0 everywhere in the window; no coordinate issue).
* Window: a coordinate box centered at (tau, r, theta, phi) = (0, r0, pi/2, 0)
  whose coordinate extents are chosen so the PROPER extent is s in every
  direction (evaluated at the center) -- an approximately proper-isotropic
  box. Refinement keeps the box fixed and refines the grid: h = s/n.
* Mesh: uniform n^4 grid of 4-cells, each split into 24 4-simplices by the
  Kuhn/Freudenthal triangulation (compatible across neighboring cells). 
* Edge lengths: proper length of the coordinate segment (6-point Gauss
  quadrature; machine-accurate for these smooth integrands).
* Dihedral angles: each 4-simplex is embedded isometrically in R^4 from its
  10 edge lengths (Gram matrix + Cholesky); the angle at a hinge (triangle)
  is the angle between the projections of the two off-hinge vertices onto
  the 2-plane orthogonal to the hinge. Deficit = 2*pi - sum over the star.
  Only hinges whose full star is present in the mesh (all three vertices
  strictly interior to the grid) are measured.
* Volume attribution: each 4-simplex donates V/10 to each of its 10 hinges,
  so densities are ratios over the SAME interior-hinge set (insensitive to
  the boundary fraction, which changes with n).

Honest scope
------------
* Edge lengths are metric chord lengths (proper length along coordinate
  segments), not exact geodesic distances. Both the true curvature deficit
  and the chordal approximation error are O(K h^2), so the measured deficits
  carry an O(1) mesh/method-dependent coefficient; the SCALING exponent is
  the invariant content. The flat-space null test (all interior deficits at
  rounding level) confirms the O(h^2) deficit content vanishes with
  curvature, i.e. there is no curvature-independent artifact.
* This measures |S_RS - S_Regge| (the RS cubic correction), which is the D2
  claim. Regge -> Einstein-Hilbert convergence itself remains the external
  Cheeger-Mueller-Schrader, although we can examine it through C1 = sum A delta / Vol as a sanity check. 

Usage
-----
    python3 scripts/qg_regge_triangulation_d2.py            # full run
    python3 scripts/qg_regge_triangulation_d2.py --quick    # smaller levels
"""

import argparse
import math
import time
from itertools import combinations, permutations

import numpy as np

# ----------------------------------------------------------------------
# constants and combinatorics
# ----------------------------------------------------------------------
R_S = 2.0  # Schwarzschild radius, M = 1, G = c = 1
kappa = 8.0 * np.pi  # EH action factor, G = c = 1
_gl_nodes, _gl_weights = np.polynomial.legendre.leggauss(6)
GL_T = (_gl_nodes + 1.0) / 2.0   # quadrature nodes on [0, 1]
GL_W = _gl_weights / 2.0

HINGES = list(combinations(range(5), 3))                       # 10 triangles
OFFV = [tuple(sorted(set(range(5)) - set(h))) for h in HINGES]  # 2 off-vertices


# ----------------------------------------------------------------------
# metrics (diagonal, vectorized): return g_dd components at points X (...,4)
# coordinate order: (tau, r, theta, phi)
# ----------------------------------------------------------------------
def gdiag_schwarzschild(X):
    r = X[..., 1]
    th = X[..., 2]
    f = 1.0 - R_S / r
    out = np.empty_like(X)
    out[..., 0] = f
    out[..., 1] = 1.0 / f
    out[..., 2] = r * r
    out[..., 3] = (r * np.sin(th)) ** 2
    return out


def gdiag_flat(X):
    return np.ones_like(X)


def edge_lengths(XA, XB, gfun):
    """Proper length of the coordinate segment XA->XB (batched, (N,4))."""
    D = XB - XA
    acc = np.zeros(XA.shape[0])
    for t, w in zip(GL_T, GL_W):
        g = gfun(XA + t * D)
        acc += w * np.sqrt(np.einsum('nd,nd->n', g, D * D))
    return acc


def sinh_minus_linear(d):
    """sinh(d) - d, series-stabilized for small d."""
    small = np.abs(d) < 1e-3
    series = d ** 3 / 6.0 * (1.0 + d * d / 20.0)
    with np.errstate(over='ignore'):
        direct = np.sinh(d) - d
    return np.where(small, series, direct)


# ----------------------------------------------------------------------
# one refinement level
# ----------------------------------------------------------------------
def run_level(n, lo, hi, gfun, cell_chunk=6000, verbose=True):
    """
    Triangulate the coordinate box [lo, hi] with an n^4 Kuhn mesh, measure
    interior deficits, and return density observables.
    """
    t0 = time.time()
    Vn = n + 1                       # vertices per dimension
    M = Vn ** 4                      # vertex-id modulus for key packing
    scale = (hi - lo) / n            # coordinate step per dimension

    # Kuhn simplex vertex offsets: one path per permutation of the 4 axes
    offsets = []
    for p in permutations(range(4)):
        o = np.zeros((5, 4), np.int64)
        for k in range(4):
            o[k + 1] = o[k]
            o[k + 1, p[k]] += 1
        offsets.append(o)

    cells = np.stack(
        np.meshgrid(*[np.arange(n)] * 4, indexing='ij'), -1
    ).reshape(-1, 4).astype(np.int64)

    keys_l, ang_l, area_l, vsh_l, rads_l = [], [], [], [], []

    for o in offsets:
        for c0 in range(0, len(cells), cell_chunk):
            G = cells[c0:c0 + cell_chunk, None, :] + o[None, :, :]   # (m,5,4) grid
            m = G.shape[0]
            X = lo[None, None, :] + G * scale[None, None, :]         # (m,5,4) coords

            # --- 10 pairwise edge lengths per simplex
            Lm = np.zeros((m, 5, 5))
            for i in range(5):
                for j in range(i + 1, 5):
                    L = edge_lengths(X[:, i, :], X[:, j, :], gfun)
                    Lm[:, i, j] = L
                    Lm[:, j, i] = L

            # --- isometric embedding in R^4 (v0 at origin)
            d0 = Lm[:, 0, 1:]                                        # (m,4)
            B = 0.5 * (d0[:, :, None] ** 2 + d0[:, None, :] ** 2
                       - Lm[:, 1:, 1:] ** 2)
            try:
                Xc = np.linalg.cholesky(B)
            except np.linalg.LinAlgError:
                w, U = np.linalg.eigh(B)
                w = np.clip(w, 0.0, None)
                Xc = U * np.sqrt(w)[:, None, :]
            P = np.zeros((m, 5, 4))
            P[:, 1:, :] = Xc
            vol = np.abs(np.linalg.det(Xc)) / 24.0                   # 4-volume

            gv = (((G[..., 0] * Vn + G[..., 1]) * Vn
                   + G[..., 2]) * Vn + G[..., 3])                    # (m,5) vertex ids

            # --- dihedral angle + area at each of the 10 hinges
            for hidx, (a, b, c) in enumerate(HINGES):
                r_hinge = (X[:, a, 1] + X[:, b, 1] + X[:, c, 1]) / 3
                p_, q_ = OFFV[hidx]
                va = P[:, a, :]
                u1 = P[:, b, :] - va
                u2 = P[:, c, :] - va
                w1 = P[:, p_, :] - va
                w2 = P[:, q_, :] - va

                n1 = np.linalg.norm(u1, axis=1)
                e1 = u1 / n1[:, None]
                u2p = u2 - np.einsum('nd,nd->n', u2, e1)[:, None] * e1
                n2 = np.linalg.norm(u2p, axis=1)
                e2 = u2p / n2[:, None]

                def perp(wv):
                    return (wv
                            - np.einsum('nd,nd->n', wv, e1)[:, None] * e1
                            - np.einsum('nd,nd->n', wv, e2)[:, None] * e2)

                t1 = perp(w1)
                t2 = perp(w2)
                ct = (np.einsum('nd,nd->n', t1, t2)
                      / (np.linalg.norm(t1, axis=1) * np.linalg.norm(t2, axis=1)))
                ang = np.arccos(np.clip(ct, -1.0, 1.0))

                uu = np.einsum('nd,nd->n', u1, u1)
                vv = np.einsum('nd,nd->n', u2, u2)
                uv = np.einsum('nd,nd->n', u1, u2)
                area = 0.5 * np.sqrt(np.maximum(uu * vv - uv * uv, 0.0))#comment on the area formula: (A.A)(B.B)-(A.B)^2=|A|^2|B|^2-|A|^2|B|^2cos^2(\theta)=|A|^2|B|^2sin^2(\theta), 
                # which is a standard parallelogram spanned by two vectors formula
                tri = np.sort(gv[:, [a, b, c]], axis=1).astype(np.int64)
                key = (tri[:, 0] * M + tri[:, 1]) * M + tri[:, 2]
                rads_l.append(r_hinge)
                keys_l.append(key)
                ang_l.append(ang)
                area_l.append(area)
                vsh_l.append(vol / 10.0)
    rads=np.concatenate(rads_l)
    keys = np.concatenate(keys_l)
    angs = np.concatenate(ang_l)
    areas = np.concatenate(area_l)
    vshs = np.concatenate(vsh_l)

    uk, first, inv = np.unique(keys, return_index=True, return_inverse=True)
    ang_sum = np.bincount(inv, weights=angs)
    v_sum = np.bincount(inv, weights=vshs)
    area_u = areas[first]
    rads_u = rads[first]
    # interior test: all three vertices strictly inside the grid
    def _interior(v):
        g3 = v % Vn
        t = v // Vn
        g2 = t % Vn
        t //= Vn
        g1 = t % Vn
        g0 = t // Vn
        return ((g0 >= 1) & (g0 <= n - 1) & (g1 >= 1) & (g1 <= n - 1)
                & (g2 >= 1) & (g2 <= n - 1) & (g3 >= 1) & (g3 <= n - 1))

    v2 = uk % M
    t = uk // M
    v1 = t % M
    v0 = t // M
    mask = _interior(v0) & _interior(v1) & _interior(v2)

    delta = 2.0 * np.pi - ang_sum[mask]
    A = area_u[mask]
    Vs = v_sum[mask]
    rads = rads_u[mask]
    res = sinh_minus_linear(delta)

    out = {
        'n': n,
        'n_int': int(mask.sum()),
        'med_abs_delta': float(np.median(np.abs(delta))) if mask.any() else float('nan'),
        'max_abs_delta': float(np.max(np.abs(delta))) if mask.any() else float('nan'),
        'med_deltasur':float(np.median(A/2*np.sqrt(48)/rads**3)) if mask.any() else float('nan'),# This is a Kretshcmann surrogate A*sqrt(K) = A*sqrt(48)/r^3,
        # which is related to the condition 1/2 R_abcd n^ab n^cd A\approx\delta, where n^ab is the normal bivector to the hinge. 
        # The square root of the Kretschmann scalar acts as a surrogate for the curvature so we can avoid computing the area bivector contracting with the Riemann tensor on 
        # each hinge. The factor of sqrt(48) comes from the fact that for Schwarzschild, K = 48 M^2 / r^6, and we set M=1.
        'max_deltasur':float(np.max(A/2*np.sqrt(48)/rads**3)) if mask.any() else float('nan'),
        'dens_abs': float(np.sum(A * np.abs(res)) /(kappa* np.sum(Vs))),
        'dens_signed': float(np.sum(A * res) / (kappa * np.sum(Vs))),
        'regge_dens': float(np.sum(A * delta) / (kappa * np.sum(Vs))),
        'seconds': time.time() - t0,
    }
    if verbose:
        print(f"    n = {n:3d}: interior hinges = {out['n_int']:8d}, "
              f"median|delta| = {out['med_abs_delta']:.3e}, "
              f"max|delta| = {out['max_abs_delta']:.3e}   "
              f"({out['seconds']:.1f} s)")
    return out


# ----------------------------------------------------------------------
# window construction: proper-isotropic coordinate box around r0
# ----------------------------------------------------------------------
def window(r0, s):
    """Coordinate box centered at (0, r0, pi/2, 0) with proper extent s per dim."""
    f0 = 1.0 - R_S / r0
    ext = np.array([s / math.sqrt(f0),     # dtau
                    s * math.sqrt(f0),     # dr
                    s / r0,                # dtheta
                    s / r0])               # dphi  (sin(pi/2) = 1)
    center = np.array([0.0, r0, math.pi / 2.0, 0.0])
    return center - ext / 2.0, center + ext / 2.0


def fit_exponent(hs, ds):
    """Least-squares slope of log(d) vs log(h)."""
    return float(np.polyfit(np.log(hs), np.log(ds), 1)[0])


# ----------------------------------------------------------------------
# main
# ----------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--levels', default='4,6,8,10,12,14,16',
                    help='comma-separated grid sizes n (h = s/n)')
    ap.add_argument('--s', type=float, default=0.8,
                    help='proper window size in units of r_s')
    ap.add_argument('--quick', action='store_true',
                    help='use levels 4,6,8 only')
    args = ap.parse_args()

    levels = [4, 6, 8] if args.quick else [int(x) for x in args.levels.split(',')]
    s = args.s * R_S

    print("=" * 72)
    print("  D2 COMPANION: MEASURED REGGE TRIANGULATION OF SCHWARZSCHILD")
    print("  (geometric deficits; the h^4 law is MEASURED here, not assumed)")
    print("=" * 72)
    print(f"  r_s = {R_S} (M = 1, G = c = 1), proper window s = {args.s} r_s")
    print(f"  levels n = {levels}  (mesh scale h = s/n)")

    # ---------------- flat-space null test ----------------
    print("\nFLAT-SPACE NULL TEST (identity metric, same pipeline):")
    lo = np.zeros(4)
    hi = np.full(4, s)
    flat = run_level(6, lo, hi, gdiag_flat, verbose=False)
    print(f"    n = 6: interior hinges = {flat['n_int']}, "
          f"max|delta| = {flat['max_abs_delta']:.2e}")
    ok = flat['max_abs_delta'] < 1e-9
    print(f"    [{'PASS' if ok else 'FAIL'}] interior deficits vanish in flat space "
          f"(threshold 1e-9)")
    if not ok:
        print("    Pipeline is broken; aborting before curved runs.")
        return

    # ---------------- Schwarzschild series ----------------
    for r0_rs, label in [(2.0, "far window, r0 = 2.0 r_s"),
                         (1.5, "near-horizon window, r0 = 1.5 r_s")]:
        r0 = r0_rs * R_S
        lo, hi = window(r0, s)
        K = 48.0 / r0 ** 6                       # Kretschmann, M = 1
        print(f"\nSCHWARZSCHILD SERIES ({label}):")
        print(f"  r in [{lo[1]/R_S:.3f}, {hi[1]/R_S:.3f}] r_s, "
              f"Kretschmann K = {K:.4e}, sqrt(K) = {math.sqrt(K):.4e}")
        results = []
        for n in levels:
            results.append(run_level(n, lo, hi, gdiag_schwarzschild))

        hs = np.array([s / r['n'] for r in results])
        med = np.array([r['med_abs_delta'] for r in results])
        dens = np.array([r['dens_abs'] for r in results])

        print("\n  h/r_s      median|delta|   median|sur|   max|sur|    C3(abs)"
              "   C2(signed)        C1(Regge)")
        for r, h in zip(results, hs):
            print(f"  {h/R_S:.4f}     {r['med_abs_delta']:.4e}      "
                  f"{r['med_deltasur']:.4e}     {r['max_deltasur']:.4e}"
                  f"     {r['dens_abs']:.4e}   {r['dens_signed']:+.4e}     {r['regge_dens']:+.4e}")

        # pairwise doubling ratios where levels double
        print("\n  Doubling ratios log2[dens(h)/dens(h/2)]:")
        by_n = {r['n']: r for r in results}
        for n in levels:
            if 2 * n in by_n:
                ratio = by_n[n]['dens_abs'] / by_n[2 * n]['dens_abs']
                print(f"    n = {n:2d} -> {2*n:2d}:  {math.log2(ratio):+.3f}")

        alpha_dens = fit_exponent(hs, dens)
        alpha_delta = fit_exponent(hs, med)
        print(f"\n  MEASURED exponents (least-squares over all levels):")
        print(f"    median|delta|        ~ h^{alpha_delta:.2f}   (power counting: 2)")
        print(f"    |S_RS - S_Regge|/Vol ~ h^{alpha_dens:.2f}   (power counting: 4)")

    print("\n" + "=" * 72)
    print("  Interpretation:")
    print("  * The deficit exponent ~2 and residual-density exponent ~4 are the")
    print("    D2 power-counting claims, now measured on a real triangulation")
    print("    with geometric dihedral angles, not generated by construction.")
    print("  * Coefficients are mesh/method dependent (chord-length edges);")
    print("    the exponent is the invariant content. Flat-space null test")
    print("    confirms zero curvature-independent artifact.")
    print("  * Regge -> EH convergence itself stays the external CMS input.")
    print("=" * 72)


if __name__ == '__main__':
    main()
