#!/usr/bin/env python3
"""
zeno_rate.py -- P0 "notation freeze": the leading Zeno tie-flip coefficient.

FROZEN CONVENTION (hbar = 1)
  Gamma_m : one-spin coherence decay rate produced by monitoring.
            Monitoring jumps L_i = sqrt(Gamma_m/2) Z_i, so D[L_i] = (Gamma_m/2)(Z.Z - .)
            and |s><s^i| decays at Gamma_m.                      (= D40 "kappa")
  nu_th   : heat-bath clock, c_i(s) = nu_th / (1 + exp(2 beta J s_i h_i(s))).
                                                                  (= D40 "nu")
  g_i     : tie strength, K = sum_i g_i X_i 1[h_i = 0].           (= D40 "g_i")
  w       : symmetric coherent tie-flip rate (each direction), so an isolated
            tie spin's z decays at nu_th + 2w.

RESULTS VERIFIED BELOW
 (a) Isolated tie (D40 eq. 7.1). Slow eigenvalue lambda_s = -(nu_th + 2w) with
        w = [Delta - sqrt(Delta^2 - 16 g^2)] / 4 ,   Delta = Gamma_m - nu_th/2
          = 2g^2/Delta + 8g^4/Delta^3 + 64g^6/Delta^5 + ...
          = 2g^2/Gamma_m + g^2 nu_th/Gamma_m^2 + O(g^2 nu^2/G^3, g^4/G^3).
     Real (overdamped/Zeno) iff Gamma_m > nu_th/2 + 4g.
     The zero-frequency Schur complement instead gives w_0f = 2g^2/(Gamma_m + nu_th/2)
     exactly; eigen and Schur agree only at leading order 2g^2/Gamma_m (the first
     thermal correction has OPPOSITE sign in the two definitions).
 (b) Many body, zero-frequency Schur complement G = L_PP - L_PQ L_QQ^-1 L_QP:
        w_i(s -> s^i) = 2 g_i^2 / (Gamma_m + Gamma_th(s, s^i)) + O(g^4/Gamma_m^3)
     R1 (rank-one, fully configuration-resolved jumps):
        Gamma_th = [Lambda(s) + Lambda(s^i)]/2,  Lambda(s) = sum_{k live} c_k(s)
        -> EXTENSIVE (every thermal jump anywhere kills the coherence).
     R2 (D40 neighbour-resolved jumps): coherence |s><s^i| is carried, not killed,
        by thermal flips outside the closed star of i. Exact O(g^2) tie-edge rate
        2 g_i^2 [T^-1]_{x0,x0},  T = Gamma_m + Gamma_th^R1 - (env hopping),
        i.e. Feynman-Kac: 2g^2 int_0^inf E_x0[exp(-Gamma_m t - int_0^t Gamma_loc) ; X_t = x0] dt.
        Edge entry  = 2g^2/(Gamma_m + Gamma_th^R1) + O(g^2 nu^2/Gamma_m^3);
        the rest goes to correlated two-spin moves; the MARGINAL flip rate of i is
        2g^2/(Gamma_m + Gamma_loc) + O(g^2 nu^2/Gamma_m^3),
        Gamma_loc = nu_th/2 + (1/2) sum_{j~i live} [c_j(s) + c_j(s^i)]   (LOCAL).
     Universal leading coefficient in every convention:  w = 2 g^2 / Gamma_m.
     D40's alpha = 4 g^2/(nu_th Gamma_m) = (2g^2/Gamma_m)/(nu_th/2)
                 = leading Zeno tie rate / thermal tie rate.

Only numpy/scipy. Small dense linear algebra (d = 16, superoperator 256 x 256).
"""
import os
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_v, "1")          # small dense blocks: extra BLAS threads only add contention
import itertools
import warnings

import numpy as np

# numpy 2.0 + macOS Accelerate emits spurious "divide by zero/overflow/invalid value encountered in
# matmul" RuntimeWarnings on finite inputs; silence them and assert finiteness explicitly instead.
warnings.filterwarnings("ignore", message=r".*encountered in matmul", category=RuntimeWarning)

np.set_printoptions(linewidth=140)
RULE = "-" * 118

# ------------------------------------------------------------------ superoperators
# column-stacking vec: vec(A X B) = (B^T kron A) vec(X); |a><b| sits at a + d*b


def ham_super(H):
    d = H.shape[0]
    I = np.eye(d)
    return -1j * (np.kron(I, H) - np.kron(H.T, I))


def diss_super(Jm):
    d = Jm.shape[0]
    I = np.eye(d)
    JdJ = Jm.conj().T @ Jm
    return np.kron(Jm.conj(), Jm) - 0.5 * (np.kron(I, JdJ) + np.kron(JdJ.T, I))


def schur_population_generator(L, d):
    P = np.array([a + d * a for a in range(d)])
    Q = np.array([k for k in range(d * d) if k not in set(P)])
    LPP, LPQ = L[np.ix_(P, P)], L[np.ix_(P, Q)]
    LQP, LQQ = L[np.ix_(Q, P)], L[np.ix_(Q, Q)]
    G = LPP - LPQ @ np.linalg.solve(LQQ, LQP)
    assert np.isfinite(G).all(), "non-finite Schur complement"
    return G, P, Q


# ------------------------------------------------------------------ (a) one tie spin
X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.array([[1, 0], [0, -1]], dtype=complex)


def single_tie_L(g, nu, Gm):
    up = np.array([[0, 1], [0, 0]], dtype=complex)   # |+><-|
    dn = np.array([[0, 0], [1, 0]], dtype=complex)   # |-><+|
    L = ham_super(g * X)
    L += diss_super(np.sqrt(nu / 2) * up) + diss_super(np.sqrt(nu / 2) * dn)
    L += diss_super(np.sqrt(Gm / 2) * Z)
    return L


def bloch_matrix(L):
    """B[a,b] = (1/2) Tr(sigma_a L(sigma_b)), a,b in (x,y,z); affine part L(I)/2."""
    S = [X, Y, Z]
    B = np.zeros((3, 3))
    for b, sb in enumerate(S):
        Ls = (L @ sb.flatten(order="F")).reshape(2, 2, order="F")
        for a, sa in enumerate(S):
            B[a, b] = 0.5 * np.trace(sa @ Ls).real
    LI = (L @ np.eye(2).flatten(order="F")).reshape(2, 2, order="F")
    return B, np.abs(LI).max()


def w_closed(g, nu, Gm):
    D = Gm - nu / 2
    return (D - np.sqrt(D * D - 16 * g * g)) / 4


def part_a(checks):
    print("=" * 118)
    print("(a) ISOLATED TIE SPIN  (h = 0, H = gX; D40 eq. 7.1)   frozen: Gamma_m = kappa, nu_th = nu")
    print("=" * 118)
    g, nu = 0.3, 1.0
    kappas = [5, 10, 20, 50, 100]
    # verify the Bloch matrix of the Lindbladian built from the jumps equals eq. (7.1)
    worst_bloch = 0.0
    for Gm in kappas:
        B, aff = bloch_matrix(single_tie_L(g, nu, Gm))
        zy = np.array([[B[2, 2], B[2, 1]], [B[1, 2], B[1, 1]]])
        ref = np.array([[-nu, 2 * g], [-2 * g, -(nu / 2 + Gm)]])
        worst_bloch = max(worst_bloch, np.abs(zy - ref).max(), abs(B[0, 0] + nu / 2 + Gm),
                          np.abs(B[0, 1:]).max(), np.abs(B[1:, 0]).max(), aff)
    print(f"Bloch matrix from jumps vs eq.(7.1) [(z,y) block, x decoupled at -(nu/2+Gamma_m), unital]:"
          f" max |diff| = {worst_bloch:.2e}")
    checks.append(dict(name="(a) Lindbladian from jumps reproduces D40 eq. 7.1 Bloch matrix",
                       expected="0", obtained=f"{worst_bloch:.2e}", pass_=worst_bloch < 1e-12))

    print(f"parameters: g = {g}, nu_th = {nu};   w_eig := -(lambda_slow + nu_th)/2 from the full 4x4 Lindbladian")
    hdr = (f"{'Gamma_m':>7} {'w_eig':>13} {'|w_eig-closed|':>14} {'2g2/(G-nu/2)':>13} "
           f"{'(w_eig-2g2/D)/(8g4/D3)':>22} {'2g2/Gamma_m':>12} {'G*(w_eig/w0-1)':>14} "
           f"{'w_Schur0':>12} {'|w_S0-2g2/(G+nu/2)|':>19} {'G*(w_S0/w0-1)':>13}")
    print(hdr)
    print(RULE)
    worst_closed, worst_schur, ratios = 0.0, 0.0, []
    for Gm in kappas:
        L = single_tie_L(g, nu, Gm)
        ev = np.linalg.eigvals(L)
        ev = ev[np.argsort(-ev.real)]
        lam_s = ev[1]
        w_eig = (-lam_s.real - nu) / 2
        wc = w_closed(g, nu, Gm)
        D = Gm - nu / 2
        lead_D = 2 * g * g / D
        ratio = (w_eig - lead_D) / (8 * g ** 4 / D ** 3)
        w0 = 2 * g * g / Gm
        G, _, _ = schur_population_generator(L, 2)
        wS = G[1, 0].real - nu / 2
        wS_th = 2 * g * g / (Gm + nu / 2)
        worst_closed = max(worst_closed, abs(w_eig - wc) / wc, abs(lam_s.imag))
        worst_schur = max(worst_schur, abs(wS - wS_th) / wS_th)
        ratios.append(ratio)
        print(f"{Gm:7.0f} {w_eig:13.10f} {abs(w_eig - wc):14.2e} {lead_D:13.10f} {ratio:22.6f} "
              f"{w0:12.8f} {Gm * (w_eig / w0 - 1):14.6f} {wS:12.10f} {abs(wS - wS_th):19.2e} "
              f"{Gm * (wS / w0 - 1):13.6f}")
    print("reading: G*(w/w0-1) -> +nu_th/2 = +0.5 for the eigenvalue definition, -> -nu_th/2 = -0.5 for the")
    print("         zero-frequency Schur definition: same leading 2g^2/Gamma_m, opposite first thermal correction.")
    Gm_c = nu / 2 + 4 * g
    ev_c = np.linalg.eigvals(single_tie_L(g, nu, 0.8))
    print(f"overdamped (real-rate) threshold Gamma_m > nu_th/2 + 4g = {Gm_c:.3f}.  At Gamma_m = 0.8 (Attack-11 value)"
          f" the (z,y) eigenvalues are {np.sort_complex(ev_c)[0]:.4f}, {np.sort_complex(ev_c)[1]:.4f}: underdamped,"
          f" no Zeno rate.")
    checks.append(dict(name="(a) exact slow eigenvalue matches closed form w=[D-sqrt(D^2-16g^2)]/4, D=Gamma_m-nu_th/2",
                       expected="rel diff < 1e-10", obtained=f"{worst_closed:.2e}", pass_=worst_closed < 1e-10))
    checks.append(dict(name="(a) next term: (w_eig - 2g^2/D)/(8g^4/D^3) -> 1",
                       expected="1 (+O(8g^2/D^2))", obtained=", ".join(f"{r:.5f}" for r in ratios),
                       pass_=abs(ratios[-1] - 1) < 1e-3 and abs(ratios[0] - 1) < 0.05))
    checks.append(dict(name="(a) zero-frequency Schur rate equals 2g^2/(Gamma_m+nu_th/2) exactly",
                       expected="rel diff < 1e-10", obtained=f"{worst_schur:.2e}", pass_=worst_schur < 1e-10))


# ------------------------------------------------------------------ (b) 2x2 plaquette
class Plaquette:
    """Sites 0=(0,0) 1=(0,1) 2=(1,0) 3=(1,1); bonds 0-1, 2-3, 0-2, 1-3.
    Each site has two frozen exterior neighbours; e_i = their sum."""
    bonds = [(0, 1), (2, 3), (0, 2), (1, 3)]
    exterior = {0: (+1, -1), 1: (-1, +1), 2: (+1, +1), 3: (-1, -1)}

    def __init__(self, beta, J, nu, gs):
        self.N, self.d = 4, 16
        self.beta, self.J, self.nu, self.g = beta, J, nu, np.array(gs, float)
        self.nbr = {i: [] for i in range(4)}
        for a, b in self.bonds:
            self.nbr[a].append(b)
            self.nbr[b].append(a)
        self.e = np.array([sum(self.exterior[i]) for i in range(4)])
        self.spin = np.array([[1 - 2 * ((b >> k) & 1) for k in range(4)] for b in range(16)])
        s = self.spin
        self.h = np.array([[sum(s[b, j] for j in self.nbr[i]) + self.e[i] for i in range(4)]
                           for b in range(16)])
        self.E = np.array([-J * sum(s[b, i] * s[b, j] for i, j in self.bonds) - J * (self.e * s[b]).sum()
                           for b in range(16)], float)
        self.c = nu / (1 + np.exp(2 * beta * J * s * self.h))          # c[b, i]
        self.Lam = self.c.sum(axis=1)
        pi = np.exp(-beta * self.E)
        self.pi = pi / pi.sum()
        self.ties = [(b, i) for b in range(16) for i in range(4) if self.h[b, i] == 0]

    def K(self):
        K = np.zeros((16, 16))
        for b, i in self.ties:
            K[b ^ (1 << i), b] = self.g[i]
        return K

    def thermal_jumps(self, conv):
        Js = []
        if conv == "R1":
            for b in range(16):
                for i in range(4):
                    Jm = np.zeros((16, 16))
                    Jm[b ^ (1 << i), b] = np.sqrt(self.c[b, i])
                    Js.append(Jm)
        else:  # R2: |-s><s|_i (x) Pi_eta(neighbours) (x) 1
            for i in range(4):
                for s in (+1, -1):
                    for eta in itertools.product((+1, -1), repeat=len(self.nbr[i])):
                        Jm = np.zeros((16, 16))
                        for b in range(16):
                            if self.spin[b, i] == s and all(self.spin[b, j] == eta[t]
                                                            for t, j in enumerate(self.nbr[i])):
                                Jm[b ^ (1 << i), b] = np.sqrt(self.c[b, i])
                        Js.append(Jm)
        return Js

    def lindbladian(self, Gm, conv, gscale=1.0):
        Hd = np.diag(self.E)
        L = ham_super(Hd + gscale * self.K())
        for Jm in self.thermal_jumps(conv):
            L += diss_super(Jm)
        for i in range(4):
            Zi = np.diag(self.spin[:, i].astype(float))
            L += diss_super(np.sqrt(Gm / 2) * Zi)
        return L

    # ---------- analytic O(g^2) predictions
    def gamma_R1(self, b, i):
        return 0.5 * (self.Lam[b] + self.Lam[b ^ (1 << i)])

    def gamma_loc(self, b, i):
        star = [i] + self.nbr[i]
        bi = b ^ (1 << i)
        return 0.5 * sum(self.c[b, k] + self.c[bi, k] for k in star)

    def r2_resolvent(self, b, i, Gm):
        """Coherence |x,s><x,s^i| transported by env flips k outside closed star(i)."""
        star = set([i] + self.nbr[i])
        env = [k for k in range(4) if k not in star]
        cfg = []
        for bits in range(2 ** len(env)):
            bb = b
            for t, k in enumerate(env):
                if (bits >> t) & 1:
                    bb ^= 1 << k
            cfg.append(bb)
        idx = {bb: n for n, bb in enumerate(cfg)}
        T = np.zeros((len(cfg), len(cfg)))
        for bb in cfg:
            a = idx[bb]
            T[a, a] = Gm + self.gamma_R1(bb, i)
            for k in env:
                T[idx[bb ^ (1 << k)], a] -= self.c[bb, k]
        M = np.linalg.inv(T)
        return cfg, M[:, 0]          # column for source x0 = b (cfg[0] == b)


def slow_manifold_generator(L, d, P):
    ev, V = np.linalg.eig(L)
    order = np.argsort(-ev.real)
    lam = ev[order]
    gap = lam[d].real / lam[d - 1].real if lam[d - 1].real != 0 else np.inf
    Vs = V[:, order[:d]]
    PV = Vs[P, :]
    Gs = PV @ np.diag(lam[:d]) @ np.linalg.inv(PV)
    return Gs, gap


def db_violation(G, pi):
    F = G * pi[None, :]           # F[t, s] = G[t, s] pi[s] = stationary flux s -> t
    off = ~np.eye(len(pi), dtype=bool)
    return np.abs(F - F.T)[off].max() / np.abs(F[off]).max()


def part_b(checks):
    beta, J, nu = 0.6, 1.0, 1.0
    gs = (0.30, 0.40, 0.25, 0.35)
    kappas = [5, 10, 20, 50, 100]
    pq = Plaquette(beta, J, nu, gs)
    print()
    print("=" * 118)
    print("(b) 2x2 PLAQUETTE, frozen exterior e = " + str(tuple(int(x) for x in pq.e)) +
          f"  (exterior pairs {pq.exterior}),  beta J = {beta*J}, nu_th = {nu}, g_i = {gs}")
    print("=" * 118)
    print(f"directed tie edges (s, i) with h_i(s) = 0: {len(pq.ties)}  "
          f"(per site: {[sum(1 for b, i in pq.ties if i == k) for k in range(4)]})")
    print("thermal rate on every tie edge c_i = nu_th/2 = 0.5 (checked):",
          np.allclose([pq.c[b, i] for b, i in pq.ties], nu / 2))
    GR1 = sorted(set(round(pq.gamma_R1(b, i), 10) for b, i in pq.ties))
    GLo = sorted(set(round(pq.gamma_loc(b, i), 10) for b, i in pq.ties))
    print(f"distinct Gamma_th^R1 values: {[f'{x:.4f}' for x in GR1]}")
    print(f"distinct Gamma_loc^R2 values: {[f'{x:.4f}' for x in GLo]}")

    results = {}
    for conv in ("R1", "R2"):
        print()
        print(f"--- convention {conv} " + ("(rank-one, fully configuration-resolved flips; Attack-11 style)"
                                         if conv == "R1" else "(neighbour-resolved flips; D40 section 0)"))
        print(f"{'Gamma_m':>7} {'stat.res':>9} {'maxrel[edge vs exact O(g2)]':>27} "
              f"{'maxrel[edge vs 2g2/(G+G_R1)]':>28} {'maxrel[edge vs 2g2/G]':>21} "
              + (f"{'maxrel[marg vs 2g2/(G+G_loc)]':>29}" if conv == "R2" else f"{'maxrel[edge vs 2g2/(G+G_loc)]':>29}")
              + f" {'O(g2) analytic-vs-num':>21}")
        print(RULE)
        res = []
        for Gm in kappas:
            L = pq.lindbladian(Gm, conv)
            G, P, Q = schur_population_generator(L, 16)
            imag = np.abs(G.imag).max()
            G = G.real
            stat = np.abs(G @ pq.pi).max()
            # exact O(g^2) numerically: -L_PQ^K (L_QQ at g=0)^-1 L_QP^K
            L0 = pq.lindbladian(Gm, conv, gscale=0.0)
            LK = L - L0
            G2 = (-LK[np.ix_(P, Q)] @ np.linalg.solve(L0[np.ix_(Q, Q)], LK[np.ix_(Q, P)])).real
            e_O2 = e_R1 = e_univ = e_loc = e_an = 0.0
            e_dbl = 0.0
            r_abs = 0.0
            for b, i in pq.ties:
                bi = b ^ (1 << i)
                gi2 = pq.g[i] ** 2
                w_edge = G[bi, b] - pq.c[b, i]
                if conv == "R1":
                    w_O2 = 2 * gi2 / (Gm + pq.gamma_R1(b, i))
                    w_marg = w_edge
                    an_list = [(bi, w_O2)]
                else:
                    cfg, col = pq.r2_resolvent(b, i, Gm)
                    w_O2 = 2 * gi2 * col[0]
                    an_list = [(bb ^ (1 << i), 2 * gi2 * col[a]) for a, bb in enumerate(cfg)]
                    w_marg = sum(G[bb ^ (1 << i), b] for bb in cfg) - pq.c[b, i]
                for tgt, val in an_list:
                    e_an = max(e_an, abs(G2[tgt, b] - val) / abs(val))
                e_O2 = max(e_O2, abs(w_edge - w_O2) / w_O2)
                r_abs = max(r_abs, abs(w_edge - w_O2))
                w1 = 2 * gi2 / (Gm + pq.gamma_R1(b, i))
                e_R1 = max(e_R1, abs(w_edge - w1) / w1)
                w0 = 2 * gi2 / Gm
                e_univ = max(e_univ, abs(w_edge - w0) / w0)
                wl = 2 * gi2 / (Gm + pq.gamma_loc(b, i))
                e_loc = max(e_loc, abs(w_marg - wl) / wl)
            # structure diagnostics
            nd = [(t, s) for t in range(16) for s in range(16)
                  if bin(t ^ s).count("1") >= 2]
            multi = max(abs(G[t, s]) for t, s in nd)
            offmin = min(G[t, s] for t in range(16) for s in range(16) if t != s)
            db = db_violation(G, pq.pi)
            res.append(dict(Gm=Gm, e_O2=e_O2, r_abs=r_abs, e_R1=e_R1, e_univ=e_univ, e_loc=e_loc,
                            e_an=e_an, stat=stat, imag=imag, multi=multi, offmin=offmin, db=db, G=G))
            print(f"{Gm:7.0f} {stat:9.1e} {e_O2:27.3e} {e_R1:28.3e} {e_univ:21.3e} {e_loc:29.3e} {e_an:21.1e}")
        # scaling
        s_O2 = np.log(res[-2]["r_abs"] / res[-1]["r_abs"]) / np.log(res[-1]["Gm"] / res[-2]["Gm"])
        s_un = np.log(res[-2]["e_univ"] / res[-1]["e_univ"]) / np.log(res[-1]["Gm"] / res[-2]["Gm"])
        print(f"log-log slope (Gamma_m 50->100): |edge - exact O(g^2)| ~ Gamma_m^-{s_O2:.3f} (expect 3: O(g^4/G^3));"
              f"  rel. dev. from 2g^2/Gamma_m ~ Gamma_m^-{s_un:.3f} (expect 1)")
        # g-halving test of the O(g^4) residual at Gamma_m = 20
        Gm = 20.0
        rr = []
        for sc in (1.0, 0.5):
            L = pq.lindbladian(Gm, conv, gscale=sc)
            G, P, Q = schur_population_generator(L, 16)
            G = G.real
            r = 0.0
            for b, i in pq.ties:
                bi = b ^ (1 << i)
                gi2 = (sc * pq.g[i]) ** 2
                if conv == "R1":
                    wO2 = 2 * gi2 / (Gm + pq.gamma_R1(b, i))
                else:
                    wO2 = 2 * gi2 * pq.r2_resolvent(b, i, Gm)[1][0]
                r = max(r, abs(G[bi, b] - pq.c[b, i] - wO2))
            rr.append(r)
        print(f"g -> g/2 at Gamma_m = 20: max|edge - exact O(g^2)| {rr[0]:.3e} -> {rr[1]:.3e}, "
              f"ratio {rr[0]/rr[1]:.3f} (expect ~16 = O(g^4))")
        print(f"structure: max|G| on >=2-spin moves: " +
              ", ".join(f"{r['multi']:.2e}" for r in res) +
              "  |  min off-diagonal G: " + ", ".join(f"{r['offmin']:+.2e}" for r in res))
        print(f"pi-detailed-balance violation of G_Schur (rel.): " + ", ".join(f"{r['db']:.1e}" for r in res) +
              f"  |  max Im G: {max(r['imag'] for r in res):.1e}")
        results[conv] = dict(res=res, s_O2=s_O2, s_un=s_un, g_ratio=rr[0] / rr[1])

    # per-edge Gamma_eff := 2 g_i^2 / w - Gamma_m  (measured thermal contribution), Gamma_m = 100
    print()
    print("Per-edge measured thermal contribution Gamma_eff := 2 g_i^2 / w - Gamma_m at Gamma_m = 100 (6 representative edges)")
    print(f"{'edge (config bits, site)':>26} {'Gamma_R1':>9} {'Gamma_loc':>9} {'R1: edge':>9} {'R2: edge':>9} {'R2: marginal':>12}")
    print(RULE)
    Gm = 100.0
    GR1m = results["R1"]["res"][-1]["G"]
    GR2m = results["R2"]["res"][-1]["G"]
    shown, seen = 0, set()
    for b, i in pq.ties:
        key = (round(pq.gamma_R1(b, i), 6), round(pq.gamma_loc(b, i), 6))
        if key in seen:
            continue
        seen.add(key)
        bi = b ^ (1 << i)
        gi2 = pq.g[i] ** 2
        g1 = 2 * gi2 / (GR1m[bi, b] - pq.c[b, i]) - Gm
        g2 = 2 * gi2 / (GR2m[bi, b] - pq.c[b, i]) - Gm
        cfg, _ = pq.r2_resolvent(b, i, Gm)
        wm = sum(GR2m[bb ^ (1 << i), b] for bb in cfg) - pq.c[b, i]
        g3 = 2 * gi2 / wm - Gm
        cfgs = "".join("+" if s > 0 else "-" for s in pq.spin[b])
        print(f"{cfgs + ' , i=' + str(i):>26} {pq.gamma_R1(b, i):9.4f} {pq.gamma_loc(b, i):9.4f} "
              f"{g1:9.4f} {g2:9.4f} {g3:12.4f}")
        shown += 1
        if shown == 6:
            break
    print("reading: R1 edge -> Gamma_R1 (extensive sum over all live sites); R2 edge -> Gamma_R1 too at this order,")
    print("         but the R2 MARGINAL (edge + correlated 2-spin moves) -> the local Gamma_loc.  Residuals are O(1/Gamma_m).")

    # bonus: exact slow-manifold population generator (frequency-shifted reduction)
    print()
    print("Bonus diagnostic: slow-manifold population generator G_slow = (P V) diag(lambda_slow) (P V)^-1 (16 slowest modes)")
    print(f"{'conv':>4} {'Gamma_m':>7} {'gap Re l17/Re l16':>17} {'max G*|w_slow/w0-1|':>19} "
          f"{'R1 pred. maxrel':>15} {'pi-DB violation':>15}")
    print(RULE)
    for conv in ("R1", "R2"):
        for Gm in (20.0, 50.0, 100.0):
            L = pq.lindbladian(Gm, conv)
            _, P, _ = schur_population_generator(L, 16)
            Gs, gap = slow_manifold_generator(L, 16, P)
            Gs = Gs.real
            dev, pred = 0.0, 0.0
            for b, i in pq.ties:
                bi = b ^ (1 << i)
                gi2 = pq.g[i] ** 2
                ws = Gs[bi, b] - pq.c[b, i]
                dev = max(dev, Gm * abs(ws / (2 * gi2 / Gm) - 1))
                if conv == "R1":   # w_slow = 2g^2/G + (g^2/G^2)(Lam(s)-Lam(s^i)) + (nu/G^2) sum_{j tie at s^i} g_j^2
                    tj = sum(pq.g[j] ** 2 for j in range(4) if pq.h[bi, j] == 0)
                    wp = 2 * gi2 / Gm + gi2 / Gm ** 2 * (pq.Lam[b] - pq.Lam[bi]) + nu / Gm ** 2 * tj
                    pred = max(pred, abs(ws - wp) / ws)
            db = db_violation(Gs, pq.pi)
            print(f"{conv:>4} {Gm:7.0f} {gap:17.2f} {dev:19.4f} "
                  f"{(f'{pred:.2e}' if conv == 'R1' else 'n/a'):>15} {db:15.2e}")
    print("reading: the O(g^2 nu_th/Gamma_m^2) correction is definition dependent (zero-frequency Schur vs slow manifold);")
    print("         only the leading 2 g^2/Gamma_m is convention- and definition-free.")

    r1, r2 = results["R1"], results["R2"]
    checks.append(dict(name="(b) G_Schur stationary on Gibbs pi (both conventions, all Gamma_m)",
                       expected="< 1e-12", obtained=f"{max(r['stat'] for c in (r1, r2) for r in c['res']):.1e}",
                       pass_=max(r['stat'] for c in (r1, r2) for r in c['res']) < 1e-12))
    checks.append(dict(name="(b) analytic exact-O(g^2) rates (R1: 2g^2/(G+Gamma_R1); R2: Feynman-Kac resolvent) vs numeric O(g^2) Schur term",
                       expected="rel diff < 1e-9",
                       obtained=f"R1 {max(r['e_an'] for r in r1['res']):.1e}, R2 {max(r['e_an'] for r in r2['res']):.1e}",
                       pass_=max(r['e_an'] for c in (r1, r2) for r in c['res']) < 1e-9))
    for name, c in (("R1", r1), ("R2", r2)):
        checks.append(dict(name=f"(b) {name}: max rel err of tie-edge excess vs exact-O(g^2) formula, Gamma_m=5,10,20,50,100",
                           expected="-> 0 like Gamma_m^-2 (abs Gamma_m^-3)",
                           obtained=", ".join(f"{r['e_O2']:.2e}" for r in c["res"]) + f"; abs slope {c['s_O2']:.3f}",
                           pass_=abs(c["s_O2"] - 3) < 0.15 and c["res"][-1]["e_O2"] < 1e-4))
        checks.append(dict(name=f"(b) {name}: max rel dev of tie-edge excess from universal 2g_i^2/Gamma_m",
                           expected="-> 0 like Gamma_m^-1",
                           obtained=", ".join(f"{r['e_univ']:.3e}" for r in c["res"]) + f"; slope {c['s_un']:.3f}",
                           pass_=abs(c["s_un"] - 1) < 0.1))
        checks.append(dict(name=f"(b) {name}: O(g^4) residual scaling under g -> g/2 at Gamma_m=20",
                           expected="ratio ~16", obtained=f"{c['g_ratio']:.3f}",
                           pass_=abs(c['g_ratio'] - 16) < 1.0))
    checks.append(dict(name="(b) R1: tie-edge excess vs 2g^2/(Gamma_m+Gamma_R1), Gamma_R1=(Lambda(s)+Lambda(s^i))/2",
                       expected="rel err O(g^2/Gamma_m^2)",
                       obtained=", ".join(f"{r['e_R1']:.2e}" for r in r1["res"]),
                       pass_=r1["res"][-1]["e_R1"] < 1e-4))
    checks.append(dict(name="(b) R2: marginal flip-i rate vs 2g^2/(Gamma_m+Gamma_loc), local Gamma_loc",
                       expected="rel err O(nu^2/Gamma_m^2)",
                       obtained=", ".join(f"{r['e_loc']:.2e}" for r in r2["res"]),
                       pass_=r2["res"][-1]["e_loc"] < 1e-3))
    checks.append(dict(name="(b) G_Schur Gibbs detailed balance (diagnostic)",
                       expected="report", obtained=f"R1 max {max(r['db'] for r in r1['res']):.1e}, R2 max {max(r['db'] for r in r2['res']):.1e}",
                       pass_=True))
    return results


def main():
    checks = []
    part_a(checks)
    part_b(checks)
    print()
    print("=" * 118)
    print("CHECK SUMMARY")
    print("=" * 118)
    for c in checks:
        print(f"[{'PASS' if c['pass_'] else 'FAIL'}] {c['name']}\n        expected {c['expected']:<32} obtained {c['obtained']}")
    print(f"\n{sum(c['pass_'] for c in checks)}/{len(checks)} checks passed")


if __name__ == "__main__":
    main()