#!/usr/bin/env python3
"""
reversibility.py -- numerical test of a candidate time-reversal / detailed-balance
property of the monitored tie model (D40 conventions, hbar = 1, J = 1).

CLAIM UNDER TEST
    Ldag_pi = Theta o L o Theta,      Theta(A) = conj(A) entrywise in the Z basis,
where L is the Heisenberg generator and Ldag_pi its adjoint in the GNS inner product
<A,B>_pi = Tr(pi A^dag B) (also checked: KMS inner product), pi = exp(-beta H0)/Z.
If true, the pointer kernel T_t(a->b) = <b| e^{tL}(|a><a|) |b> is exactly
pi-reversible at every t:  pi_a T_t(a->b) = pi_b T_t(b->a).

MODEL (D40 sec. 0)
    H0 = -J sum_<ij> Z_i Z_j - J sum_i e_i Z_i        (e_i = frozen exterior field)
    K  = sum_i g_i X_i 1[h_i = 0],  h_i = sum_{live j~i} Z_j + e_i
    thermal rates c_i(s) = nu / (1 + exp(2 beta J s_i h_i))
        R1: rank-one jumps sqrt(c_i(s)) |s^i><s|  for every configuration s
        R2: D40 neighbour-resolved jumps sqrt(c_i(s,eta)) |-s><s|_i (x) Pi_eta(nbrs) (x) 1
        R3: (EXTRA control, not D40) site/spin-resolved but neighbour-UNresolved jumps
            |-s><s|_i (x) sqrt(c_i(s, h_i-operator))   (one jump per (i,s))
    monitoring L_i = sqrt(kappa/2) Z_i
    Liouvillian L(rho) = -i[H0+K, rho] + sum_J D[J] rho, row-major vec(rho)[s*d+t] = rho_st.

TESTS
    T0   ||L(pi)||_max
    Tid  operator identity residual  max|Ldag_pi - Theta L Theta| / max|L|  (GNS and KMS)
    T1   max_{a!=b} |pi_a T_t(a->b) - pi_b T_t(b->a)| / max_{a!=b} pi_a T_t(a->b),  t = 0.1, 1, 5
         + Kolmogorov affinities log(prod fwd / prod bwd) on a fundamental cycle basis of
           the single-flip graph (measure-free)
    T2   zero-frequency Schur complement G = L_PP - L_PQ L_QQ^{-1} L_QP on populations
    T3   slow-manifold generator G_s = R_P Lam R_P^{-1} (n_config slowest modes, only if
         |Re lam_{n+1}| >= 5 |Re lam_n|), and the slip-symmetrised G^ = X^{-1/2} G_s X^{1/2},
         X = R_P W_P (population block of the slow spectral projector = "initial slip").
    T4   controls: (a) S1 non-gaugeable loop phase  g' e^{i Phi} |s3><s1| + h.c.,
                   (b) K = g sum X_i without tie projector (reference = true stationary state),
                   (c) EXTRA: thermal convention R3.
    T5   site phases g_i -> g_i w_i e^{i phi_i} (inhomogeneous w_i): plain Theta vs gauge-dressed
         Theta_U = U Theta U^dag (U = prod exp(-i phi_i Z_i / 2)), and kernel identity vs phi = 0.

Run:  python3 reversibility.py      (numpy + scipy only; parallel over up to 12 processes)
"""
import os
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS",
           "VECLIB_MAXIMUM_THREADS", "NUMEXPR_NUM_THREADS"):
    os.environ.setdefault(_v, "1")
import sys
import time
import itertools
import warnings
import numpy as np
import scipy.linalg as sla
import scipy.sparse as sp
import scipy.sparse.linalg as spla
from scipy.sparse.csgraph import connected_components
from multiprocessing import Pool

# numpy 2.0 + macOS Accelerate emits spurious FP warnings inside matmul; every reported
# number is checked for finiteness instead (see 'non-finite outputs' at the end).
warnings.filterwarnings("ignore", message=".*encountered in matmul")

JJ = 1.0
TIMES = (0.1, 1.0, 5.0)
BETAS, GS, NUS, KAPPAS = (0.3, 0.7, 1.2), (0.3, 1.0, 3.0), (0.2, 1.0), (0.5, 2.0, 10.0, 50.0)
SEP_RATIO = 5.0          # slow/fast separation required for T3
AFF_RELTOL = 1e-4        # edges used for affinities: both directions >= reltol * largest edge
DENSE_MAX = 1000         # dense linear algebra on the population sector up to this size
EXPM_DENSE_MAX = 300     # dense expm up to this size, else expm_multiply
GAUGE_W = np.array([1.0, 0.8, 1.25, 0.9, 1.1, 0.7])          # inhomogeneous |g_i| / g
GAUGE_PHI = np.array([0.37, 2.21, -1.13, 2.90, 0.61, -2.47])  # fixed site phases phi_i

BONDS4 = [(0, 1), (0, 2), (1, 3), (2, 3)]            # 2x2 plaquette: 0=(0,0) 1=(0,1) 2=(1,0) 3=(1,1)
GEOM = {
    "S1":    (2, [(0, 1)], (1, 1)),                   # each spin: 3 frozen nbrs, 2 up 1 down
    "S2/B1": (4, BONDS4, (2, 2, 2, 2)),               # all-plus frozen boundary
    "S2/B2": (4, BONDS4, (0, 0, 0, 0)),               # each site one + and one - frozen nbr
    "S2/B3": (4, BONDS4, (2, 0, 0, -2)),              # diagonal domain-wall boundary
    "S2/B4": (4, BONDS4, (2, -2, 0, 0)),              # mixed boundary
    "S3":    (6, [(0, 1), (2, 3), (0, 2), (1, 3), (1, 4), (3, 5), (4, 5)], (2, 1, 2, 1, 2, 2)),
}
# S1 labels (bit k <-> site k, bit 1 = spin down): s0=(+,+)=0, s1=(-,+)=1, s2=(-,-)=3, s3=(+,-)=2
S1_LOOP = (0, 1, 3, 2)


class Model:
    def __init__(self, key):
        n, bonds, e = GEOM[key]
        d = 2 ** n
        self.key, self.n, self.d, self.bonds = key, n, d, bonds
        self.e = np.asarray(e, float)
        a = np.arange(d)
        self.sig = np.array([1 - 2 * ((a >> k) & 1) for k in range(n)]).T          # (d, n)
        self.nbrs = [[] for _ in range(n)]
        for i, j in bonds:
            self.nbrs[i].append(j)
            self.nbrs[j].append(i)
        self.h = np.array([self.sig[:, self.nbrs[i]].sum(1) + self.e[i] for i in range(n)]).T
        self.E = (-JJ * sum(self.sig[:, i] * self.sig[:, j] for i, j in bonds)
                  - JJ * (self.sig * self.e).sum(1)).astype(float)
        self.P = a * d + a
        mask = np.ones(d * d, bool)
        mask[self.P] = False
        self.Q = np.nonzero(mask)[0]
        self.edges = [(x, x ^ (1 << k)) for x in range(d) for k in range(n) if x < (x ^ (1 << k))]
        self.ntie = int((self.h == 0).sum())


def hamiltonian(m, g, w, phi, notie=False, loop=None):
    H = np.diag(m.E).astype(complex)
    for i in range(m.n):
        for a in range(m.d):
            if notie or m.h[a, i] == 0:
                H[a ^ (1 << i), a] += g * w[i] * np.exp(1j * phi[i] * m.sig[a, i])
    if loop is not None:
        gp, Phi, s1, s3 = loop
        H[s3, s1] += gp * np.exp(1j * Phi)
        H[s1, s3] += gp * np.exp(-1j * Phi)
    return H


def thermal_jumps(m, beta, nu, conv):
    c = nu / (1.0 + np.exp(2.0 * beta * JJ * m.sig * m.h))                      # (d, n)
    groups = {}
    for i in range(m.n):
        for a in range(m.d):
            if conv == "R1":
                key = (i, a)
            elif conv == "R2":
                key = (i, int(m.sig[a, i]), tuple(int(x) for x in m.sig[a, m.nbrs[i]]))
            elif conv == "R3":
                key = (i, int(m.sig[a, i]))
            else:
                raise ValueError(conv)
            groups.setdefault(key, []).append(a)
    jumps = []
    for key, al in groups.items():
        a_ = np.array(al)
        jumps.append((a_ ^ (1 << key[0]), a_, np.sqrt(c[a_, key[0]])))
    return jumps


def lindbladian(m, H, jumps, kappa):
    d = m.d
    D2 = d * d
    Hs = sp.csr_matrix(H)
    I = sp.identity(d, dtype=complex, format="csr")
    L = -1j * (sp.kron(Hs, I) - sp.kron(I, Hs.T))
    rr, cc, vv = [], [], []
    A = np.zeros(d)
    for b_, a_, w_ in jumps:                       # J = sum_k w_k |b_k><a_k| ; J (x) conj(J)
        rr.append((b_[:, None] * d + b_[None, :]).ravel())
        cc.append((a_[:, None] * d + a_[None, :]).ravel())
        vv.append((w_[:, None] * np.conj(w_)[None, :]).ravel())
        np.add.at(A, a_, np.abs(w_) ** 2)          # J^dag J is diagonal
    S = sp.coo_matrix((np.concatenate(vv), (np.concatenate(rr), np.concatenate(cc))),
                      shape=(D2, D2)).tocsr()
    zz = (m.sig @ m.sig.T).astype(float)
    diag = -0.5 * (A[:, None] + A[None, :]).ravel() + 0.5 * kappa * (zz.ravel() - m.n)
    return (L + S + sp.diags(diag)).tocsr()


# ------------------------------------------------------------------ diagnostics
def identity_residuals(L, pi, d, W=None):
    """max|Ldag - Theta_W L Theta_W| / max|L| for GNS and KMS adjoints."""
    perm = np.arange(d * d).reshape(d, d).T.ravel()
    Lh = L.T.tocsr()[perm][:, perm]                          # Heisenberg matrix
    tgt = Lh.conj()
    if W is not None:
        tgt = sp.diags(W) @ tgt @ sp.diags(W.conj())
    piR = np.tile(pi, d)                                      # (s,t) -> pi_t
    sK = np.kron(np.sqrt(pi), np.sqrt(pi))                    # (s,t) -> sqrt(pi_s pi_t)
    Mg = sp.diags(1.0 / piR) @ L @ sp.diags(piR)
    Mk = sp.diags(1.0 / sK) @ L @ sp.diags(sK)
    nL = abs(L).max()
    rg = abs(Mg - tgt).max() / nL
    rk = abs(Mk - tgt).max() / nL
    return float(rg), float(rk)


def db_residual(M, mu):
    F = M * mu[None, :]                                       # F[b,a] = mu_a M(a->b)
    off = ~np.eye(len(mu), dtype=bool)
    den = np.max(np.abs(F[off]))
    return float(np.max(np.abs(F - F.T)[off]) / den) if den > 0 else 0.0


def affinities(M, edges, reltol=AFF_RELTOL):
    """Fundamental-cycle affinities log(prod M(fwd)/prod M(bwd)); M[b,a] = weight a->b.
    Returns (array of affinities, number of sign-mismatched edges)."""
    n = M.shape[0]
    sc = max(max(abs(M[b, a]), abs(M[a, b])) for a, b in edges)
    keep, mism = [], 0
    for a, b in edges:
        x, y = M[b, a], M[a, b]
        if abs(x) < reltol * sc or abs(y) < reltol * sc:
            continue
        if x * y <= 0:
            mism += 1
            continue
        keep.append((a, b, np.log(abs(x) / abs(y))))
    adj = [[] for _ in range(n)]
    for k, (a, b, w) in enumerate(keep):
        adj[a].append((b, w, k))
        adj[b].append((a, -w, k))
    phi = [None] * n
    tree = set()
    for r in range(n):
        if phi[r] is not None:
            continue
        phi[r] = 0.0
        queue = [r]
        while queue:
            u = queue.pop(0)
            for v, w, k in adj[u]:
                if phi[v] is None:
                    phi[v] = phi[u] + w
                    tree.add(k)
                    queue.append(v)
    aff = np.array([w - (phi[b] - phi[a]) for k, (a, b, w) in enumerate(keep) if k not in tree])
    return aff, mism


def loop_affinity(M):
    a, b, c, e = S1_LOOP                                      # s0 -> s1 -> s2 -> s3 -> s0
    fwd = M[b, a] * M[c, b] * M[e, c] * M[a, e]
    bwd = M[a, b] * M[b, c] * M[c, e] * M[e, a]
    # sign-consistent products (G may have negative off-diagonals); nan = sign mismatch
    return float(np.log(fwd / bwd)) if fwd * bwd > 0 else float("nan")


def maxabs(x):
    return float(np.max(np.abs(x))) if len(x) else 0.0




def perron(T):
    w, V = np.linalg.eig(T)
    v = np.real(V[:, np.argmin(np.abs(w - 1.0))])
    return v / v.sum()


# ------------------------------------------------------------------ population sector
class Sector:
    """L is block diagonal over the weakly connected components of its sparsity graph.
    Keep the components that contain populations (exact reduction; e.g. E_s - E_t is
    conserved in the model, so only equal-energy coherences couple to populations)."""

    def __init__(self, L, m):
        self.ncomp, lab = connected_components(abs(L) > 0, directed=False)
        popl = np.unique(lab[m.P])
        self.keep = np.nonzero(np.isin(lab, popl))[0]
        pos = -np.ones(L.shape[0], int)
        pos[self.keep] = np.arange(len(self.keep))
        self.P = pos[m.P]
        mask = np.ones(len(self.keep), bool)
        mask[self.P] = False
        self.Q = np.nonzero(mask)[0]
        self.L = L[self.keep][:, self.keep].tocsr()
        self.size = len(self.keep)
        self.dense = self.size <= DENSE_MAX
        self._lab, self._popl, self._full = lab, popl, L

    def excluded_max_re(self):
        """largest real part of any eigenvalue in the discarded (pure-coherence) sectors"""
        lab, L = self._lab, self._full
        best = -np.inf
        order = np.argsort(lab, kind="stable")
        bounds = np.searchsorted(lab[order], np.arange(self.ncomp + 1))
        for c in range(self.ncomp):
            if c in set(self._popl.tolist()):
                continue
            idx = order[bounds[c]:bounds[c + 1]]
            blk = L[idx][:, idx].toarray()
            best = max(best, float(np.max(np.linalg.eigvals(blk).real)))
        return best


def kernels(S, d):
    Ts = []
    if S.size <= EXPM_DENSE_MAX:
        Ld = S.L.toarray()
        for t in TIMES:
            Ts.append(sla.expm(t * Ld)[np.ix_(S.P, S.P)])
    else:
        B = np.zeros((S.size, d), complex)
        B[S.P, np.arange(d)] = 1.0
        tp = 0.0
        for t in TIMES:
            B = spla.expm_multiply((t - tp) * S.L, B)
            tp = t
            Ts.append(B[S.P, :].copy())
    return Ts


def schur(S):
    P, Q = S.P, S.Q
    if S.dense:
        Ld = S.L.toarray()
        X = np.linalg.solve(Ld[np.ix_(Q, Q)], Ld[np.ix_(Q, P)])
        return Ld[np.ix_(P, P)] - Ld[np.ix_(P, Q)] @ X
    Lr, LQ = S.L[P, :], S.L[Q, :]
    X = spla.splu(LQ[:, Q].tocsc()).solve(LQ[:, P].toarray())
    return Lr[:, P].toarray() - Lr[:, Q] @ X


def stationary(S):
    p0 = S.P[0]
    if S.dense:
        A = S.L.toarray()
        A[p0, :] = 0.0
        A[p0, S.P] = 1.0
        rhs = np.zeros(S.size, complex)
        rhs[p0] = 1.0
        return np.linalg.solve(A, rhs)
    A = S.L.tolil()
    A[p0, :] = 0
    A[p0, S.P] = 1.0
    rhs = np.zeros(S.size, complex)
    rhs[p0] = 1.0
    return spla.spsolve(A.tocsc(), rhs)


def slow_generator(S, n):
    """n slowest modes of the population sector; separation also checked against the
    slowest eigenvalue of every discarded sector."""
    if S.dense:
        lam, R = sla.eig(S.L.toarray())
        o = np.argsort(-lam.real)
        lam, R = lam[o], R[:, o]
        Wm = np.linalg.inv(R)[:n]
    else:
        k, s0 = n + 10, 1e-2
        lam, R = spla.eigs(S.L.tocsc(), k=k, sigma=s0, which="LM")
        mu, Y = spla.eigs(S.L.conj().T.tocsc(), k=k, sigma=s0, which="LM")
        o = np.argsort(-lam.real)
        lam, R = lam[o], R[:, o]
        Y = Y[:, np.argsort(-mu.real)][:, :n]
        Wm = np.linalg.solve(Y.conj().T @ R[:, :n], Y.conj().T)
    fast = lam[n].real if len(lam) > n else -np.inf
    fast = max(fast, S.excluded_max_re())
    sep = fast / lam[n - 1].real
    if not sep >= SEP_RATIO:
        return None, sep
    R, lam = R[:, :n], lam[:n]
    RP, WP = R[S.P, :], Wm[:, S.P]
    Gs = RP @ np.diag(lam) @ np.linalg.inv(RP)
    X = RP @ WP
    Xh = sla.sqrtm(X)
    Gh = np.linalg.solve(Xh, Gs @ Xh)
    return (Gs, Gh, X), sep


# ------------------------------------------------------------------ one parameter point
def run_point(spec):
    key, conv, variant, beta, g, nu, kappa, do_t3 = spec
    m = Model(key)
    n, d = m.n, m.d
    w, phi, notie, loop = np.ones(n), np.zeros(n), False, None
    if variant == "gauge":
        w, phi = GAUGE_W[:n], GAUGE_PHI[:n]
    elif variant.startswith("loop"):
        loop = (g, float(variant.split(":")[1]), 1, 2)            # s1 <-> s3, g' = g
    elif variant == "notie":
        notie = True
    H = hamiltonian(m, g, w, phi, notie, loop)
    jumps = thermal_jumps(m, beta, nu, conv)
    L = lindbladian(m, H, jumps, kappa)
    pw = np.exp(-beta * (m.E - m.E.min()))
    pi = pw / pw.sum()
    vpi = np.zeros(d * d, complex)
    vpi[m.P] = pi
    o = dict(spec=spec)
    o["T0"] = float(np.max(np.abs(L @ vpi)))
    o["Tid_gns"], o["Tid_kms"] = identity_residuals(L, pi, d)
    if variant == "gauge":
        u = np.exp(-0.5j * (m.sig * phi[None, :]).sum(1))
        Wg = np.kron(u ** 2, np.conj(u) ** 2)
        o["TidU_gns"], o["TidU_kms"] = identity_residuals(L, pi, d, Wg)
    S = Sector(L, m)
    o["sector"] = S.size
    if notie:
        x = stationary(S)
        o["ss_res"] = float(np.max(np.abs(S.L @ x)))
        mu = np.real(x[S.P])
        o["ss_coh"] = float(np.max(np.abs(x[S.Q])))
    else:
        mu = pi
    allpairs = [(a, b) for a in range(d) for b in range(a + 1, d)]
    Ts = kernels(S, d)
    o["kern_err"] = max(max(float(np.max(np.abs(T.imag))), float(np.max(np.abs(T.real.sum(0) - 1))))
                        for T in Ts)
    Tr = [T.real for T in Ts]
    o["T1_db_t"] = [db_residual(T, mu) for T in Tr]
    o["T1_db"] = max(o["T1_db_t"])
    o["T1_dbP"] = max(db_residual(T, perron(T)) for T in Tr)      # w.r.t. the kernel's own invariant law
    affs = [affinities(T, allpairs) for T in Tr]
    o["T1_aff"] = max(maxabs(a) for a, _ in affs)
    o["T1_ncyc"] = min(len(a) for a, _ in affs)
    o["T1_mism"] = max(mm for _, mm in affs)
    if key == "S1":
        o["T1_loop"] = [loop_affinity(T) for T in Tr]
    if variant == "gauge":
        L0 = lindbladian(m, hamiltonian(m, g, w, np.zeros(n)), jumps, kappa)
        T0s = kernels(Sector(L0, m), d)
        o["gauge_kdiff"] = max(float(np.max(np.abs(a - b))) for a, b in zip(Ts, T0s))
    G = schur(S)
    nG = np.max(np.abs(G))
    o["T2_imag"] = float(np.max(np.abs(G.imag)) / nG)
    G = G.real
    off = ~np.eye(d, dtype=bool)
    o["T2_colsum"] = float(np.max(np.abs(G.sum(0))) / nG)
    o["T2_stat"] = float(np.max(np.abs(G @ mu)) / (nG * mu.max()))
    o["T2_minoff"] = float(G[off].min() / np.max(np.abs(G[off])))
    o["T2_db"] = db_residual(G, mu)
    aG, mmG = affinities(G, allpairs)
    o["T2_aff"], o["T2_mism"], o["T2_ncyc"] = maxabs(aG), mmG, len(aG)
    if key == "S1":
        o["T2_loop"] = loop_affinity(G)
    if do_t3:
        res, sep = slow_generator(S, d)
        o["T3_sep"] = float(sep)
        if res is not None:
            Gs, Gh, X = res
            o["T3_imag"] = float(max(np.max(np.abs(Gs.imag)), np.max(np.abs(Gh.imag))) / np.max(np.abs(Gs)))
            Gs, Gh = Gs.real, Gh.real
            o["T3_raw_db"] = db_residual(Gs, mu)
            o["T3_sym_db"] = db_residual(Gh, mu)
            o["T3_slip"] = float(np.max(np.abs(X - np.eye(d))))
            ar, mr = affinities(Gs, allpairs)
            ah, mh = affinities(Gh, allpairs)
            o["T3_raw_aff"], o["T3_raw_mism"] = maxabs(ar), mr
            o["T3_sym_aff"], o["T3_sym_mism"] = maxabs(ah), mh
            o["T3_minoff"] = float(Gs[off].min() / np.max(np.abs(Gs[off])))
            o["T3_G_vs_Gs"] = float(np.max(np.abs(G - Gs)) / np.max(np.abs(Gs)))
            Fs = Gs * mu[None, :]
            Dv = np.abs(Fs - Fs.T)
            np.fill_diagonal(Dv, 0.0)
            bw, aw = np.unravel_index(np.argmax(Dv), Dv.shape)
            gmax = np.max(np.abs(Gs[off]))
            o["T3_worst"] = (bin(int(aw) ^ int(bw)).count("1"), float(Gs[bw, aw] / gmax), float(Gs[aw, bw] / gmax))
            if key == "S1":
                o["T3_raw_loop"] = loop_affinity(Gs)
    return o


# ------------------------------------------------------------------ driver
def grid(sub=False):
    if sub:
        return list(itertools.product((0.3, 1.2), (1.0, 3.0), (1.0,), (0.5, 10.0)))
    return list(itertools.product(BETAS, GS, NUS, KAPPAS))


PHI_HALF, PHI_THIRD, PHI_PI = "loop:1.5707963267948966", "loop:1.0471975511965976", "loop:3.141592653589793"
S2KEYS = ["S2/B1", "S2/B2", "S2/B3", "S2/B4"]


def build_specs():
    specs = []
    for b, g, nu, k in grid():
        for var in ("model", "gauge", PHI_HALF, PHI_THIRD, PHI_PI, "notie"):
            specs.append(("S1", "R1", var, b, g, nu, k, True))          # R1 == R2 for S1
        specs.append(("S1", "R3", "model", b, g, nu, k, True))
        for s in S2KEYS:
            specs.append((s, "R1", "model", b, g, nu, k, True))
            specs.append((s, "R2", "model", b, g, nu, k, True))
            specs.append((s, "R1", "gauge", b, g, nu, k, True))
            specs.append((s, "R2", "gauge", b, g, nu, k, True))
            specs.append((s, "R2", "notie", b, g, nu, k, True))
            specs.append((s, "R3", "model", b, g, nu, k, True))
        specs.append(("S3", "R1", "model", b, g, nu, k, k >= 10))
        specs.append(("S3", "R2", "model", b, g, nu, k, k >= 10))
    for b, g, nu, k in grid(sub=True):
        specs.append(("S3", "R2", "gauge", b, g, nu, k, False))
        specs.append(("S3", "R2", "notie", b, g, nu, k, False))
        specs.append(("S3", "R3", "model", b, g, nu, k, False))
    return specs


SCALE_KAPPAS = (10.0, 20.0, 50.0, 100.0, 200.0, 500.0, 1000.0)


def build_scaling_specs():
    specs = []
    for k in SCALE_KAPPAS:
        for g in (0.3, 1.0):
            specs.append(("S1", "R1", "model", 0.7, g, 1.0, k, True))
            specs.append(("S2/B2", "R2", "model", 0.7, g, 1.0, k, True))
        specs.append(("S3", "R2", "model", 0.7, 1.0, 1.0, k, True))
        specs.append(("S1", "R1", PHI_HALF, 0.7, 1.0, 1.0, k, True))
        specs.append(("S2/B2", "R3", "model", 0.7, 1.0, 1.0, k, True))
    return specs


def safe_run(spec):
    try:
        return run_point(spec)
    except Exception as exc:                                  # report, never hide
        return dict(spec=spec, error=repr(exc))


def e(x):
    return "   n/a  " if x is None or (isinstance(x, float) and np.isnan(x)) else f"{x:8.1e}"


def nonfinite(o):
    bad = 0
    for k, v in o.items():
        if k == "spec":
            continue
        for x in (v if isinstance(v, list) else [v]):
            if isinstance(x, float) and not np.isfinite(x) and not (k == "T3_sep"):
                bad += 1
    return bad


NAMES = {PHI_HALF: "loop Phi=pi/2", PHI_THIRD: "loop Phi=pi/3", PHI_PI: "loop Phi=pi",
         "notie": "no tie proj.", "gauge": "site phases", "model": "model"}


def label(k):
    if k[1] == "R3" and k[2] == "model":
        return "R3 unresolved"
    return NAMES.get(k[2], k[2])


def main():
    t0 = time.time()
    main_specs, scale_specs = build_specs(), build_scaling_specs()
    specs = list(dict.fromkeys(main_specs + scale_specs))
    nproc = min(12, os.cpu_count() or 1)
    specs.sort(key=lambda s: 0 if s[0] == "S3" else 1)       # heavy tasks first
    with Pool(nproc) as pool:
        outs = pool.map(safe_run, specs, chunksize=1)
    errs = [o for o in outs if "error" in o]
    good = {o["spec"]: o for o in outs if "error" not in o}
    main_outs = [good[s] for s in main_specs if s in good]
    scale_outs = [good[s] for s in scale_specs if s in good]
    nbad = sum(nonfinite(o) for o in good.values())
    badspecs = [o["spec"] for o in good.values() if nonfinite(o)]

    def groups(pred):
        gr = {}
        for o in main_outs:
            k = o["spec"][:3]
            if pred(k):
                gr.setdefault(k, []).append(o)
        return gr

    def vals(lst, f):
        return [o[f] for o in lst if f in o and o[f] is not None and np.isfinite(o[f])]

    def mx(lst, f):
        v = vals(lst, f)
        return max(v) if v else None

    def mn(lst, f):
        v = vals(lst, f)
        return min(v) if v else None

    def cyc(lst, f):
        return f"{min(o[f] for o in lst)}-{max(o[f] for o in lst)}"

    W = 124
    print("=" * W)
    print("reversibility.py -- candidate identity  Ldag_pi = Theta L Theta  (monitored tie model, D40 conventions, J=1)")
    print(f"grid: betaJ {BETAS} x g {GS} x nu {NUS} x kappa {KAPPAS} = 72 points  (S3 controls: 8-point subgrid"
          " betaJ{0.3,1.2} g{1,3} nu=1 kappa{0.5,10})")
    print(f"t = {TIMES};  T3 needs |Re lam_(n+1)| >= {SEP_RATIO:.0f}|Re lam_n| (all sectors);  affinity edges: both"
          f" weights >= {AFF_RELTOL:.0e} x largest")
    print("DB residual = max_(a!=b)|mu_a M(a->b) - mu_b M(b->a)| / max_(a!=b) mu_a M(a->b);  mu = pi unless stated.")
    print("=" * W)
    for key in GEOM:
        m = Model(key)
        lv = sorted(set(np.round(m.E, 9)))
        print(f"  {key:6s} live spins={m.n} d={m.d:2d} superop={m.d**2:5d}  e={tuple(int(x) for x in m.e)}  "
              f"#(site,config) ties={m.ntie:3d}  Ising levels={len(lv)}")
    print()

    # ---------------- Table A
    print("TABLE A  THE MODEL (real g, D40 thermal conventions).  MAX over the 72-point grid")
    print(f"{'system':7s}{'conv':5s}{'sector':>7s}{'T0':>9s}{'Tid.GNS':>9s}{'Tid.KMS':>9s}{'kernErr':>9s}{'T1.DB':>9s}"
          f"{'T1.|A|':>9s}{'#cyc':>8s}{'T2.csum':>9s}{'T2.DB':>9s}{'T2.|A|':>9s}{'#cyc':>8s}{'T2.minoff':>10s}"
          f"{'T3.rawDB':>9s}{'T3.symDB':>9s}{'#sep':>8s}")
    gA = groups(lambda k: k[2] == "model" and k[1] in ("R1", "R2"))
    for k in sorted(gA):
        lst = gA[k]
        nsep = sum(1 for o in lst if "T3_raw_db" in o)
        nt3 = sum(1 for o in lst if "T3_sep" in o)
        print(f"{k[0]:7s}{k[1]:5s}{mx(lst,'sector'):7d}{e(mx(lst,'T0'))}{e(mx(lst,'Tid_gns'))}{e(mx(lst,'Tid_kms'))}"
              f"{e(mx(lst,'kern_err'))}{e(mx(lst,'T1_db'))}{e(mx(lst,'T1_aff'))}{cyc(lst,'T1_ncyc'):>8s}"
              f"{e(mx(lst,'T2_colsum'))}{e(mx(lst,'T2_db'))}{e(mx(lst,'T2_aff'))}{cyc(lst,'T2_ncyc'):>8s}"
              f"{mn(lst,'T2_minoff'):10.2e}{e(mx(lst,'T3_raw_db'))}{e(mx(lst,'T3_sym_db'))}{nsep:4d}/{nt3:<3d}")
    print("  sector = size of the population-coupled invariant block of the superoperator (exact reduction).")
    print("  kernErr = max(|Im T|, |column sum - 1|).  T1 = max over t.  #cyc = min-max number of fundamental cycles tested.")
    print("  T2.minoff = min over grid of min_(a!=b) G_ab / max|G_ab|  (< 0  =>  G has negative off-diagonal entries).")
    print("  T3: rawDB for G_s = R_P Lam R_P^-1 ; symDB for X^-1/2 G_s X^1/2 ; #sep = separated points / points tried.")
    mism = max(max(o.get("T1_mism", 0), o.get("T2_mism", 0)) for lst in gA.values() for o in lst)
    print(f"  sign-mismatched edges (G_ab G_ba < 0 or T) anywhere in Table A: {mism}")
    print()

    # ---------------- Table B
    print("TABLE B  CONTROLS.  [min .. max] over the grid.  T1.DB(own) = DB residual w.r.t. the kernel's own invariant law")
    print(f"{'system':7s}{'conv':5s}{'variant':14s}{'Tid.GNS':>20s}{'Tid.KMS':>20s}{'T1.DB':>20s}{'T1.DB(own)':>20s}"
          f"{'T1.|A|':>20s}{'T2.DB':>20s}{'T2.|A|':>20s}{'T2.mism':>9s}")
    gB = groups(lambda k: k[2] == "notie" or k[2].startswith("loop") or k[1] == "R3")

    def rng(lst, f):
        return f"[{e(mn(lst, f))}..{e(mx(lst, f))}]"
    for k in sorted(gB, key=lambda k: (k[1] == "R3", k[2] == "notie", k[0], k[2])):
        lst = gB[k]
        print(f"{k[0]:7s}{k[1]:5s}{label(k):14s}{rng(lst,'Tid_gns'):>20s}{rng(lst,'Tid_kms'):>20s}{rng(lst,'T1_db'):>20s}"
              f"{rng(lst,'T1_dbP'):>20s}{rng(lst,'T1_aff'):>20s}{rng(lst,'T2_db'):>20s}{rng(lst,'T2_aff'):>20s}"
              f"{mn(lst,'T2_mism'):4d}-{mx(lst,'T2_mism'):<4d}")
    print("  loop Phi=pi: real amplitude (gaugeable) sanity case.  no tie proj.: mu = diag(rho_ss), rho_ss = true stationary state.")
    print("  R3 (extra, not D40): one neighbour-unresolved thermal jump per (site, spin value).")
    print("  T2.mism = min-max number of pairs with G_ab G_ba < 0 (a sign-level DB violation; such pairs are excluded from |A|,")
    print("  so |A| = 0 with mism > 0 is NOT reversibility).  T1.DB(own) accuracy is limited to ~1e-9 by the Perron eigenvector.")
    nt = [o for o in main_outs if o["spec"][2] == "notie"]
    if nt:
        print(f"  no tie proj.: ||L pi|| in [{e(mn(nt,'T0'))}..{e(mx(nt,'T0'))}], ||L rho_ss|| <= {e(mx(nt,'ss_res'))},"
              f" max|offdiag rho_ss| in [{e(mn(nt,'ss_coh'))}..{e(mx(nt,'ss_coh'))}]")
    for k in sorted(gB):
        lst = gB[k]
        if any("T3_sym_db" in o for o in lst):
            print(f"  T3 on {k[0]} {k[1]} {label(k)}: rawDB [{e(mn(lst,'T3_raw_db'))}..{e(mx(lst,'T3_raw_db'))}]  "
                  f"symDB [{e(mn(lst,'T3_sym_db'))}..{e(mx(lst,'T3_sym_db'))}]  "
                  f"({sum(1 for o in lst if 'T3_sym_db' in o)} separated pts)")
    print()

    # ---------------- Table C
    print("TABLE C  T5 SITE PHASES  g_i -> g w_i exp(i phi_i),  w = %s,  phi = %s   (max over grid)"
          % (tuple(float(x) for x in GAUGE_W), tuple(float(x) for x in GAUGE_PHI)))
    print(f"{'system':7s}{'conv':5s}{'Tid(plain Theta)':>20s}{'Tid.GNS(Th_U)':>14s}{'Tid.KMS(Th_U)':>14s}{'T1.DB':>9s}"
          f"{'T1.|A|':>9s}{'T2.DB':>9s}{'T2.|A|':>9s}{'T3.symDB':>9s}{'|T^phi-T^0|':>12s}")
    gC = groups(lambda k: k[2] == "gauge")
    for k in sorted(gC):
        lst = gC[k]
        print(f"{k[0]:7s}{k[1]:5s}{'[' + e(mn(lst,'Tid_gns')) + '..' + e(mx(lst,'Tid_gns')) + ']':>20s}"
              f"{e(mx(lst,'TidU_gns')):>14s}{e(mx(lst,'TidU_kms')):>14s}{e(mx(lst,'T1_db'))}{e(mx(lst,'T1_aff'))}"
              f"{e(mx(lst,'T2_db'))}{e(mx(lst,'T2_aff'))}{e(mx(lst,'T3_sym_db'))}{e(mx(lst,'gauge_kdiff')):>12s}")
    print("  Theta_U = U Theta U^dag with U = prod_i exp(-i phi_i Z_i/2);  |T^phi - T^0| = max kernel difference vs phi = 0.")
    print()

    # ---------------- Table D
    print("TABLE D  S1 LOOP s0->s1->s2->s3->s0:  A = log(prod fwd / prod bwd)   (betaJ=0.7, g=g'=1, nu=1)")
    print(f"{'variant':15s}{'kappa':>6s}{'A[T_0.1]':>11s}{'A[T_1]':>11s}{'A[T_5]':>11s}{'A[G]':>11s}{'A[G_s]':>11s}"
          f"{'T1.DB':>9s}{'T1.DB(own)':>11s}{'T2.DB':>9s}")
    for conv, var in (("R1", "model"), ("R1", "gauge"), ("R1", PHI_HALF), ("R1", PHI_THIRD), ("R1", PHI_PI),
                      ("R1", "notie"), ("R3", "model")):
        for k in (0.5, 2.0, 10.0, 50.0):
            sel = [o for o in main_outs if o["spec"] == ("S1", conv, var, 0.7, 1.0, 1.0, k, True)]
            if not sel:
                continue
            o = sel[0]
            la = o.get("T1_loop", [np.nan] * 3)
            print(f"{label(('S1', conv, var)):15s}{k:6.1f}{la[0]:11.3e}{la[1]:11.3e}{la[2]:11.3e}"
                  f"{o.get('T2_loop', np.nan):11.3e}{o.get('T3_raw_loop', np.nan):11.3e}{o['T1_db']:9.1e}"
                  f"{o['T1_dbP']:11.1e}{o['T2_db']:9.1e}")
    print("  A[G_s] only where the slow manifold is separated (else nan).")
    print()

    # ---------------- Table E
    print("TABLE E  T3 SLOW-MANIFOLD GENERATOR vs kappa (betaJ=0.7, nu=1).  p = local slope -dlog(rawDB)/dlog(kappa)")
    print(f"{'system':7s}{'conv':5s}{'variant':14s}{'g':>5s}{'kappa':>7s}{'sep':>8s}{'|X-1|':>10s}{'|G-G_s|':>10s}"
          f"{'rawDB':>10s}{'raw|A|':>10s}{'symDB':>10s}{'sym|A|':>10s}{'minoff':>10s}{'p':>6s}{'C':>7s}  worst raw pair (dH,G_ba,G_ab)/max")
    series = {}
    for o in scale_outs:
        s = o["spec"]
        series.setdefault((s[0], s[1], s[2], s[4]), []).append(o)
    for sk in sorted(series, key=lambda s: (s[2] != "model" or s[1] == "R3", s[0], s[1], s[3])):
        lst = sorted(series[sk], key=lambda o: o["spec"][6])
        prev = None
        lab = label(sk)
        for o in lst:
            k = o["spec"][6]
            if "T3_raw_db" in o:
                p = ""
                if prev is not None and o["T3_raw_db"] > 0 and prev[1] > 0:
                    p = f"{-np.log(o['T3_raw_db'] / prev[1]) / np.log(k / prev[0]):6.2f}"
                print(f"{sk[0]:7s}{sk[1]:5s}{lab[:13]:14s}{sk[3]:5.1f}{k:7.0f}{o['T3_sep']:8.1f}{o['T3_slip']:10.2e}"
                      f"{o['T3_G_vs_Gs']:10.2e}{o['T3_raw_db']:10.2e}{o['T3_raw_aff']:10.2e}{o['T3_sym_db']:10.2e}"
                      f"{o['T3_sym_aff']:10.2e}{o['T3_minoff']:10.2e}{p:>6s}{o['T3_raw_db'] * (k / sk[3]) ** 2:7.2f}  "
                      f"({o['T3_worst'][0]},{o['T3_worst'][1]:+.1e},{o['T3_worst'][2]:+.1e})")
                prev = (k, o["T3_raw_db"])
            else:
                print(f"{sk[0]:7s}{sk[1]:5s}{lab[:13]:14s}{sk[3]:5.1f}{k:7.0f}{o.get('T3_sep', float('nan')):8.1f}"
                      "   (slow manifold not separated)")
    print("  |X-1| = size of the initial slip X = R_P W_P;  |G-G_s| = max|Schur G(0) - G_s| / max|G_s|.")
    print("  C = rawDB (kappa/g)^2.  |A| uses two-way edges only; a one-way or negative entry of G_s shows up in rawDB/minoff/worst pair.")
    print()
    print(f"errors: {len(errs)}   non-finite outputs: {nbad}   tasks: {len(outs)}   wall time: {time.time() - t0:.0f} s")
    for s in badspecs[:10]:
        print("  non-finite (loop affinity undefined: sign mismatch on the loop) at", s)
    for o in errs[:10]:
        print("  ERROR", o["spec"], o["error"])


if __name__ == "__main__":
    main()