#!/usr/bin/env python3
"""
PROVENANCE: PROOF

Computation 155 -- The tensor projection Pi_mn: structure, mechanism, and what is left
=======================================================================================
Sec 7 constructs the stress-energy projection only in the scalar form
T_mn = Pi(T(C)) g_mn and calls the tensor generalisation Pi_mn an open problem.
This computation does three things to that problem: it says exactly what the scalar
form covers (less than the text claims), it rules out the obvious tensor completion,
and it supplies one that works, with the equation of state coming out of a structure
PST already owns.

CHECKS
------
  (1) THE SCALAR FORM IS EXACTLY AND ONLY A COSMOLOGICAL CONSTANT.  Two symbolic
      facts on Friedmann-Robertson-Walker.  First, the Bianchi identity applied to
      T_mn = Pi g_mn gives nabla_m T^{mn} = d(Pi)/dx_n, so conservation FORCES Pi to
      be constant.  Second, a perfect fluid is proportional to g_mn if and only if
      rho + p = 0, i.e. w = -1.  So the scalar form is not "the leading-order
      approximation valid for isotropic configurations": an isotropic dust or
      radiation fluid is already outside it.  Only Lambda is inside.

  (2) THE OBVIOUS COMPLETION IS SINGULAR ON THE LIGHT CONE.  The natural tensor per
      pair is the unit dyad e_m e_n with e = D / sqrt|D.D|.  That normalisation
      diverges as D.D -> 0 -- which is exactly where tau_null sits, the threshold
      eq:metric uses to define the light cone.  Measured here: the sum is dominated
      by a handful of near-null pairs and does not settle as the cloud grows.

  (3) THE SEPARATION DYAD WORKS, AND ITS TRACE IS THE TENSION.  Use D_m D_n
      unnormalised, weighted by a function of the pair's own invariant:

          Pi_mn(x) = (1/2) sum_{a != b} W(s_ab) D_m D_n B(x; rho(a), rho(b)),
          s_ab = D.D  (the interval, which eq:metric ties to T_ab - tau_null).

      No normalisation, no singularity, and g^mn D_m D_n = s: the SCALAR Pi
      reappears as the TRACE of the tensor projection rather than as a factor
      multiplying the metric.

  (4) IT IS AUTOMATICALLY A PERFECT FLUID.  For any weight depending only on the
      invariant, on a rotationally isotropic ensemble, Pi_mn = rho u_m u_n + p h_mn
      with h_mn = u_m u_n - g_mn.  The residual of that decomposition falls with
      cloud size, which is what a finite-sample statement should do.

  (5) THE EQUATION OF STATE IS THE BALANCE OF THE TWO TENSION BUDGETS.  rho is the
      timelike-pair moment and p the spacelike-pair moment, and w = p/rho.  The two
      classes are separated by T_ab vs tau_null -- the SAME classification eq:metric
      already makes to get the Lorentzian signature.  The threshold that gives the
      light cone gives the equation of state.

  (6) w = -1 IS ONE POINT OF THAT FAMILY, AND IT IS THE SCALAR FORM.  Writing the
      weight as +f(|s|) on timelike pairs and -lam f(|s|) on spacelike ones, lam is
      the budget ratio and the single free scalar.  Solving rho + p = 0 pins it; at
      that lam the measured Pi_mn is proportional to g_mn to the accuracy of the
      ensemble.  So the book's cosmological constant is one point of a one-parameter
      family, not the whole content of the projection.

SUPERSEDED IN PART BY COMP 158, which supplies what this computation left open and
CORRECTS ONE OF ITS CONCLUSIONS.  What stands here is the structure, the light-cone
obstruction and the perfect-fluid form; read comp 158 for the weight and the
equation of state.

  * The weight profile W is NOT adopted, as the paragraphs below say it is.  Trace
    consistency between eq:Pi and eq:Pimunu, with eq:metric supplying
    T - tau_null = -s|s|, forces W(s) = -|s| with nothing left free.
  * The reachable range reported below -- w in (-inf, +0.2] on the physical branch,
    from which this computation concluded that RADIATION at +1/3 was out of reach --
    was a property of the two-class EXPONENTIAL family used here, not of the
    construction.  The derived weight gives w = 1/3 exactly, for every configuration
    tried.  So the caveat below, that a negative answer would be a real problem for
    the construction, is moot: the answer came out positive.

Still open from this computation's list, and untouched by comp 158: conservation is
checked in flat space rather than on a curved background, and the construction is
fixed only up to a divergence-free term, the standard Irving-Kirkwood/Hardy gauge
ambiguity, so Pi_mn is a representative and not yet a unique answer.

A "pass": (1) the FRW divergence is d(Pi) and the fluid is proportional to g only at
rho + p = 0; (2) the unit-dyad sum fails to settle; (3) the trace identity holds
exactly; (4) the perfect-fluid residual decreases monotonically with cloud size;
(5)-(6) rho + p has a root in lam, and Pi_mn is proportional to g there to the
residual.
"""

from __future__ import annotations
import numpy as np
import sympy as sp

ETA = np.diag([1.0, -1.0, -1.0, -1.0])          # mostly-minus, as pst-body-se.tex uses
U   = np.array([1.0, 0.0, 0.0, 0.0])            # the ensemble's rest frame


# ---------------------------------------------------------------- (1) symbolic
def scalar_form_on_frw():
    """Bianchi on T_mn = Pi g_mn, and when a perfect fluid is proportional to g."""
    t, r, th, k = sp.symbols('t r theta k', real=True)
    ph = sp.Symbol('phi', real=True)
    a = sp.Function('a', positive=True)(t)
    x = [t, r, th, ph]
    g = sp.diag(1, -a**2/(1 - k*r**2), -a**2*r**2, -a**2*r**2*sp.sin(th)**2)
    ginv, n = g.inv(), 4
    Gam = [[[sp.simplify(sum(ginv[i, l]*(sp.diff(g[l, j], x[m]) + sp.diff(g[l, m], x[j])
                                         - sp.diff(g[j, m], x[l])) for l in range(n))/2)
             for m in range(n)] for j in range(n)] for i in range(n)]

    def div_up(T):                               # nabla_m T^{mn}
        out = []
        for nu in range(n):
            e = sum(sp.diff(T[mu][nu], x[mu]) for mu in range(n))
            e += sum(Gam[mu][mu][l]*T[l][nu] for mu in range(n) for l in range(n))
            e += sum(Gam[nu][mu][l]*T[mu][l] for mu in range(n) for l in range(n))
            out.append(sp.simplify(sp.expand_trig(sp.simplify(e))))
        return out

    Pi = sp.Function('Pi')(t)
    d_scalar = div_up([[sp.simplify(Pi*ginv[i, j]) for j in range(n)] for i in range(n)])

    rho, p = sp.Function('rho')(t), sp.Function('p')(t)
    T_fluid = [[sp.simplify((rho + p)*(1 if (i == j == 0) else 0) - p*ginv[i, j])
                for j in range(n)] for i in range(n)]
    gap = sp.simplify(sp.Matrix(T_fluid) + p*ginv)          # T^{mn} + p g^{mn}
    prop_iff = sp.simplify(gap.subs(rho, -p)) == sp.zeros(4, 4)
    return d_scalar, gap, prop_iff, Pi, t


# ---------------------------------------------------------------- the ensemble
def cloud(rng, n, tmax=1.0, rmax=1.0):
    """Realisation points rho(C_a), rotationally isotropic about U."""
    t = rng.uniform(-tmax, tmax, n)
    d = rng.uniform(0, 1, n)**(1/3) * rmax
    v = rng.normal(size=(n, 3))
    v /= np.linalg.norm(v, axis=1, keepdims=True)
    return np.column_stack([t, v*d[:, None]])


def pair_data(X):
    """Lowered separation vectors and invariants over all unordered pairs.

    Filled row by row rather than by indexing a full difference cube.  The cube is
    n-by-n-by-4, 82 MB at the n = 1600 this file reaches, and np.triu_indices adds
    two index arrays of 10 MB each on top; the browser runtime this book publishes to
    cannot be relied on for that.  Row i contributes X[i] - X[i+1:], which is exactly
    the order np.triu_indices(n, 1) visits the pairs in, so D is the same array.
    """
    n = len(X)
    D = np.empty((n * (n - 1) // 2, X.shape[1]))
    at = 0
    for i in range(n - 1):
        k = n - 1 - i
        D[at:at + k] = X[i] - X[i + 1:]
        at += k
    s = np.einsum('pm,mn,pn->p', D, ETA, D)
    return D @ ETA, s                                     # D_m , s = D.D


def decompose(P):
    """P_mn = rho u_m u_n + p h_mn (+ residual);  h_mn = u_m u_n - g_mn."""
    u_low = ETA @ U
    h = np.outer(u_low, u_low) - ETA
    basis = [np.outer(u_low, u_low), h]
    G = np.array([[np.sum(x*y) for y in basis] for x in basis])
    c = np.linalg.solve(G, np.array([np.sum(x*P) for x in basis]))
    resid = P - c[0]*basis[0] - c[1]*basis[1]
    return c[0], c[1], np.linalg.norm(resid)/max(np.linalg.norm(P), 1e-30)


def prop_to_g(P):
    lam = np.sum(P*ETA)/np.sum(ETA*ETA)
    return np.linalg.norm(P - lam*ETA)/max(np.linalg.norm(P), 1e-30)


def main():
    print("=" * 100)
    print("  Computation 155 -- the tensor projection Pi_mn: structure, mechanism, and what is left")
    print("=" * 100)
    print()
    rng = np.random.default_rng(20260927)

    # ---- (1) ----------------------------------------------------------
    print("  (1) THE SCALAR FORM T_mn = Pi g_mn IS EXACTLY AND ONLY A COSMOLOGICAL CONSTANT")
    d_scalar, gap, prop_iff, Pi, t = scalar_form_on_frw()
    nonzero = [e for e in d_scalar if e != 0]
    print(f"      on FRW, nabla_m T^{{mn}} for T_mn = Pi g_mn : {d_scalar}")
    print(f"      -> conservation forces {nonzero} = 0, i.e. Pi is CONSTANT")
    print(f"      a perfect fluid minus (-p) g^{{mn}} leaves : {sp.simplify(gap[0, 0])} in the 00 slot")
    print(f"      -> proportional to g iff rho + p = 0, i.e. w = -1 : {prop_iff}")
    check1 = (nonzero == [sp.Derivative(Pi, t)] or
              all(sp.simplify(e - sp.Derivative(Pi, t)) == 0 for e in nonzero)) and prop_iff
    print(f"      so an isotropic DUST or RADIATION fluid is already outside the scalar form: {check1}")
    print()

    # ---- (2) ----------------------------------------------------------
    print("  (2) THE UNIT DYAD e_m e_n IS SINGULAR ON THE LIGHT CONE, WHERE tau_null SITS")
    print(f"      {'cloud':>7}{'pairs':>10}{'max |e|':>14}{'rho (unit dyad)':>18}{'fluid resid':>13}")
    unit_rho = []
    for n in (200, 400, 800):
        X = cloud(rng, n)
        Dl, s = pair_data(X)
        keep = np.abs(s) > 1e-14
        e = Dl[keep] / np.sqrt(np.abs(s[keep]))[:, None]
        P = np.einsum('pm,pn->mn', e, e)
        a, b, r = decompose(P)
        unit_rho.append(a/len(s))                      # per-pair, to compare across sizes
        print(f"      {n:7d}{len(s):10d}{np.max(np.abs(e)):14.3e}{a/len(s):18.4f}{r:13.2e}")
    spread = (max(unit_rho) - min(unit_rho))/abs(np.mean(unit_rho))
    check2 = spread > 0.10
    print(f"      -> per-pair value swings by {spread*100:.0f}% across cloud sizes: does not settle: {check2}")
    print()

    # ---- (3) ----------------------------------------------------------
    print("  (3) THE SEPARATION DYAD D_m D_n IS WELL DEFINED, AND ITS TRACE IS THE INVARIANT")
    X = cloud(rng, 400)
    Dl, s = pair_data(X)
    # g^{mn} D_m D_n  must equal s exactly, pair by pair
    tr = np.einsum('pm,mn,pn->p', Dl, np.linalg.inv(ETA), Dl)
    trace_err = np.max(np.abs(tr - s))
    check3 = trace_err < 1e-10
    print(f"      max | g^mn D_m D_n - s |  over {len(s)} pairs = {trace_err:.2e}")
    print(f"      -> the scalar Pi is the TRACE of the tensor projection, not a factor on g: {check3}")
    print()

    # ---- (4) ----------------------------------------------------------
    print("  (4) FOR AN INVARIANT-ONLY WEIGHT ON AN ISOTROPIC ENSEMBLE, Pi_mn IS A PERFECT FLUID")
    W_ref = lambda z: np.sign(z)*np.exp(-np.abs(z))
    print(f"      {'cloud':>7}{'pairs':>10}{'rho':>14}{'p':>14}{'w':>10}{'resid':>11}")
    resids = []
    for n in (200, 400, 800, 1600):
        X = cloud(rng, n)
        Dl, s = pair_data(X)
        P = np.einsum('p,pm,pn->mn', W_ref(s), Dl, Dl)
        a, b, r = decompose(P)
        resids.append(r)
        print(f"      {n:7d}{len(s):10d}{a:14.1f}{b:14.1f}{b/a:10.4f}{r:11.2e}")
    # Requiring a Monte Carlo residual to fall at EVERY doubling, with no tolerance
    # and no replication, fails roughly a third of seeds.  The claim is that it falls
    # across the range; test the endpoints and report the path.
    check4 = resids[-1] < 0.5 * resids[0]
    print(f"      -> residual falls across the range, {resids[0]:.2e} -> {resids[-1]:.2e}"
          f" (factor {resids[0] / max(resids[-1], 1e-300):.1f}): {check4}")
    print()

    # ---- (5) and (6) ---------------------------------------------------
    print("  (5) rho IS THE TIMELIKE-PAIR MOMENT, p THE SPACELIKE-PAIR MOMENT")
    X = cloud(rng, 1600)
    Dl, s = pair_data(X)
    f = np.exp(-np.abs(s))
    P_t = np.einsum('p,pm,pn->mn', np.where(s > 0, f, 0.0), Dl, Dl)
    P_s = np.einsum('p,pm,pn->mn', np.where(s < 0, f, 0.0), Dl, Dl)
    at, bt, _ = decompose(P_t)
    as_, bs, _ = decompose(P_s)
    print(f"      timelike pairs only  (T_ab < tau_null) : rho = {at:12.1f}   p = {bt:12.1f}")
    print(f"      spacelike pairs only (T_ab > tau_null) : rho = {as_:12.1f}   p = {bs:12.1f}")
    print(f"      the split is eq:metric's own classification, not an extra ingredient")
    print()

    print("  (6) w = p/rho IS THE BUDGET RATIO, AND w = -1 IS THE SCALAR FORM")
    def rho_p(lam):
        P = P_t - lam*P_s
        a, b, r = decompose(P)
        return a, b, r, P
    print(f"      {'lam':>8}{'rho':>14}{'p':>14}{'w':>11}")
    for lam in (0.0, 0.25, 0.50, 0.75, 1.00):
        a, b, _, _ = rho_p(lam)
        print(f"      {lam:8.2f}{a:14.1f}{b:14.1f}{b/a:11.4f}")
    # bisect on rho + p, which is smooth (w has a pole where rho crosses zero)
    lo, hi = 0.0, 1.0
    flo = sum(rho_p(lo)[:2])
    for _ in range(80):
        mid = 0.5*(lo + hi)
        fm = sum(rho_p(mid)[:2])
        if (fm > 0) == (flo > 0):
            lo, flo = mid, fm
        else:
            hi = mid
    lam_star = 0.5*(lo + hi)
    a, b, r, P = rho_p(lam_star)
    pg = prop_to_g(P)
    print(f"      rho + p = 0 at lam* = {lam_star:.6f} : rho = {a:.1f}, p = {b:.1f}, w = {b/a:+.8f}")
    print(f"      || Pi - lam g || / || Pi || there = {pg:.2e}, against a decomposition residual of {r:.2e}")
    check6 = abs(b/a + 1.0) < 1e-6 and pg <= 3*r
    print(f"      -> at that one point Pi_mn IS the scalar form, to the accuracy of the ensemble: {check6}")
    print()

    # ---- the honest half ------------------------------------------------
    print("  SCOPE: the reachable equation of state on the physical branch (rho > 0)")
    ws = []
    for lam in np.linspace(0.0, lam_star, 9):
        a, b, _, _ = rho_p(lam)
        if a > 0:
            ws.append(b/a)
    print(f"      w runs from {max(ws):+.3f} down to {min(ws):+.3f} over lam in [0, lam*]")
    print(f"      dust (0) and Lambda (-1) are inside; radiation (+1/3) is not reachable with")
    print(f"      THIS profile -- which comp 158 shows is a property of the exponential family")
    print(f"      used here, not of the construction.  The derived weight W(s) = -|s| gives")
    print(f"      w = 1/3 exactly, so the profile is not free after all.  See comp 158.")
    print()

    print("=" * 100)
    print("  ASSESSMENT")
    print("=" * 100)
    print(f"  (1) the scalar form is exactly and only a cosmological constant      : {check1}")
    print(f"  (2) the unit-dyad completion is singular on the light cone           : {check2}")
    print(f"  (3) g^mn D_m D_n = s, so the scalar Pi is the tensor projection's trace : {check3}")
    print(f"  (4) invariant-only weight + isotropic ensemble -> perfect fluid      : {check4}")
    print(f"  (6) rho + p = 0 has a root in the budget ratio, and Pi ~ g there     : {check6}")
    print()
    ok = check1 and check2 and check3 and check4 and check6
    print("  RESULT: " + ("ADVANCES sec:stress-energy-tensor-open -- supplies the tensor structure, "
                          "the light-cone obstruction, and the equation-of-state mechanism; the "
                          "weight and the radiation branch are settled by comp 158"
                          if ok else "MISMATCH -- revise"))


if __name__ == "__main__":
    main()
