#!/usr/bin/env python
r"""BNS parameter estimation benchmark: Cosmic Explorer injection, FD relative
binning, and the jaxpe JAX HMC sampler.

Setup
-----
Injection: m1 = m2 = 1.4 Msun (chirp mass ~ 1.2188 Msun, eta = 0.25), zero spins,
f_lower = 10 Hz, sampling rate 4096 Hz, zero-noise data in a single detector with the
Cosmic Explorer P1600143 PSD (H1 site geometry as the CE stand-in). The segment is
2048 s, comfortably longer than the ~1000 s inspiral from 10 Hz, so the FD waveform is
wrap-free and df = 1/2048 Hz resolves the signal's frequency structure.

Priors (all uniform): chirp mass in [0.9, 1.1] x true, eta in [0.2, 0.25],
spin1z and spin2z in [--spin-min, --spin-max] (default [0, 0.05], matching the
zero-spin BNS reference above; ``--spin1z``/``--spin2z`` set a non-zero injected
truth and ``--spin-min``/``--spin-max`` widen the prior to a symmetric aligned-spin
range, e.g. [-0.9, 0.9] for BBH-mass sources -- see run_mass_sweep_pe.py for the
mass-dependent NS/BH convention used across a sweep). Extrinsic parameters are
fixed at the injected values.

Likelihood: :class:`~jaxpe.gw.likelihood.RelativeBinningFDLikelihood` summary data
built once on the dense grid (CPU), then a lean jitted log-posterior that evaluates
IMRPhenomD only at the ~n_bins bin edges. The lean path is asserted equal to the class
implementation at machine precision, and the heterodyne is validated against the dense
:class:`~jaxpe.gw.likelihood.FDNetworkLikelihood` on draws spanning the posterior bulk.

Sampler: jaxpe's HMC kernel with a *dense* mass matrix (the Laplace covariance at the
unconstrained-space MAP, found by damped Newton from the fiducial, with the soft
eigenvalues floored at 1 for the boundary-tail directions) and long leapfrog
trajectories that bend along the curved chirp-mass/eta/spin degeneracy valley. What
matters is the integration time T = eps * n_leapfrog, and eps does NOT transfer
between trajectory lengths, so warmup adapts eps at the production n_leapfrog and
averages log(eps) over its post-transient blocks (a single final iterate oscillates
enough to swing the run 7-14 min). Chains stranded on secondary ripples of the
oscillatory matched-filter likelihood (Delta lnL ~ -10^3) are re-seeded.

An equilibration phase then runs discarded flow rounds: each spreads the chains and
is refit on that spread, which both bootstraps the flow out of the poor fit warmup
alone provides and starts production at stationarity (a burn-in transient inside the
kept series is indistinguishable from non-convergence to Rhat). Production then
interleaves local HMC with flow global independence proposals (``jaxpe.flows`` +
``jaxpe.sampler``'s global block), which teleport chains along the boundary-piled
eta -> 1/4 and spin -> 0 tails that no fixed mass matrix equilibrates.

Nothing here is fitted to a particular source: masses are ``--mass1/--mass2``, the
priors and the optimiser start are derived from them, and every adaptation is driven
by measured acceptance. Verified by rerunning at 1.35 + 1.25 Msun with no retuning.

``--kernel`` selects the local transition kernel from ``jaxpe.kernels``: ``hmc``
(default, the only one with validated numbers on docs/bns_ce_pe_benchmark.md),
``mala``, ``mmala`` (constant-metric mode -- no per-point Fisher/metric estimator
exists in this pipeline, so this is NOT the full Riemannian variant, just dense
MALA under another name), ``random-walk``, or ``uld``. The four non-HMC kernels
reuse the same MAP-Laplace mass matrix and the same flow-based equilibration, but
adapt their step size to jaxpe's own literature-default target acceptance
(``jaxpe.kernels.adaptation.TARGET_ACCEPTANCE``, overridable via
``--target-acceptance``) rather than a target measured on this posterior, and skip
the HMC-specific trajectory-length/eps-cycling machinery entirely (there is no
trajectory length or leapfrog resonance to manage outside HMC). ``uld`` is
unadjusted -- no Metropolis-Hastings step, so ``--step-size``/``--friction`` are
held fixed for the whole run (nothing to adapt an acceptance rate toward) and the
kept posterior carries an uncorrected O(step_size^2) discretization bias by
construction (see ``jaxpe/kernels/uld.py``); it is exposed for a fast/approximate
look, not as a substitute for HMC/MALA's exact posterior.

Convergence gate, evaluated per block: rank-normalized split-Rhat over the
*global-subseries* (near-independent draws; split-Rhat over the raw autocorrelated
series only re-measures tau) < 1.01, Geyer min ESS over the full series >= target,
and no stuck chains. All design decisions above were measured, not assumed -- see
docs/bns_ce_pe_benchmark.md for the experiment ledger.

The heavy setup runs on CPU regardless of the default device; the sampling hot loop
runs on the default device (GPU when available) touching only O(n_bins) constants.

Run:  python bin/run_bns_ce_pe.py                        (default, ~7-10 min on a T2000)
      python bin/run_bns_ce_pe.py --reference           (the 15.4-min round-1 config)
      python bin/run_bns_ce_pe.py --mass1 1.35 --mass2 1.25   (a different source)
      python bin/run_bns_ce_pe.py --quick               (reduced CPU validation config)
      python bin/run_bns_ce_pe.py --setup-only          (profile setup + RB validation)
      python bin/run_bns_ce_pe.py --target-snr 20       (rescale distance to hit SNR 20)
      python bin/run_bns_ce_pe.py --kernel mala         (a different local kernel)

The first invocation pays JIT compilation; the persistent XLA cache makes every
later one compile-free, which is the honest basis for the quoted timings.
"""

import argparse
import json
import os
import time
from pathlib import Path

os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.15")

import jax

jax.config.update("jax_enable_x64", True)
# persistent XLA compilation cache: repeat invocations skip all jit compiles
# (the 20-minute benchmark budget excludes compile time; a warm second run is
# the honest measurement of it)
# JAXPE_XLA_CACHE_DIR relocates it: the cache grows to several GB, and the home
# filesystem is not always where there is room for it.
jax.config.update(
    "jax_compilation_cache_dir",
    os.environ.get("JAXPE_XLA_CACHE_DIR", os.path.expanduser("~/.cache/jaxpe_xla")),
)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 1.0)

import jax.numpy as jnp
import numpy as np

from jaxpe.core.priors import JointPrior, Uniform
from jaxpe.core.problem import InferenceProblem
from jaxpe.drivers.relative_binning_pe import eta_to_q, map_laplace, run_pe
from jaxpe.gw import (
    IMRPhenomD,
    distance_for_target_snr,
    lalsim_psd,
    make_injection,
    network_snr,
)
from jaxpe.gw.detectors import EARTH_OMEGA
from jaxpe.gw.likelihood import RelativeBinningFDLikelihood
from jaxpe.gw.likelihood.base import project_to_detector

# --------------------------------------------------------------------------- setup


def build_loglike(rb, fixed, f32: bool = False):
    r"""Lean jitted-friendly log-likelihood over (chirp_mass, eta, spin1z, spin2z).

    Mathematically identical to ``rb.log_likelihood`` but closed over ONLY the
    O(n_bins) summary arrays (as numpy constants, baked into the jit on the sampling
    device), so the multi-million-point setup grids never reach the GPU.

    With ``f32``, the waveform and the per-bin heterodyne products are evaluated in
    single precision -- measured 3x faster here, because this GPU is a consumer part
    with 1/32 fp64 throughput and the cost is dominated by IMRPhenomD's ~3600-op
    scalar coefficient algebra. Two things make that safe rather than reckless:

    * **The coalescence-time phase is removed, not approximated.** ``geocent_time``
      is ~1.19e9 s, so ``exp(-2 pi i f t_c)`` reaches ~1e13 rad and is meaningless in
      f32. But t_c is *fixed*, not sampled, so its phase factor is identical in the
      trial and in the fiducial and cancels exactly in the ratio r = h / h_0. Both
      are therefore evaluated at ``geocent_time = 0``: the ratio is unchanged, and
      the summary data (built from the true h_0 against the data) are untouched.
    * **Only the products are single precision; the sums are not.** Each per-bin
      term is ~<d|d>/(2 n_bins), and the f64 accumulation keeps the large
      cancellation against ``half_dd`` exact, so the f32 error enters as ~1e-3 in
      lnL -- an order of magnitude below the relative-binning truncation error the
      run already accepts, and checked at runtime by ``validate_rb`` against the
      dense f64 likelihood.
    """
    st = rb._static()
    half_dd = float(st["rb_half_dd"])
    gmst = rb.gmst_ref + EARTH_OMEGA * (fixed["geocent_time"] - rb.t_ref)
    waveform = rb.waveform
    ra, dec, psi = fixed["ra"], fixed["dec"], fixed["psi"]
    rdt, cdt = (np.float32, np.complex64) if f32 else (np.float64, np.complex128)

    edge_freqs = np.asarray(st["rb_edge_freqs"], rdt)
    dfbin = np.asarray(st["rb_dfbin"], rdt)
    # tc=0 base point: see the docstring -- the fixed-tc phase cancels in the ratio
    base = {k: v for k, v in fixed.items()}
    base["geocent_time"] = 0.0 if f32 else fixed["geocent_time"]

    dets = []
    for det in rb.detectors:
        if f32:  # ratio denominator must use the SAME tc=0 convention as the trial
            fid = dict(base)
            hp0, hc0 = waveform(fid, jnp.asarray(st["rb_edge_freqs"]))
            h0 = np.asarray(
                project_to_detector(
                    det,
                    hp0,
                    hc0,
                    jnp.asarray(st["rb_edge_freqs"]),
                    ra,
                    dec,
                    psi,
                    gmst,
                )
            )
        else:
            h0 = np.asarray(st["rb_h0_edges"][det.name])
        dets.append(
            (
                det,
                h0.astype(cdt),
                np.asarray(st["rb_A0"][det.name], cdt),
                np.asarray(st["rb_A1"][det.name], cdt),
                np.asarray(st["rb_B0"][det.name], rdt),
                np.asarray(st["rb_B1"][det.name], rdt),
            )
        )

    f64 = jnp.float64

    def loglike(p):
        full = dict(base)
        full["chirp_mass"] = p["chirp_mass"]
        full["mass_ratio"] = eta_to_q(p["eta"])
        full["spin1z"] = p["spin1z"]
        full["spin2z"] = p["spin2z"]
        if f32:
            full = {k: jnp.asarray(v, jnp.float32) for k, v in full.items()}
        hp, hc = waveform(full, edge_freqs)
        lnl = -half_dd
        for det, h0, A0, A1, B0, B1 in dets:
            h = project_to_detector(det, hp, hc, edge_freqs, ra, dec, psi, gmst)
            r = h / h0
            r0 = 0.5 * (r[1:] + r[:-1])
            r1 = (r[1:] - r[:-1]) / dfbin
            # products in the working precision, reduction always in float64
            zdh = jnp.sum(
                (A0 * jnp.conj(r0) + A1 * jnp.conj(r1)).astype(jnp.complex128)
            )
            hh = jnp.sum(
                (
                    B0 * (r0.real**2 + r0.imag**2)
                    + 2.0 * B1 * jnp.real(r0 * jnp.conj(r1))
                ).astype(f64)
            )
            lnl = lnl + jnp.real(zdh) - 0.5 * hh
        return lnl

    return loglike


# ----------------------------------------------------------------------- validation
def validate_rb(rb, dense_like, loglike, prior, truth, x_true, sigma, rng, f32=False):
    """RB vs dense parity on draws spanning the posterior bulk and moderate tails.

    ``x_true`` is the truth in sampled-space order (chirp_mass, eta, spin1z, spin2z).
    Tolerance follows the Zackay error model (error ~ beta * |lnL|): require
    |RB - dense| < 0.1 in the bulk (|lnL| < 50) and < 5e-3 * |lnL| further out.
    Returns the worst (bulk_err, model_ratio) seen; raises on failure.
    """
    x_true = np.asarray(x_true, float)
    _dense_jit = jax.jit(dense_like.log_likelihood)  # 4M-point graph: jit pays off

    def dense_eval(params):  # jnp-array leaves so repeated calls do not retrace
        return float(_dense_jit({k: jnp.asarray(v) for k, v in params.items()}))

    lnl_rb_true = float(loglike(prior.as_dict(jnp.asarray(x_true))))
    lnl_dense_true = dense_eval(dict(truth))
    # both sides are O(<d|d>/2 ~ SNR^2/2) sums over ~1e5-1e6 points reduced in
    # different orders (numpy summary data vs fused XLA): exact-at-fiducial holds
    # to ~1e-8 relative, so the absolute tolerance must carry the <d|d>/2 scale
    tol_fid = 1e-8 * (1.0 + abs(rb._static()["rb_half_dd"]))
    if abs(lnl_rb_true - lnl_dense_true) > tol_fid:
        raise RuntimeError(
            f"exact-at-fiducial violated: RB {lnl_rb_true} vs dense "
            f"{lnl_dense_true} (tol {tol_fid:.2e})"
        )

    # Lean closure == class implementation. In f64 they must agree to machine
    # precision; in f32 the closure is deliberately a lower-precision evaluation of
    # the same expression, so the meaningful check is the parity-vs-dense loop
    # below (which bounds the TOTAL error, binning plus arithmetic) -- here we only
    # require agreement at the f32 level so a gross wiring error still trips.
    tol_lean = 3e-3 if f32 else 1e-6
    for _ in range(3):
        x = x_true + sigma * rng.standard_normal(x_true.size)
        x = np.clip(x, [p.low for p in prior.priors], [p.high for p in prior.priors])
        p = prior.as_dict(jnp.asarray(x))
        full = dict(truth)
        full.update(
            chirp_mass=x[0], mass_ratio=float(eta_to_q(x[1])), spin1z=x[2], spin2z=x[3]
        )
        a, b = float(loglike(p)), float(rb.log_likelihood(full))
        if abs(a - b) > tol_lean * (1.0 + abs(b)):
            raise RuntimeError(f"lean loglike != class loglike: {a} vs {b}")

    worst_bulk, worst_ratio = 0.0, 0.0
    for s in (0.5, 2.0):
        for _ in range(2):
            x = x_true + s * sigma * rng.standard_normal(x_true.size)
            x = np.clip(
                x, [p.low for p in prior.priors], [p.high for p in prior.priors]
            )
            full = dict(truth)
            full.update(
                chirp_mass=x[0],
                mass_ratio=float(eta_to_q(x[1])),
                spin1z=x[2],
                spin2z=x[3],
            )
            lnl_rb = float(loglike(prior.as_dict(jnp.asarray(x))))
            lnl_d = dense_eval(full)
            err = abs(lnl_rb - lnl_d)
            if abs(lnl_d) < 50.0:
                worst_bulk = max(worst_bulk, err)
            else:
                worst_ratio = max(worst_ratio, err / abs(lnl_d))
    if worst_bulk > 0.1 or worst_ratio > 5e-3:
        raise RuntimeError(
            f"relative-binning parity too loose: bulk {worst_bulk:.3g} (tol 0.1), "
            f"tail ratio {worst_ratio:.3g} (tol 5e-3); decrease --epsilon"
        )
    return worst_bulk, worst_ratio


# ----------------------------------------------------------------------------- main
def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--outdir", default="examples/output/bns_ce_rb_hmc")
    ap.add_argument("--duration", type=float, default=2048.0)
    ap.add_argument("--sampling-rate", type=float, default=4096.0)
    ap.add_argument("--f-min", type=float, default=10.0)
    ap.add_argument("--f-max", type=float, default=None, help="default 0.45*rate")
    ap.add_argument("--distance", type=float, default=200.0, help="Mpc")
    ap.add_argument(
        "--target-snr",
        type=float,
        default=None,
        help="rescale --distance so the injected network SNR equals this value "
        "(exact: SNR is proportional to 1/distance for a fixed source, so one "
        "rebuild at the solved distance suffices); default keeps --distance as given",
    )
    ap.add_argument("--mass1", type=float, default=1.4, help="component mass [Msun]")
    ap.add_argument("--mass2", type=float, default=1.4, help="component mass [Msun]")
    ap.add_argument(
        "--eta-min", type=float, default=0.2, help="lower edge of the eta prior"
    )
    ap.add_argument(
        "--spin-min", type=float, default=0.0, help="lower edge of the spin priors"
    )
    ap.add_argument(
        "--spin-max", type=float, default=0.05, help="upper edge of the spin priors"
    )
    ap.add_argument(
        "--spin1z",
        type=float,
        default=0.0,
        help="injected aligned-spin truth, component 1",
    )
    ap.add_argument(
        "--spin2z",
        type=float,
        default=0.0,
        help="injected aligned-spin truth, component 2",
    )
    ap.add_argument("--chi", type=float, default=1.0)
    ap.add_argument("--epsilon", type=float, default=0.25, help="RB phase per bin")
    ap.add_argument(
        "--max-rb-refinements",
        type=int,
        default=4,
        help="times to quarter --epsilon and retry when the RB-vs-dense parity "
        "guard fails, before giving up (the tolerances are never relaxed)",
    )
    ap.add_argument("--n-chains", type=int, default=64)
    ap.add_argument(
        "--kernel",
        choices=["hmc", "mala", "mmala", "uld", "random-walk"],
        default="hmc",
        help="local transition kernel (jaxpe.kernels). HMC is the only one with "
        "validated numbers on docs/bns_ce_pe_benchmark.md; the other four use "
        "jaxpe's library-default adaptation targets, unbenchmarked on this "
        "posterior. uld has no MH step -- see --friction and the printed note",
    )
    ap.add_argument(
        "--friction",
        type=float,
        default=1.0,
        help="uld only: BAOAB friction coefficient (held fixed, no adaptation)",
    )
    ap.add_argument(
        "--target-acceptance",
        type=float,
        default=None,
        help="override the Robbins-Monro target acceptance for non-hmc kernels "
        "(default: jaxpe.kernels.adaptation.TARGET_ACCEPTANCE's literature value "
        "per kernel); has no effect for hmc (uses this script's own measured "
        "0.75) or uld (no acceptance to target)",
    )
    # --n-leapfrog / --warmup-leapfrog are HMC-only (trajectory length); unused
    # by the other four kernels, which have no such concept.
    ap.add_argument("--n-leapfrog", type=int, default=32)
    ap.add_argument(
        "--warmup-leapfrog",
        type=int,
        default=48,
        help="shorter trajectories suffice to adapt eps and seed the flow",
    )
    ap.add_argument("--flow-epochs", type=int, default=25)
    ap.add_argument(
        "--f32",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="single-precision waveform + per-bin products (3x on this GPU)",
    )
    # 3, not 5. Under the eager-init overhead in run_chains this looked like an
    # 8 s difference ("within noise", so 5 was kept); with that overhead removed
    # the gap is 3.42 vs 4.01 min at the same seed -- cheaper equilibration
    # (34.7 vs 44.6 s) AND fewer production blocks (25 vs 31). A third stale
    # measurement taken in a regime a fixed cost dominated.
    ap.add_argument("--equil-rounds", type=int, default=3)
    ap.add_argument(
        "--ensemble-metric",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="replace the Laplace metric with the equilibrated ensemble covariance",
    )
    ap.add_argument("--retune-blocks", type=int, default=3)
    ap.add_argument(
        "--adapt-gain",
        type=float,
        default=2.0,
        help="Robbins-Monro gain for warmup step-size adaptation",
    )
    ap.add_argument("--flow-acc-target", type=float, default=0.65)
    ap.add_argument("--flow-interval", type=float, default=8.0)
    # The global block is ~3.4 s of an ~8.3 s production block, and only 0.6 s of
    # that is the likelihood -- the rest is the flow's two passes per step
    # (sample + log_prob). Measured warm, per 1200-step block: 8 layers/width 64
    # = 3.38 s, 4/64 = 2.00 s, 4/32 = 1.64 s, 2/64 = 1.31 s. 8 coupling layers is
    # generous for a 4-dim posterior, so this is exposed to trade capacity for
    # speed -- but a weaker flow means worse proposals, so it must be judged end
    # to end (block COUNT), never on per-block cost alone.
    # 4 is the measured knee: 8 -> 4 holds the block count (24 vs 25) and is
    # faster at both seeds, but 2 layers/width 32 degrades capacity and costs
    # 33 blocks against 24, so the cheaper-block/more-blocks trade returns below 4.
    ap.add_argument("--flow-layers", type=int, default=4)
    ap.add_argument("--flow-width", type=int, default=64)
    # > 0 enables a SECOND flow at this wider interval, cycled with the narrow one
    # in production. Measured motivation: a single wide flow reaches the boundary
    # tails and cuts blocks 25 -> 8, but collapses acceptance to ~0.1 and makes the
    # run a lottery (3.87 min at one seed, 10.71 min at another). Cycling keeps the
    # narrow kernel's acceptance and reproducibility while still reaching the tails.
    ap.add_argument("--flow-interval-wide", type=float, default=0.0)
    ap.add_argument(
        "--reference",
        action="store_true",
        help="reproduce the 15.4-minute reference run's sampler settings",
    )
    ap.add_argument("--step-size", type=float, default=0.5)
    ap.add_argument("--warmup-blocks", type=int, default=5)
    ap.add_argument("--warmup-steps", type=int, default=15)
    # 12, not 25. This LOST under the eager-init overhead in run_chains (5.27 min
    # vs 4.83), because a 2.25 s fixed cost per call meant halving the steps could
    # only ever save 0.65 s of a 6.6 s block -- the test was rigged against itself.
    # With the overhead removed, measured over three seeds (42/7/13):
    #   ps=25: 3.42 / 4.38 / 3.43 min  (25 / 38 / 25 blocks)  worst 4.38
    #   ps=12: 3.49 / 2.95 / 2.89 min  (25 / 22 / 21 blocks)  worst 3.49
    ap.add_argument("--production-steps", type=int, default=12)
    ap.add_argument("--thin", type=int, default=2)
    ap.add_argument("--n-global", type=int, default=1200)
    # Equilibration and production both spend flow proposals, but for different
    # reasons -- equilibration to TRAIN the flow and spread the chains, production
    # to accumulate near-independent draws for Rhat. Tying them to one knob makes
    # any sweep of the production count silently pay for it twice in setup, so the
    # equilibration count is separable (defaults to --n-global for continuity).
    ap.add_argument("--n-global-equil", type=int, default=None)
    # A safety stop only -- --max-minutes is the real budget guard. It was 40, which
    # a 1.35+1.25 Msun source hit at Rhat 1.0107 and so reported as NOT converged
    # despite being ~2 blocks short; a cap that turns "needs a bit longer" into
    # "failed" is measuring the cap, not the sampler.
    # A CAP, not a budget: --max-minutes is the real guard (see the note above about
    # 40 turning "needs slightly longer" into "failed"). 80 is still low enough to
    # bind on single-step kernels -- measured, MALA hits it at R-hat 1.012 having
    # spent only 6.4 of its 12 allowed minutes -- which makes a cross-kernel
    # comparison at a fixed cap measure the CAP, penalising exactly the kernels that
    # take smaller steps per block. Raised so wall clock is the binding constraint.
    ap.add_argument("--max-production-blocks", type=int, default=400)
    ap.add_argument(
        "--max-saved-samples",
        type=int,
        default=400_000,
        help="uniform stride the kept series down to at most this many rows before "
        "writing samples.npz; diagnostics are always computed on the full chains",
    )
    ap.add_argument("--rhat-target", type=float, default=1.01)
    ap.add_argument("--ess-target", type=float, default=2000.0)
    ap.add_argument("--max-minutes", type=float, default=20.0)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--quick", action="store_true", help="small CPU smoke test")
    ap.add_argument(
        "--setup-only", action="store_true", help="stop after setup + RB validation"
    )
    ap.add_argument(
        "--setup-cache",
        default=None,
        help="npz path: reuse (or, if absent, write) this injection's solved "
        "distance, refined --epsilon and MAP+Laplace mode/covariance. All three "
        "depend only on the injection, never on --kernel, so sharing them across "
        "kernels is both cheaper and strictly more like-for-like",
    )
    args = ap.parse_args()
    if args.reference:  # the 15.4-minute configuration, for like-for-like reruns
        args.epsilon, args.n_chains, args.n_leapfrog = 0.1, 256, 128
        args.n_global, args.flow_epochs = 300, 40
        args.warmup_blocks, args.warmup_steps = 5, 25
        args.adapt_gain = 1.0
        args.equil_rounds, args.retune_blocks, args.f32 = 0, 0, False
        args.ensemble_metric, args.n_leapfrog, args.step_size = False, 128, 0.1
    if not args.warmup_leapfrog:
        args.warmup_leapfrog = args.n_leapfrog
    if args.quick:
        args.duration, args.sampling_rate, args.f_min = 128.0, 2048.0, 25.0
        args.n_chains, args.ess_target = 64, 500.0
        args.warmup_blocks, args.max_production_blocks = 5, 40
        args.production_steps = 50

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    timings: dict = {}
    t_start = time.perf_counter()
    print(f"jax {jax.__version__}, default backend: {jax.default_backend()}")

    m1, m2 = max(args.mass1, args.mass2), min(args.mass1, args.mass2)
    mc_true = (m1 * m2) ** 0.6 / (m1 + m2) ** 0.2
    eta_true = m1 * m2 / (m1 + m2) ** 2
    if not (args.spin_min <= args.spin1z <= args.spin_max):
        raise ValueError(
            f"--spin1z={args.spin1z} outside the ({args.spin_min}, {args.spin_max}) "
            "prior; widen --spin-min/--spin-max"
        )
    if not (args.spin_min <= args.spin2z <= args.spin_max):
        raise ValueError(
            f"--spin2z={args.spin2z} outside the ({args.spin_min}, {args.spin_max}) "
            "prior; widen --spin-min/--spin-max"
        )
    truth = dict(
        chirp_mass=mc_true,
        mass_ratio=m2 / m1,
        spin1z=args.spin1z,
        spin2z=args.spin2z,
        luminosity_distance=args.distance,
        geocent_time=1187008882.43,
        phase=1.3,
        inclination=0.4,
        ra=3.446,
        dec=-0.408,
        psi=0.8,
    )

    # Everything the setup phase derives depends only on the INJECTION -- the
    # rescaled distance, the refined relative-binning resolution, and the
    # MAP+Laplace mode and covariance -- never on which transition kernel will
    # sample it. Caching them makes a multi-kernel comparison on one injection
    # both cheaper (the MAP and the parity refinement are minutes at high bin
    # count) and strictly more like-for-like: every kernel then samples a
    # bit-identical likelihood from a bit-identical mass matrix, so a difference
    # between two runs is a difference between two kernels and nothing else.
    cache_path = Path(args.setup_cache) if args.setup_cache else None
    cached = None
    if cache_path is not None and cache_path.exists():
        with np.load(cache_path) as c:
            cached = {k: c[k] for k in c.files}
        print(f"setup cache: reusing {cache_path}")

    # ---- heavy setup pinned to CPU: dense grids never touch the GPU ----
    cpu = jax.devices("cpu")[0]
    with jax.default_device(cpu):
        t0 = time.perf_counter()
        n = int(args.duration * args.sampling_rate)
        freqs = np.fft.rfftfreq(n, d=1.0 / args.sampling_rate)
        psd = lalsim_psd("CE", freqs)
        timings["psd"] = time.perf_counter() - t0

        # On a cache hit the solved distance is already known, so the injection is
        # built once at that distance instead of twice (once to measure the SNR,
        # once at the rescaled distance).
        if cached is not None:
            args.distance = float(cached["distance"])
            truth["luminosity_distance"] = args.distance

        # Hoisted so the rescale rebuild below is conditioned identically to this one
        # by construction rather than by the two argument lists agreeing.
        inj_kwargs = dict(
            detector_names=("H1",),  # H1 site geometry as the single-CE stand-in
            duration=args.duration,
            sampling_rate=args.sampling_rate,
            f_min=args.f_min,
            f_max=args.f_max,
            psd_fn=lambda f: np.interp(f, freqs, psd),
            noise_seed=None,  # zero-noise injection: lnL peaks at exactly 0 at truth
        )

        t0 = time.perf_counter()
        dense_like = make_injection(IMRPhenomD(f_ref=args.f_min), truth, **inj_kwargs)
        snr = dense_like.optimal_snr(truth)
        net_snr = network_snr(dense_like, truth)
        timings["injection"] = time.perf_counter() - t0
        print(
            f"injection: Mc={mc_true:.5f} Msun, D={args.distance:.0f} Mpc, "
            f"SNR {snr} (network {net_snr:.1f})  [{timings['injection']:.1f}s]"
        )

        if args.target_snr is not None and cached is None:
            # h(f) scales as 1/D_L for a fixed source and orientation (the only
            # distance dependence in the detector response), so SNR = sqrt(<h|h>)
            # does too -- this rescale is EXACT, not an iterative or approximate
            # search, and one rebuild at the solved distance suffices.
            t0 = time.perf_counter()
            args.distance = distance_for_target_snr(dense_like, truth, args.target_snr)
            truth["luminosity_distance"] = args.distance
            dense_like = make_injection(
                IMRPhenomD(f_ref=args.f_min), truth, **inj_kwargs
            )
            snr = dense_like.optimal_snr(truth)
            net_snr = network_snr(dense_like, truth)
            timings["snr_rescale"] = time.perf_counter() - t0
            print(
                f"rescaled to target SNR {args.target_snr:.1f}: distance -> "
                f"{args.distance:.1f} Mpc, achieved network SNR {net_snr:.2f}  "
                f"[{timings['snr_rescale']:.1f}s]"
            )

        prior = JointPrior(
            {
                "chirp_mass": Uniform(0.9 * mc_true, 1.1 * mc_true),
                "eta": Uniform(args.eta_min, 0.25),
                "spin1z": Uniform(args.spin_min, args.spin_max),
                "spin2z": Uniform(args.spin_min, args.spin_max),
            }
        )
        if not (args.eta_min < eta_true <= 0.25):
            raise ValueError(
                f"eta_true={eta_true:.4f} outside the eta prior "
                f"({args.eta_min}, 0.25]; widen --eta-min"
            )

        def build_rb(epsilon):
            rb = RelativeBinningFDLikelihood.from_likelihood(
                dense_like, truth, chi=args.chi, epsilon=epsilon
            )
            loglike = build_loglike(rb, truth, f32=args.f32)
            return rb, loglike, InferenceProblem(prior=prior, log_likelihood=loglike)

        t0 = time.perf_counter()
        # A cache hit skips straight to the resolution the parity guard already
        # settled on for this injection; rebuilding the summary data there is
        # seconds, while REDERIVING it costs the whole refinement ladder.
        epsilon = float(cached["epsilon"]) if cached is not None else args.epsilon
        rb, loglike, problem = build_rb(epsilon)
        n_bins = rb.n_bins
        timings["rb_setup"] = time.perf_counter() - t0
        n_band = int(np.sum((freqs >= args.f_min) & (freqs <= rb.f_max)))
        print(
            f"relative binning: {n_band} band points -> {n_bins} bins "
            f"[{timings['rb_setup']:.1f}s]"
        )

        # Optimizer start: the fiducial (trigger) point itself, inset off any prior
        # edge it happens to lie on -- the sigmoid bijection sends the open bounds
        # to +-inf, so a start exactly on a boundary is not representable. Derived
        # from the prior support and the injection, with no numbers specific to a
        # particular binary: an equal-mass system starts inset from eta = 1/4, an
        # unequal-mass one starts at its own (interior) eta.
        #
        # The inset FRACTION is not fixed: the optimiser is run from a ladder of
        # them and the best converged mode is kept (see below). A single fixed
        # 0.02 is safe only where the likelihood is gentle across that offset,
        # which is a mass-dependent accident: for a zero-noise injection lnL peaks
        # at exactly 0 at the fiducial, and at 55 Msun the M_c-eta ridge is sharp
        # enough that insetting eta by 0.02 * 0.05 = 1e-3 costs 392 nats. Newton
        # then starts 392 below the peak, correctly climbs monotonically, and
        # still ends at a prior CORNER with a degenerate (sigma ~ 1e-13) Laplace
        # covariance -- a silently unusable mass matrix.
        lo = np.array([p.low for p in prior.priors])
        hi = np.array([p.high for p in prior.priors])
        x_fid = np.array(
            [truth["chirp_mass"], eta_true, truth["spin1z"], truth["spin2z"]]
        )
        t0 = time.perf_counter()
        if cached is not None:
            y_map, cov0 = cached["y_map"], cached["cov0"]
            logp_map = float(cached["logp_map"])
            y_init = None
        else:
            # Run the optimiser from EVERY rung and keep the best converged mode --
            # not the rung whose STARTING log-posterior is highest. Those are
            # different things, and the difference is not academic: for the
            # zero-spin BNS the highest-starting rung (2e-4) climbs to a mode
            # pinned against the eta boundary with log-posterior -26.8 and a
            # Laplace sigma 100-2400x too narrow, while the 2e-2 rung reaches
            # -11.5 with a usable covariance. Newton only accepts improvements, so
            # each rung is a valid local ascent; which BASIN it lands in is what
            # the start selects, and only the final mode reveals that.
            #
            # Cost is n_rungs MAP solves, paid once per injection and shared across
            # kernels by --setup-cache, against a mass matrix that is otherwise
            # silently degenerate (see the boundary-corner failure at 55 Msun).
            best = None
            for frac in (0.02, 2e-3, 2e-4, 2e-5, 2e-6):
                x_c = np.clip(x_fid, lo + frac * (hi - lo), hi - frac * (hi - lo))
                y_c = np.asarray(prior.to_unconstrained(jnp.asarray(x_c)))
                if not np.all(np.isfinite(y_c)):
                    continue
                y_m, cov_m, lp_m = map_laplace(problem, y_c)
                # Reject a covariance the Laplace approximation has collapsed:
                # a mode on the prior edge has a vanishing Jacobian, and the
                # resulting metric is numerically a delta function.
                sig = np.sqrt(np.diag(cov_m))
                ok = np.all(np.isfinite(sig)) and np.all(sig > 1e-8)
                print(
                    f"  MAP from inset {frac:g}: log-posterior {lp_m:.2f}, "
                    f"min sigma_y {sig.min():.2e}{'' if ok else '  [rejected: degenerate]'}"
                )
                if ok and (best is None or lp_m > best[0]):
                    best = (lp_m, y_m, cov_m, y_c, frac)
            if best is None:
                raise RuntimeError(
                    "every optimizer start gave a degenerate Laplace covariance; "
                    "the mode is on a prior edge for all of them"
                )
            logp_map, y_map, cov0, y_init, inset_frac = best
            print(f"optimizer start: kept inset fraction {inset_frac:g}")
        timings["map_laplace"] = time.perf_counter() - t0
        x_map = np.asarray(prior.to_physical(jnp.asarray(y_map)))
        print(
            f"MAP (unconstrained-space mode): x = {np.array2string(x_map, precision=6)}, "
            f"log-posterior {logp_map:.2f}  [{timings['map_laplace']:.1f}s]"
        )

        t0 = time.perf_counter()
        rng = np.random.default_rng(args.seed)
        # sampled-space truth (eta from the actual component masses, not assumed 1/4)
        x_true = np.array([mc_true, eta_true, truth["spin1z"], truth["spin2z"]])

        def posterior_scale(y, cov):
            """physical-space sigma: sigma_y * |dx/dy| (bijections are elementwise)"""
            jac = np.asarray(jax.jacfwd(prior.to_physical)(jnp.asarray(y)))
            return np.sqrt(np.diag(cov)) * np.abs(np.diag(jac))

        sigma_phys = posterior_scale(y_map, cov0)

        # The bin scheme's required resolution is a property of the PRIOR VOLUME,
        # not of the source alone: the linear-in-f ratio model has to hold across
        # the parameters actually proposed, so a +-0.9 aligned-spin prior at
        # SNR ~20 needs far finer bins than the +-0.05 BNS prior this pipeline was
        # first validated against (measured: tail ratio 0.22-0.50 at epsilon 0.25,
        # 40-100x outside tolerance). Refining automatically holds the ACCURACY
        # contract fixed and lets the bin count -- i.e. the cost -- absorb the
        # difference, instead of sampling a likelihood that provably does not
        # reproduce the dense one. The tolerances themselves are never relaxed.
        # A cache hit re-uses a resolution that ALREADY passed this guard on this
        # injection, so the guard is not re-run: it is a property of (injection,
        # epsilon), both of which the cache pins.
        for attempt in range(0 if cached is not None else args.max_rb_refinements + 1):
            try:
                wb, wr = validate_rb(
                    rb,
                    dense_like,
                    loglike,
                    prior,
                    truth,
                    x_true,
                    sigma_phys,
                    rng,
                    args.f32,
                )
                break
            except RuntimeError as exc:
                if attempt == args.max_rb_refinements:
                    raise RuntimeError(
                        f"{exc}\n(already refined epsilon {args.epsilon:g} -> "
                        f"{epsilon:g} over {attempt} attempts; raise "
                        "--max-rb-refinements or widen the analysis assumptions)"
                    ) from exc
                epsilon /= 4.0
                rb, loglike, problem = build_rb(epsilon)
                print(
                    f"  RB parity failed ({exc.args[0].split(';')[0]}); refining "
                    f"epsilon -> {epsilon:g} ({rb.n_bins} bins) and retrying"
                )
        if rb.n_bins != n_bins:
            # The mass matrix must be the curvature of the likelihood actually
            # sampled, so the MAP is redone once at the FINAL resolution -- once,
            # not per refinement, since the coarse-epsilon sigma above is only
            # needed to set the scale of the validation draws.
            n_bins = rb.n_bins
            y_map, cov0, logp_map = map_laplace(problem, y_init)
            sigma_phys = posterior_scale(y_map, cov0)
            x_map = np.asarray(prior.to_physical(jnp.asarray(y_map)))
            wb, wr = validate_rb(
                rb, dense_like, loglike, prior, truth, x_true, sigma_phys, rng, args.f32
            )
            print(
                f"re-MAP at epsilon {epsilon:g} ({n_bins} bins): "
                f"x = {np.array2string(x_map, precision=6)}, "
                f"log-posterior {logp_map:.2f}"
            )
        timings["rb_validation"] = time.perf_counter() - t0
        timings["rb_epsilon"] = epsilon
        timings["rb_n_bins"] = n_bins
        if cached is not None:
            wb = wr = float("nan")  # not re-measured; see the cache note above
        elif cache_path is not None:
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            np.savez(
                cache_path,
                distance=args.distance,
                epsilon=epsilon,
                y_map=y_map,
                cov0=cov0,
                logp_map=logp_map,
                net_snr=net_snr,
                n_bins=n_bins,
                parity_bulk=wb,
                parity_tail=wr,
            )
            print(f"setup cache: wrote {cache_path}")
        print(
            f"RB parity vs dense: bulk err {wb:.2e} (tol 0.1), tail ratio {wr:.2e} "
            f"(tol 5e-3); sigma_phys ~ {np.array2string(sigma_phys, precision=2)} "
            f"[{timings['rb_validation']:.1f}s]"
        )

    if args.setup_only:
        timings["total"] = time.perf_counter() - t_start
        print(f"setup-only: done in {timings['total'] / 60.0:.2f} min")
        print(f"timings: { {k: round(v, 2) for k, v in timings.items()} }")
        return 0

    # ---- sampling on the default device (GPU when available) ----
    # the wall-clock budget covers the WHOLE run: hand production what remains
    budget_min = args.max_minutes
    args.max_minutes = max(2.0, budget_min - (time.perf_counter() - t_start) / 60.0)
    phys, lps, rhat, ess, converged, kernel = run_pe(
        problem, y_map, cov0, args, timings
    )
    timings["total"] = time.perf_counter() - t_start

    # ---- report ----
    names = list(prior.names)
    flat = phys.reshape(-1, phys.shape[-1])
    q05, q50, q95 = np.percentile(flat, [5, 50, 95], axis=0)
    print("\n===== results =====")
    print(f"converged: {converged}  (Rhat {rhat.max():.4f}, min ESS {ess.min():.0f})")
    for i, nme in enumerate(names):
        print(
            f"  {nme:>11s}: median {q50[i]:.6g}  90% CI [{q05[i]:.6g}, {q95[i]:.6g}]"
            f"  truth {x_true[i]:.6g}"
        )
    imax = int(np.argmax(lps))
    print(f"  max log-posterior sampled: {lps.reshape(-1)[imax]:.3f}")
    total_min = timings["total"] / 60.0
    print(f"wall time: {total_min:.2f} min (budget {budget_min:.0f} min)")
    print(f"timings: { {k: round(v, 2) for k, v in timings.items()} }")

    # Decimate before saving. R-hat, ESS and every posterior summary above are
    # computed on the FULL chains; what lands on disk only has to support plots and
    # sample-to-sample comparisons downstream. The kept series is ~1200 global draws
    # per block per chain, so a long run writes 200+ MB per injection -- across a
    # multi-kernel suite that is tens of GB of a 4-parameter posterior, for no
    # statistical gain (the Geyer ESS is O(10^4), so a few 10^5 stored draws are
    # already far past the point where storing more reduces Monte Carlo error).
    # A uniform stride preserves the distribution; it is not a burn-in cut.
    stride = max(1, flat.shape[0] // max(1, args.max_saved_samples))
    print(
        f"saving {flat[::stride].shape[0]:,} of {flat.shape[0]:,} samples "
        f"(stride {stride}); diagnostics above used all of them"
    )
    np.savez(
        outdir / "samples.npz",
        names=names,
        samples=flat[::stride],
        log_prob=lps.reshape(-1)[::stride],
        truth=x_true,
        rhat=rhat,
        ess=ess,
        snr=net_snr,
        save_stride=stride,
        n_samples_full=flat.shape[0],
    )
    with open(outdir / "timings.json", "w") as f:
        json.dump(
            {
                **{k: float(v) for k, v in timings.items()},
                "converged": converged,
                "backend": jax.default_backend(),
                "n_bins": int(n_bins),
                "n_chains": int(args.n_chains),
                "network_snr": net_snr,
                "step_size": float(kernel.step_size),
            },
            f,
            indent=2,
        )

    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        d = len(names)
        fig, axes = plt.subplots(d, d, figsize=(2.4 * d, 2.4 * d))
        sub = flat[:: max(1, flat.shape[0] // 20000)]
        for i in range(d):
            for j in range(d):
                ax = axes[i, j]
                if j > i:
                    ax.axis("off")
                    continue
                if i == j:
                    ax.hist(flat[:, i], bins=80, histtype="step", color="C0")
                    ax.axvline(x_true[i], color="k", ls="--", lw=1)
                else:
                    ax.plot(sub[:, j], sub[:, i], ",", color="C0", alpha=0.3)
                    ax.plot(x_true[j], x_true[i], "k+", ms=10)
                if i == d - 1:
                    ax.set_xlabel(names[j])
                if j == 0 and i > 0:
                    ax.set_ylabel(names[i])
        fig.suptitle(
            f"BNS/CE FD relative binning + HMC (SNR {net_snr:.0f}, "
            f"{'converged' if converged else 'NOT converged'}, {total_min:.1f} min)"
        )
        fig.tight_layout()
        fig.savefig(outdir / "corner.png", dpi=110)
        print(f"saved {outdir}/corner.png")
    except ImportError:
        pass

    return 0 if converged and total_min < budget_min else 1


if __name__ == "__main__":
    raise SystemExit(main())