#!/usr/bin/env python3
"""Local-stress observable v1 — observation-only addition on unchanged CPUREF-0.3.

Frozen spec: 18_local_stress_observable_spec_v1.md (Drive 1IgggnRfJrdvRtRo-07pmpUtILRjFmY2V)
Thread: C0C78UFKGRM / 1791869055.115119

Unchanged from shear_diagnostic_v0.2: same 7x32 runs, same 22 hard checks, same
resource bounds, reference file unmodified. Addition: at coarse step 32, from
ARRIVAL (post-stream, pre-collision) populations, report scalar R and one signed
horizontal-face profile. Diagnostic values only — no thresholds / mechanism claims.
"""
from __future__ import annotations

import hashlib
import importlib.util
import os
import sys
import time
from dataclasses import dataclass

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
REF_PATH = os.environ.get(
    "CPU_REFERENCE_PATH", os.path.join(HERE, "cpu_reference_v0.3.py")
)

U0 = 1e-3
N_COARSE = 16
N_FINE_DOM = 32
PATCH_LO, PATCH_HI = 4, 12
TAU_C, TAU_F = 0.8, 1.1
N_STEPS = 32
FIT_LO, FIT_HI = 4, 32
CONS_TOL = 1e-12
TRANS_TOL = 1e-12 * U0
INIT_TOL = CONS_TOL
T0_TOL = TRANS_TOL
K_PHYS = 2.0 * np.pi / 16.0
CS2 = 1.0 / 3.0
NU = CS2 * (TAU_C - 0.5)  # 0.1 coarse units
RHO0 = 1.0

CX = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1], dtype=np.int64)
CY = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1], dtype=np.int64)
W = np.array([4 / 9] + [1 / 9] * 4 + [1 / 36] * 4, dtype=np.float64)
NQ = 9


def _load_ref():
    spec = importlib.util.spec_from_file_location("cpu_ref_v03", REF_PATH)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = mod
    spec.loader.exec_module(mod)
    return mod


def sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 16), b""):
            h.update(chunk)
    return h.hexdigest()


def sinc(x):
    x = np.asarray(x, dtype=np.float64)
    out = np.empty_like(x, dtype=np.float64)
    small = np.abs(x) < 1e-14
    out[small] = 1.0
    xs = x[~small]
    out[~small] = np.sin(xs) / xs
    return float(out) if out.ndim == 0 else out


def feq(rho, ux, uy):
    rho = np.asarray(rho, dtype=np.float64)
    ux = np.asarray(ux, dtype=np.float64)
    uy = np.asarray(uy, dtype=np.float64)
    cu = CX[:, None, None] * ux + CY[:, None, None] * uy
    usq = ux * ux + uy * uy
    return W[:, None, None] * rho * (1.0 + 3.0 * cu + 4.5 * cu * cu - 1.5 * usq)


def srt_collide(f, tau, mask=None):
    g = f.copy()
    if mask is None:
        mask = np.ones(f.shape[1:], dtype=bool)
    fm = f[:, mask]
    rho = fm.sum(axis=0)
    ux = (CX[:, None] * fm).sum(axis=0) / rho
    uy = (CY[:, None] * fm).sum(axis=0) / rho
    rho_g = np.zeros(f.shape[1:])
    ux_g = np.zeros(f.shape[1:])
    uy_g = np.zeros(f.shape[1:])
    rho_g[mask] = rho
    ux_g[mask] = ux
    uy_g[mask] = uy
    eq_full = feq(rho_g, ux_g, uy_g)
    g[:, mask] = fm - (fm - eq_full[:, mask]) / tau
    return g


def stream_periodic(g):
    out = np.empty_like(g)
    for i in range(NQ):
        out[i] = np.roll(g[i], (int(CY[i]), int(CX[i])), axis=(0, 1))
    return out


def moments(f):
    rho = f.sum(axis=0)
    jx = (CX[:, None, None] * f).sum(axis=0)
    jy = (CY[:, None, None] * f).sum(axis=0)
    return rho, jx, jy


def cell_avg_ux(y_centers, dx, phi):
    factor = float(sinc(K_PHYS * dx / 2.0))
    return U0 * factor * np.sin(K_PHYS * y_centers + phi)


def mode_template(phi):
    yc = np.arange(N_COARSE, dtype=np.float64) + 0.5
    m1d = float(sinc(K_PHYS / 2.0)) * np.sin(K_PHYS * yc + phi)
    return np.broadcast_to(m1d[:, None], (N_COARSE, N_COARSE)).copy()


def project_amp(ux, phi):
    m = mode_template(phi)
    return float(np.sum(ux * m) / np.sum(m * m))


def restrict_fine_to_coarse(rho_f, jx_f, jy_f):
    assert rho_f.shape == (N_FINE_DOM, N_FINE_DOM)
    rho = rho_f.reshape(N_COARSE, 2, N_COARSE, 2).mean(axis=(1, 3))
    jx = jx_f.reshape(N_COARSE, 2, N_COARSE, 2).mean(axis=(1, 3))
    jy = jy_f.reshape(N_COARSE, 2, N_COARSE, 2).mean(axis=(1, 3))
    return rho, jx, jy, jx / rho, jy / rho


def volume_avg_2x2(field_f):
    """2x2 volume-average of a fine 32^2 field onto 16^2."""
    assert field_f.shape == (N_FINE_DOM, N_FINE_DOM)
    return field_f.reshape(N_COARSE, 2, N_COARSE, 2).mean(axis=(1, 3))


def sigma_xy_cell(f, tau):
    """Per-cell corrected stress from ARRIVAL populations.
    Π^neq_xy = Σ c_x c_y f − ρ u_x u_y
    σ_xy = −(1 − 1/(2τ)) Π^neq_xy
    No extra dx/dt factor ((dx/dt)^2 = 1).
    """
    rho, jx, jy = moments(f)
    ux = jx / rho
    uy = jy / rho
    pi_neq = (CX[:, None, None] * CY[:, None, None] * f).sum(axis=0) - rho * ux * uy
    sigma = -(1.0 - 1.0 / (2.0 * tau)) * pi_neq
    return sigma, rho, ux, uy


def project_a_u_P_sigma(u_field, sigma, phi, transpose=False):
    """Fourier projections over all 16^2 cells (spec §3).
    a_u = (2/256) Σ u · sin(k· + φ);  P_σ = (2/256) Σ σ · cos(k· + φ)
    Non-transpose: u_x vs y; transpose: u_y vs x.
    """
    centers = np.arange(N_COARSE, dtype=np.float64) + 0.5
    if transpose:
        # u_y(x): template varies with x (columns)
        sin_t = np.sin(K_PHYS * centers[None, :] + phi)
        cos_t = np.cos(K_PHYS * centers[None, :] + phi)
    else:
        # u_x(y): template varies with y (rows)
        sin_t = np.sin(K_PHYS * centers[:, None] + phi)
        cos_t = np.cos(K_PHYS * centers[:, None] + phi)
    a_u = (2.0 / 256.0) * float(np.sum(u_field * sin_t))
    P_s = (2.0 / 256.0) * float(np.sum(sigma * cos_t))
    return a_u, P_s


def profile_s(sigma, a_u, transpose=False):
    """Signed horizontal-face profile: mean over x∈[4,12), / (ρ0 ν k a_u).
    Transpose: no profile (returns None).
    """
    if transpose:
        return None
    denom = RHO0 * NU * K_PHYS * a_u
    row_mean = sigma[:, PATCH_LO:PATCH_HI].mean(axis=1)  # j=0..15
    return row_mean / denom


@dataclass
class Series:
    name: str
    phi: float
    A: np.ndarray
    e_if: np.ndarray
    e_cb: np.ndarray
    e_if_patch: np.ndarray
    e_if_out: np.ndarray
    nonmodal: np.ndarray
    max_uy: np.ndarray
    max_drho: np.ndarray
    mass: np.ndarray
    momx: np.ndarray
    momy: np.ndarray
    pi_xy_mean: np.ndarray
    finite_ok: bool
    cons_ok: bool
    mass_drift: float = 0.0
    momx_drift: float = 0.0
    momy_drift: float = 0.0
    notes: str = ""


def init_uniform(n, dx, phi, transpose=False):
    centers = (np.arange(n, dtype=np.float64) + 0.5) * dx
    Y, X = np.meshgrid(centers, centers, indexing="ij")
    rho = np.ones((n, n), dtype=np.float64)
    if transpose:
        ux = np.zeros_like(rho)
        uy = cell_avg_ux(X, dx, phi)
    else:
        ux = cell_avg_ux(Y, dx, phi)
        uy = np.zeros_like(rho)
    return feq(rho, ux, uy)


def advance_uniform(f, tau, steps):
    for _ in range(steps):
        f = stream_periodic(srt_collide(f, tau))
    return f


def observe_uniform(f, dx, phi, transpose=False):
    rho, jx, jy = moments(f)
    if dx < 1.0 - 1e-12:
        rho_c, jx_c, jy_c, ux, uy = restrict_fine_to_coarse(rho, jx, jy)
    else:
        rho_c, jx_c, jy_c = rho, jx, jy
        ux, uy = jx / rho, jy / rho
    if transpose:
        xc = np.arange(N_COARSE, dtype=np.float64) + 0.5
        mx = float(sinc(K_PHYS / 2.0)) * np.sin(K_PHYS * xc + phi)
        mT = np.broadcast_to(mx[None, :], (N_COARSE, N_COARSE)).copy()
        amp = float(np.sum(uy * mT) / np.sum(mT * mT))
        modal = amp * mT
        nonmodal = float(np.linalg.norm(uy - modal) / U0)
        max_cross = float(np.max(np.abs(ux)) / U0)
    else:
        amp = project_amp(ux, phi)
        m = mode_template(phi)
        modal = amp * m
        nonmodal = float(np.linalg.norm(ux - modal) / U0)
        max_cross = float(np.max(np.abs(uy)) / U0)
    return {
        "amp": amp,
        "ux": ux,
        "uy": uy,
        "rho": rho_c,
        "jx": jx_c,
        "jy": jy_c,
        "nonmodal": nonmodal,
        "max_cross": max_cross,
        "max_drho": float(np.max(np.abs(rho_c - 1.0))),
        "mass": float(rho_c.sum()),
        "momx": float(jx_c.sum()),
        "momy": float(jy_c.sum()),
    }


def run_uniform(kind, phi, transpose=False):
    if kind == "B":
        n, dx, tau, sub = N_FINE_DOM, 0.5, TAU_F, 2
    else:
        n, dx, tau, sub = N_COARSE, 1.0, TAU_C, 1
    f = init_uniform(n, dx, phi, transpose=transpose)
    amps, nonmodal, max_cross, max_drho = [], [], [], []
    mass, momx, momy = [], [], []
    fields = []
    finite_ok = True
    for tick in range(N_STEPS + 1):
        if not np.isfinite(f).all():
            finite_ok = False
        obs = observe_uniform(f, dx, phi, transpose=transpose)
        amps.append(obs["amp"])
        nonmodal.append(obs["nonmodal"])
        max_cross.append(obs["max_cross"])
        max_drho.append(obs["max_drho"])
        mass.append(obs["mass"])
        momx.append(obs["momx"])
        momy.append(obs["momy"])
        fields.append(
            (obs["ux"].copy(), obs["uy"].copy(), obs["rho"].copy(),
             obs["jx"].copy(), obs["jy"].copy())
        )
        if tick < N_STEPS:
            f = advance_uniform(f, tau, sub)
    # After N_STEPS advances, f is ARRIVAL at coarse step 32 (post-stream).
    arrival_name = (
        f"uniform_{kind}_f_after_{N_STEPS * sub}_stream_steps"
        f"_(=coarse_step_{N_STEPS}_arrival; collide-then-stream loop)"
    )
    mass = np.asarray(mass)
    momx = np.asarray(momx)
    momy = np.asarray(momy)
    mass_drift = float(np.max(np.abs(mass - mass[0])))
    momx_drift = float(np.max(np.abs(momx - momx[0])))
    momy_drift = float(np.max(np.abs(momy - momy[0])))
    tol_abs = CONS_TOL * abs(mass[0])
    cons_ok = (
        mass_drift <= tol_abs
        and momx_drift <= tol_abs
        and momy_drift <= tol_abs
    )
    ser = Series(
        name=f"{kind}_phi{phi:.4f}{'_T' if transpose else ''}",
        phi=phi,
        A=np.asarray(amps),
        e_if=np.full(N_STEPS + 1, np.nan),
        e_cb=np.full(N_STEPS + 1, np.nan),
        e_if_patch=np.full(N_STEPS + 1, np.nan),
        e_if_out=np.full(N_STEPS + 1, np.nan),
        nonmodal=np.asarray(nonmodal),
        max_uy=np.asarray(max_cross),
        max_drho=np.asarray(max_drho),
        mass=mass,
        momx=momx,
        momy=momy,
        pi_xy_mean=np.full(N_STEPS + 1, np.nan),
        finite_ok=finite_ok,
        cons_ok=bool(cons_ok),
        mass_drift=mass_drift,
        momx_drift=momx_drift,
        momy_drift=momy_drift,
    )
    return ser, fields, f, tau, arrival_name


def init_composite(ref, spec, maps, phi, transpose=False):
    n, nf = spec.n_coarse, spec.n_fine
    coarse = np.full((NQ, n, n), np.nan)
    fine = np.full((NQ, nf, nf), np.nan)
    yc = np.arange(n, dtype=np.float64) + 0.5
    xc = np.arange(n, dtype=np.float64) + 0.5
    Yc, Xc = np.meshgrid(yc, xc, indexing="ij")
    if transpose:
        ux_c = np.zeros((n, n))
        uy_c = cell_avg_ux(Xc, 1.0, phi)
    else:
        ux_c = cell_avg_ux(Yc, 1.0, phi)
        uy_c = np.zeros((n, n))
    coarse[:, :, :] = feq(np.ones((n, n)), ux_c, uy_c)
    coarse[:, maps.coarse_covered] = np.nan
    fy = maps.fine_centre
    fx = maps.fine_centre
    Yf, Xf = np.meshgrid(fy, fx, indexing="ij")
    if transpose:
        ux_f = np.zeros((nf, nf))
        uy_f = cell_avg_ux(Xf, 0.5, phi)
    else:
        ux_f = cell_avg_ux(Yf, 0.5, phi)
        uy_f = np.zeros((nf, nf))
    eq_f = feq(np.ones((nf, nf)), ux_f, uy_f)
    fine[:, maps.fine_active] = eq_f[:, maps.fine_active]
    return ref.State(coarse_f=coarse, fine_f=fine, phase="arrival_f", tick=0)


def observe_composite(ref, maps, state, phi, transpose=False):
    fc, ff = state.coarse_f, state.fine_f
    n = N_COARSE
    rho_c = np.full((n, n), np.nan)
    jx_c = np.full((n, n), np.nan)
    jy_c = np.full((n, n), np.nan)
    for y in range(n):
        for x in range(n):
            if maps.coarse_active[y, x]:
                sl = fc[:, y, x]
                rho_c[y, x] = sl.sum()
                jx_c[y, x] = (CX * sl).sum()
                jy_c[y, x] = (CY * sl).sum()
    for ky in range(PATCH_LO, PATCH_HI):
        for kx in range(PATCH_LO, PATCH_HI):
            y0 = spec_children_start(ky)
            x0 = spec_children_start(kx)
            kids = []
            for dy in (0, 1):
                for dx in (0, 1):
                    sl = ff[:, y0 + dy, x0 + dx]
                    kids.append((sl.sum(), (CX * sl).sum(), (CY * sl).sum()))
            rho_c[ky, kx] = np.mean([k[0] for k in kids])
            jx_c[ky, kx] = np.mean([k[1] for k in kids])
            jy_c[ky, kx] = np.mean([k[2] for k in kids])
    ux = jx_c / rho_c
    uy = jy_c / rho_c
    if transpose:
        xc = np.arange(n, dtype=np.float64) + 0.5
        mx = float(sinc(K_PHYS / 2.0)) * np.sin(K_PHYS * xc + phi)
        mT = np.broadcast_to(mx[None, :], (n, n)).copy()
        amp = float(np.sum(uy * mT) / np.sum(mT * mT))
        modal = amp * mT
        nonmodal = float(np.linalg.norm(uy - modal) / U0)
        max_cross = float(np.max(np.abs(ux)) / U0)
    else:
        amp = project_amp(ux, phi)
        m = mode_template(phi)
        modal = amp * m
        nonmodal = float(np.linalg.norm(ux - modal) / U0)
        max_cross = float(np.max(np.abs(uy)) / U0)
    leaf = ref._leaf_totals(fc, ff, maps)
    return {
        "amp": amp,
        "ux": ux,
        "uy": uy,
        "rho": rho_c,
        "jx": jx_c,
        "jy": jy_c,
        "nonmodal": nonmodal,
        "max_cross": max_cross,
        "max_drho": float(np.max(np.abs(rho_c - 1.0))),
        "mass": float(leaf[0]),
        "momx": float(leaf[1]),
        "momy": float(leaf[2]),
        "finite_active": bool(
            np.isfinite(fc[:, maps.coarse_active]).all()
            and np.isfinite(ff[:, maps.fine_active]).all()
        ),
    }


def spec_children_start(k):
    return 4 + 2 * (k - PATCH_LO)


def sigma_composite_16(maps, state):
    """Build 16^2 corrected σ from arrival_f composite state.
    Patch cells: 2x2 volume-average of fine σ (τ_f). Outside: coarse σ (τ_c).
    Never uses coarse underlay under the patch.
    """
    fc, ff = state.coarse_f, state.fine_f
    n = N_COARSE
    sig = np.full((n, n), np.nan)
    # outside / active coarse
    for y in range(n):
        for x in range(n):
            if maps.coarse_active[y, x]:
                f1 = fc[:, y : y + 1, x : x + 1]
                s1, _, _, _ = sigma_xy_cell(f1, TAU_C)
                sig[y, x] = float(s1[0, 0])
    # patch from fine
    for ky in range(PATCH_LO, PATCH_HI):
        for kx in range(PATCH_LO, PATCH_HI):
            y0 = spec_children_start(ky)
            x0 = spec_children_start(kx)
            kids = []
            for dy in (0, 1):
                for dx in (0, 1):
                    f1 = ff[:, y0 + dy : y0 + dy + 1, x0 + dx : x0 + dx + 1]
                    s1, _, _, _ = sigma_xy_cell(f1, TAU_F)
                    kids.append(float(s1[0, 0]))
            sig[ky, kx] = float(np.mean(kids))
    return sig


def run_composite(ref, phi, transpose=False):
    spec = ref.make_spec(collision="srt")
    assert abs(spec.tau_coarse - TAU_C) < 1e-15 and abs(spec.tau_fine - TAU_F) < 1e-15
    maps = ref.build_maps(spec)
    st = init_composite(ref, spec, maps, phi, transpose=transpose)
    yc = np.arange(N_COARSE) + 0.5
    xc = np.arange(N_COARSE) + 0.5
    init_mismatch = 0.0
    for ky in range(PATCH_LO, PATCH_HI):
        for kx in range(PATCH_LO, PATCH_HI):
            y0 = spec_children_start(ky)
            x0 = spec_children_start(kx)
            rhos, jxs, jys = [], [], []
            for dy in (0, 1):
                for dx in (0, 1):
                    sl = st.fine_f[:, y0 + dy, x0 + dx]
                    rhos.append(sl.sum())
                    jxs.append((CX * sl).sum())
                    jys.append((CY * sl).sum())
            rho_r, jx_r, jy_r = np.mean(rhos), np.mean(jxs), np.mean(jys)
            if transpose:
                ux, uy = 0.0, float(cell_avg_ux(xc[kx], 1.0, phi))
            else:
                ux, uy = float(cell_avg_ux(yc[ky], 1.0, phi)), 0.0
            init_mismatch = max(
                init_mismatch, abs(rho_r - 1.0), abs(jx_r - ux), abs(jy_r - uy)
            )
    amps, nonmodal, max_cross, max_drho = [], [], [], []
    mass, momx, momy = [], [], []
    fields = []
    finite_ok = True
    for tick in range(N_STEPS + 1):
        obs = observe_composite(ref, maps, st, phi, transpose=transpose)
        if not obs["finite_active"]:
            finite_ok = False
        amps.append(obs["amp"])
        nonmodal.append(obs["nonmodal"])
        max_cross.append(obs["max_cross"])
        max_drho.append(obs["max_drho"])
        mass.append(obs["mass"])
        momx.append(obs["momx"])
        momy.append(obs["momy"])
        fields.append(
            (obs["ux"].copy(), obs["uy"].copy(), obs["rho"].copy(),
             obs["jx"].copy(), obs["jy"].copy())
        )
        if tick < N_STEPS:
            st, _ = ref.advance_coarse_once(spec, maps, st)
    # After 32 advances: tick=64, phase=arrival_f — post-stream, pre-collision.
    arrival_name = (
        f"State.coarse_f + State.fine_f at phase={st.phase!r} tick={st.tick} "
        f"(after {N_STEPS} advance_coarse_once; arrival_f = post-stream pre-collision)"
    )
    mass = np.asarray(mass)
    momx = np.asarray(momx)
    momy = np.asarray(momy)
    mass_drift = float(np.max(np.abs(mass - mass[0])))
    momx_drift = float(np.max(np.abs(momx - momx[0])))
    momy_drift = float(np.max(np.abs(momy - momy[0])))
    tol_abs = CONS_TOL * max(abs(mass[0]), 1.0)
    cons_ok = (
        mass_drift <= tol_abs
        and momx_drift <= tol_abs
        and momy_drift <= tol_abs
    )
    ser = Series(
        name=f"A_phi{phi:.4f}{'_T' if transpose else ''}",
        phi=phi,
        A=np.asarray(amps),
        e_if=np.full(N_STEPS + 1, np.nan),
        e_cb=np.full(N_STEPS + 1, np.nan),
        e_if_patch=np.full(N_STEPS + 1, np.nan),
        e_if_out=np.full(N_STEPS + 1, np.nan),
        nonmodal=np.asarray(nonmodal),
        max_uy=np.asarray(max_cross),
        max_drho=np.asarray(max_drho),
        mass=mass,
        momx=momx,
        momy=momy,
        pi_xy_mean=np.full(N_STEPS + 1, np.nan),
        finite_ok=finite_ok,
        cons_ok=bool(cons_ok),
        mass_drift=mass_drift,
        momx_drift=momx_drift,
        momy_drift=momy_drift,
        notes=f"init_rho_j_mismatch_vs_coarse_formula={init_mismatch:.3e}",
    )
    return ser, fields, st, maps, init_mismatch, arrival_name


def fit_lambda(amps):
    t = np.arange(FIT_LO, FIT_HI + 1, dtype=np.float64)
    y = np.log(np.abs(amps[FIT_LO : FIT_HI + 1]))
    n = len(t)
    tbar, ybar = t.mean(), y.mean()
    sxx = np.sum((t - tbar) ** 2)
    slope = np.sum((t - tbar) * (y - ybar)) / sxx
    intercept = ybar - slope * tbar
    pred = intercept + slope * t
    resid = y - pred
    dof = max(n - 2, 1)
    s_err = float(np.sqrt(np.sum(resid**2) / dof / sxx))
    ss_tot = np.sum((y - ybar) ** 2)
    r2 = float(1.0 - np.sum(resid**2) / ss_tot) if ss_tot > 0 else 1.0
    return float(-slope), s_err, r2


def field_l2(a, b):
    return float(np.linalg.norm(a - b))


def stress_from_uniform_arrival(f, tau, phi, transpose=False):
    """σ on 16^2 from uniform arrival f; also ux/uy restricted for a_u."""
    sig, rho, ux, uy = sigma_xy_cell(f, tau)
    if f.shape[1] == N_FINE_DOM:
        sig16 = volume_avg_2x2(sig)
        rho_f, jx_f, jy_f = moments(f)
        _, _, _, ux16, uy16 = restrict_fine_to_coarse(rho_f, jx_f, jy_f)
    else:
        sig16 = sig
        ux16, uy16 = ux, uy
    u_proj = uy16 if transpose else ux16
    a_u, P_s = project_a_u_P_sigma(u_proj, sig16, phi, transpose=transpose)
    R = P_s / (RHO0 * NU * K_PHYS * a_u)
    s_prof = profile_s(sig16, a_u, transpose=transpose)
    return {
        "sigma16": sig16,
        "ux16": ux16,
        "uy16": uy16,
        "a_u": a_u,
        "P_sigma": P_s,
        "R": R,
        "s": s_prof,
        "finite": bool(np.isfinite(sig16).all() and np.isfinite(a_u) and np.isfinite(P_s)
                       and np.isfinite(R) and (s_prof is None or np.isfinite(s_prof).all())),
    }


def stress_from_composite_arrival(maps, state, phi, transpose=False):
    sig16 = sigma_composite_16(maps, state)
    # u from observe path (same restriction rule)
    obs = None
    # rebuild ux/uy on 16^2 from arrival populations (owned cells only)
    n = N_COARSE
    rho_c = np.full((n, n), np.nan)
    jx_c = np.full((n, n), np.nan)
    jy_c = np.full((n, n), np.nan)
    fc, ff = state.coarse_f, state.fine_f
    for y in range(n):
        for x in range(n):
            if maps.coarse_active[y, x]:
                sl = fc[:, y, x]
                rho_c[y, x] = sl.sum()
                jx_c[y, x] = (CX * sl).sum()
                jy_c[y, x] = (CY * sl).sum()
    for ky in range(PATCH_LO, PATCH_HI):
        for kx in range(PATCH_LO, PATCH_HI):
            y0 = spec_children_start(ky)
            x0 = spec_children_start(kx)
            kids = []
            for dy in (0, 1):
                for dx in (0, 1):
                    sl = ff[:, y0 + dy, x0 + dx]
                    kids.append((sl.sum(), (CX * sl).sum(), (CY * sl).sum()))
            rho_c[ky, kx] = np.mean([k[0] for k in kids])
            jx_c[ky, kx] = np.mean([k[1] for k in kids])
            jy_c[ky, kx] = np.mean([k[2] for k in kids])
    ux16 = jx_c / rho_c
    uy16 = jy_c / rho_c
    u_proj = uy16 if transpose else ux16
    a_u, P_s = project_a_u_P_sigma(u_proj, sig16, phi, transpose=transpose)
    R = P_s / (RHO0 * NU * K_PHYS * a_u)
    s_prof = profile_s(sig16, a_u, transpose=transpose)
    return {
        "sigma16": sig16,
        "ux16": ux16,
        "uy16": uy16,
        "a_u": a_u,
        "P_sigma": P_s,
        "R": R,
        "s": s_prof,
        "finite": bool(np.isfinite(sig16).all() and np.isfinite(a_u) and np.isfinite(P_s)
                       and np.isfinite(R) and (s_prof is None or np.isfinite(s_prof).all())),
    }


def main():
    t0 = time.time()
    print("=== local_stress_observable_v1 against CPUREF-0.3 ===")
    print(f"REF_PATH={REF_PATH}")
    ref_sha = sha256_file(REF_PATH)
    print(f"REF_SHA256={ref_sha}")
    expected_ref = "bba2a06a9f6517b32079ffff15e4060d8c3b7ed35f16ca134e1877226c9e4e1e"
    print(f"REF_SHA_MATCH={ref_sha == expected_ref}")
    ref = _load_ref()
    print(f"API_VERSION={getattr(ref, 'API_VERSION', '?')}")
    assert getattr(ref, "API_VERSION", None) == "CPUREF-0.3"
    print(f"U0={U0} tau_c={TAU_C} tau_f={TAU_F} steps={N_STEPS} nu={NU} k={K_PHYS}")
    print(f"nu*k^2={NU * K_PHYS**2:.8f}")
    print("ARRIVAL_STATE_POLICY: read CPUREF-0.3 State.coarse_f/fine_f at phase='arrival_f'")
    print("  after N advance_coarse_once calls (tick=2N). That array IS post-stream pre-collision.")
    print("  Uniform B/C use local collide-then-stream; f after 2N (B) or N (C) streams is arrival.")
    print("  If arrival were unreachable without changing the reference, this script would STOP.")
    print("STRESS: Pi_neq=sum(cx*cy*f)-rho*ux*uy; sigma=-(1-1/(2*tau))*Pi_neq; 2x2 vol-avg of sigma.")
    print("  No extra dx/dt factor. a_u=(2/256)sum u sin; P_sigma=(2/256)sum sigma cos; own a_u.")
    print("  Profile s(j)=mean_x[4,12)(sigma)/ (rho0*nu*k*a_u). Transposed: R only, no profile.")
    print()

    phis = [0.0, 0.5 * np.pi]
    results = {}
    fields = {}
    init_mm = {}
    arrivals = {}
    stress = {}

    for phi in phis:
        tag = f"phi{phi:.4f}"
        print(f"-- running C (uniform-coarse) {tag}")
        ser_c, fld_c, f_c, tau_c, an_c = run_uniform("C", phi)
        results[f"C_{tag}"] = ser_c
        fields[f"C_{tag}"] = fld_c
        arrivals[f"C_{tag}"] = an_c
        stress[f"C_{tag}"] = stress_from_uniform_arrival(f_c, tau_c, phi)
        print(f"  ARRIVAL_ARRAY: {an_c}")

        print(f"-- running B (uniform-fine) {tag}")
        ser_b, fld_b, f_b, tau_b, an_b = run_uniform("B", phi)
        results[f"B_{tag}"] = ser_b
        fields[f"B_{tag}"] = fld_b
        arrivals[f"B_{tag}"] = an_b
        stress[f"B_{tag}"] = stress_from_uniform_arrival(f_b, tau_b, phi)
        print(f"  ARRIVAL_ARRAY: {an_b}")

        print(f"-- running A (composite CPUREF) {tag}")
        ser_a, fld_a, st_a, maps_a, mism, an_a = run_composite(ref, phi)
        init_mm[tag] = mism
        results[f"A_{tag}"] = ser_a
        fields[f"A_{tag}"] = fld_a
        arrivals[f"A_{tag}"] = an_a
        stress[f"A_{tag}"] = stress_from_composite_arrival(maps_a, st_a, phi)
        print(f"  ARRIVAL_ARRAY: {an_a}")
        print(f"  state.phase={st_a.phase} state.tick={st_a.tick}")

        patch = np.zeros((N_COARSE, N_COARSE), dtype=bool)
        patch[PATCH_LO:PATCH_HI, PATCH_LO:PATCH_HI] = True
        out = ~patch
        for t in range(N_STEPS + 1):
            du = fld_a[t][0] - fld_b[t][0]
            ser_a.e_if[t] = float(np.linalg.norm(du))
            ser_a.e_if_patch[t] = float(np.linalg.norm(du[patch]))
            ser_a.e_if_out[t] = float(np.linalg.norm(du[out]))
            ser_a.e_cb[t] = field_l2(fld_c[t][0], fld_b[t][0])

    phi_t = 0.5 * np.pi
    print(f"-- running A transpose phi={phi_t}")
    ser_at, fld_at, st_at, maps_at, mism_t, an_at = run_composite(ref, phi_t, transpose=True)
    init_mm["A_T"] = mism_t
    results["A_T_phi1.5708"] = ser_at
    fields["A_T_phi1.5708"] = fld_at
    arrivals["A_T_phi1.5708"] = an_at
    stress["A_T_phi1.5708"] = stress_from_composite_arrival(
        maps_at, st_at, phi_t, transpose=True
    )
    print(f"  ARRIVAL_ARRAY: {an_at}")

    ser_a_ref = results["A_phi1.5708"]
    fld_a_ref = fields["A_phi1.5708"]
    trans_err = 0.0
    for t in range(N_STEPS + 1):
        ux, uy, rho, jx, jy = fld_a_ref[t]
        ux_t, uy_t, rho_t, jx_t, jy_t = fld_at[t]
        err = max(
            float(np.max(np.abs(ux_t - uy.T))),
            float(np.max(np.abs(uy_t - ux.T))),
            float(np.max(np.abs(rho_t - rho.T))),
        )
        trans_err = max(trans_err, err)
    transpose_ok = trans_err <= TRANS_TOL

    print("\n=== HARD CHECKS (unchanged v0.2 set) ===")
    hard = {}
    hard_keys = []
    for phi in phis:
        tag = f"phi{phi:.4f}"
        hard_keys.extend([f"A_{tag}", f"B_{tag}", f"C_{tag}"])
    hard_keys.append("A_T_phi1.5708")
    for key in hard_keys:
        s = results[key]
        hard[f"finite_{key}"] = s.finite_ok
        hard[f"cons_{key}"] = s.cons_ok
        tol_abs = CONS_TOL * max(abs(s.mass[0]), 1.0)
        print(
            f"  finite[{key}]={s.finite_ok} cons={s.cons_ok} "
            f"mass_drift={s.mass_drift:.3e} momx_drift={s.momx_drift:.3e} "
            f"momy_drift={s.momy_drift:.3e} tol_abs={tol_abs:.3e} {s.notes}"
        )
    hard["transpose_ok"] = transpose_ok
    print(f"  transpose_ok={transpose_ok} max_err={trans_err:.3e} (tol={TRANS_TOL:.3e})")

    for phi in phis:
        tag = f"phi{phi:.4f}"
        ux_a = fields[f"A_{tag}"][0][0]
        ux_b = fields[f"B_{tag}"][0][0]
        ux_c = fields[f"C_{tag}"][0][0]
        patch = np.zeros((16, 16), dtype=bool)
        patch[PATCH_LO:PATCH_HI, PATCH_LO:PATCH_HI] = True
        d_ab_p = float(np.max(np.abs(ux_a[patch] - ux_b[patch])))
        d_ac_o = float(np.max(np.abs(ux_a[~patch] - ux_c[~patch])))
        hard[f"t0_AB_patch_ok_{tag}"] = bool(d_ab_p <= T0_TOL)
        hard[f"t0_AC_out_ok_{tag}"] = bool(d_ac_o <= T0_TOL)
        hard[f"init_mismatch_ok_{tag}"] = bool(init_mm[tag] <= INIT_TOL)
        print(
            f"  t0 {tag}: AB_patch={d_ab_p:.3e} ok={hard[f't0_AB_patch_ok_{tag}']} "
            f"AC_out={d_ac_o:.3e} ok={hard[f't0_AC_out_ok_{tag}']} "
            f"init_mm={init_mm[tag]:.3e} ok={hard[f'init_mismatch_ok_{tag}']}"
        )
    hard["init_mismatch_ok_A_T_phi1.5708"] = bool(init_mm["A_T"] <= INIT_TOL)
    print(f"  init_mismatch[A_T]={init_mm['A_T']:.3e} ok={hard['init_mismatch_ok_A_T_phi1.5708']}")

    # New finite checks for stress arrays
    print("\n=== STRESS FINITE CHECKS (addition) ===")
    for key, stv in stress.items():
        ok = stv["finite"]
        hard[f"stress_finite_{key}"] = ok
        print(
            f"  stress_finite[{key}]={ok} R={stv['R']:.8f} a_u={stv['a_u']:.6e} "
            f"P_sigma={stv['P_sigma']:.6e}"
        )

    _bool_keys = sorted(k for k, v in hard.items() if isinstance(v, bool))
    # Report both: original-22-style (without stress_finite) and all including stress
    v02_bool = [k for k in _bool_keys if not k.startswith("stress_finite_")]
    print(f"  N_BOOLEAN_HARD_CHECKS_V02={len(v02_bool)}: " + ",".join(v02_bool))
    print(f"  N_BOOLEAN_HARD_CHECKS_WITH_STRESS={len(_bool_keys)}")

    print("\n=== LOCAL STRESS OBSERVABLE at coarse step 32 ===")
    print("FORMULA: R = P_sigma / (rho0 * nu * k * a_u); a_u from SAME run/tick/arrival")
    print("WORDING: reported diagnostics only; phi probes phase structure only; no mechanism claim")
    table_rows = []
    for phi in phis:
        tag = f"phi{phi:.4f}"
        sa, sb, sc = stress[f"A_{tag}"], stress[f"B_{tag}"], stress[f"C_{tag}"]
        ra, rb, rc = sa["R"], sb["R"], sc["R"]
        print(f"\n-- {tag} --")
        print(f"  R_A={ra:.10f}  R_B={rb:.10f}  R_C={rc:.10f}")
        print(f"  R_A/R_B={ra/rb:.10f}  R_A/R_C={ra/rc:.10f}")
        print(f"  a_u A/B/C: {sa['a_u']:.8e} {sb['a_u']:.8e} {sc['a_u']:.8e}")
        print(f"  P_sigma A/B/C: {sa['P_sigma']:.8e} {sb['P_sigma']:.8e} {sc['P_sigma']:.8e}")
        # profiles
        sA, sB, sC = sa["s"], sb["s"], sc["s"]
        dAB = sA - sB
        dAC = sA - sC
        face = {3, 4, 11, 12}
        print("  s_A(j): " + " ".join(f"{v:.6e}" for v in sA))
        print("  s_B(j): " + " ".join(f"{v:.6e}" for v in sB))
        print("  s_C(j): " + " ".join(f"{v:.6e}" for v in sC))
        print("  s_A-s_B: " + " ".join(f"{v:.6e}" for v in dAB))
        print("  s_A-s_C: " + " ".join(f"{v:.6e}" for v in dAC))
        print(
            "  FACE_ROWS j=3,4,11,12 flagged: "
            + " ".join(
                f"j{j}:A={sA[j]:.5e},B={sB[j]:.5e},C={sC[j]:.5e},"
                f"A-B={dAB[j]:.5e},A-C={dAC[j]:.5e}"
                for j in (3, 4, 11, 12)
            )
        )
        table_rows.append(
            {
                "phi": phi,
                "tag": tag,
                "R_A": ra,
                "R_B": rb,
                "R_C": rc,
                "R_A/R_B": ra / rb,
                "R_A/R_C": ra / rc,
                "s_A": sA,
                "s_B": sB,
                "s_C": sC,
                "dAB": dAB,
                "dAC": dAC,
            }
        )

    stT = stress["A_T_phi1.5708"]
    print(f"\n-- A_T phi=1.5708 (R only, no profile) --")
    print(f"  R_A_T={stT['R']:.10f}  a_u={stT['a_u']:.8e}  P_sigma={stT['P_sigma']:.8e}")
    print(f"  profile={'None (transpose)'}")

    print("\n=== DECAY FITS (informational; same as v0.2) ===")
    lam_analytic = NU * K_PHYS**2
    print(f"  analytic lambda = {lam_analytic:.8f}")
    for phi in phis:
        tag = f"phi{phi:.4f}"
        for kind in ("A", "B", "C"):
            s = results[f"{kind}_{tag}"]
            lam, se, r2 = fit_lambda(s.A)
            print(
                f"  {kind} {tag}: lambda={lam:.8f} stderr={se:.3e} R2={r2:.6f} "
                f"A0={s.A[0]:.6e} A32={s.A[-1]:.6e}"
            )

    all_hard_pass = all(v for k, v in hard.items() if isinstance(v, bool))
    v02_pass = all(hard[k] for k in v02_bool)
    print("\n=== NUMERIC SUMMARY ===")
    print(f"ALL_HARD_CHECKS_PASS_V02={v02_pass}")
    print(f"ALL_HARD_CHECKS_PASS_WITH_STRESS={all_hard_pass}")
    print(f"runtime_sec={time.time()-t0:.2f}")

    print("SUMMARY_LINES_BEGIN")
    for row in table_rows:
        print(
            f"R {row['tag']} A={row['R_A']:.12f} B={row['R_B']:.12f} C={row['R_C']:.12f} "
            f"A/B={row['R_A/R_B']:.12f} A/C={row['R_A/R_C']:.12f}"
        )
        print(
            f"PROFILE_A {row['tag']} " + " ".join(f"{v:.8e}" for v in row["s_A"])
        )
        print(
            f"PROFILE_B {row['tag']} " + " ".join(f"{v:.8e}" for v in row["s_B"])
        )
        print(
            f"PROFILE_C {row['tag']} " + " ".join(f"{v:.8e}" for v in row["s_C"])
        )
        print(
            f"PROFILE_A-B {row['tag']} " + " ".join(f"{v:.8e}" for v in row["dAB"])
        )
        print(
            f"PROFILE_A-C {row['tag']} " + " ".join(f"{v:.8e}" for v in row["dAC"])
        )
    print(f"R A_T_phi1.5708 R={stT['R']:.12f}")
    print(f"TRANSPOSE err={trans_err:.6e} ok={transpose_ok}")
    print(f"REF_SHA256={ref_sha}")
    print("SUMMARY_LINES_END")
    return 0 if all_hard_pass else 1


if __name__ == "__main__":
    sys.exit(main())