#!/usr/bin/env python
"""
integrate_likelihood_extrinsic_jax
==================================

A JAX, automatic-differentiation-based driver for the RIFT ILE extrinsic
likelihood.  It mirrors the *structure* and the relevant CLI/output conventions
of ``integrate_likelihood_extrinsic_batchmode`` but replaces the Monte-Carlo
likelihood evaluation with the differentiable JAX likelihood in
``RIFT.likelihood.jax_ile`` (a reimplementation of the production
``DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop`` "fused" code branch,
validated to ~1e-13 against the numpy reference).

REUSED unchanged from production RIFT (deliberately not reinvented):
  * frame reading            -> lalsimutils.frame_data_to_non_herm_hoff
  * PSD handling             -> lalsimutils.get_psd_series_from_xmldoc / resample
  * waveform + inner products-> factored_likelihood.PrecomputeLikelihoodTerms
  * array packing / epoch    -> factored_likelihood.PackLikelihoodDataStructuresAsArrays

NEW (the point of this driver):
  * the extrinsic -> lnL map is pure-JAX and AD-compatible, so we can
    gradient-ascend the (narrow) likelihood peak, form a Fisher approximation,
    and run efficient (gradient-aware) importance sampling.
  * DISTANCE MARGINALIZATION is performed analytically inside the JAX
    likelihood when ``--distance-marginalization`` is passed (ILE-compatible
    flag, default OFF).  The bare factored likelihood diverges on amplitude
    slivers (template power -> 0, e.g. inclination -> pi); marginalizing
    distance against the volumetric prior regulates this and yields a smooth,
    bounded objective over the five angular/sky parameters.  Required by
    ``--mode nuts``.

Full ILE argument compatibility: every integrate_likelihood_extrinsic_batchmode
option is accepted so this can be dropped into a production command line.
Implemented options are used; the rest are silently ignored (and reported);
science-changing options we do not implement (calibration marginalization, ROM,
NR templates, supplementary-likelihood, zero-likelihood, maximize-only) cause a
hard failure rather than a misleading result.  --sim-xml / --sim-grid and
--n-events-to-analyze enable batch processing of multiple intrinsic templates.

Run modes (``--mode``):
  prior-mc    : brute-force importance sampling from the physical prior
                (robust, slow; the classic ILE estimator).
  laplace-is  : prior-seeded *adaptive* Gaussian importance sampling
                (efficient for narrow peaks).  [default]
  map         : gradient-ascend the angular peak + report Fisher (AD demo).

Output (matches ILE conventions):
  <output>_<event>_.dat :
     event_id m1 m2 s1x s1y s1z s2x s2y s2z  lnL sigma_lnL ntotal neff
     ``lnL`` = log marginal likelihood (evidence) over extrinsic parameters,
     ``sigma_lnL = sqrt(var)/Z`` (ILE's ``sqrt(var)/res``).  Masses in M_sun.
  <output>_<event>_samples.dat (--save-samples) : per-sample extrinsic params +
     loglikelihood with ILE-style column names.
"""

from __future__ import print_function

import sys
from optparse import OptionParser, OptionGroup

import numpy as np

import jax
jax.config.update("jax_enable_x64", True)

import lal
import lalsimulation as lalsim

import RIFT.lalsimutils as lalsimutils
from RIFT.likelihood.jax_ile import build_data_from_precompute
from RIFT.likelihood.jax_ile.wrapper import (
    JAXExtrinsicLikelihood, JAXDistanceMarginalizedLikelihood,
)

MSUN = lal.MSUN_SI
PC = lal.PC_SI

# Angular parameter order for the distance-marginalized (default) path.
ANG_NAMES = ("ra", "dec", "psi", "incl", "phiref")
# Full extrinsic order for the fixed-distance path.
FULL_NAMES = ("ra", "dec", "psi", "incl", "phiref", "distMpc")


# ---------------------------------------------------------------------------
# Drop-in ILE argument compatibility
# ---------------------------------------------------------------------------
# Every option of integrate_likelihood_extrinsic_batchmode is accepted so this
# driver can be substituted into an existing production command line.  Options
# we implement are defined explicitly in build_parser(); the rest are registered
# here as accepted-but-ignored (with the correct arity so parsing succeeds), and
# a small set that would silently change the *science* if ignored is failed on.

# Boolean (zero-argument) ILE options (action=store_true/false).
_ILE_BOOL_OPTS = {
    "--check-good-enough", "--zero-likelihood", "--random-event",
    "--soft-fail-event-range", "--fmin-template-correct-for-lmax",
    "--internal-use-gwpy", "--nr-lookup", "--nr-hybrid-use", "--rom-use-basis",
    "--rom-integrate-intrinsic", "--nr-perturbative-extraction",
    "--nr-perturbative-extraction-full", "--nr-use-provided-strain",
    "--no-memory", "--use-gwsignal", "--maximize-only", "--time-marginalization",
    "--resample-time-marginalization", "--distance-marginalization",
    "--calibration-fused-kernel", "--calibration-export-posterior",
    "--extrinsic-proposal-adapt", "--vectorized", "--gpu", "--force-gpu-only",
    "--force-xpy", "--save-samples", "--save-samples-process-params",
    "--internal-hard-fail-on-error", "--internal-soft-fail-on-cuda-error",
    "--internal-make-empty-file-on-error", "--internal-waveform-fd-no-condition",
    "--verbose", "--save-eccentricity", "--fairdraw-extrinsic-output",
    "--convergence-tests-on", "--no-adapt", "--force-adapt-all",
    "--force-reset-all", "--no-adapt-distance", "--no-adapt-after-first",
    "--adapt-adapt", "--adapt-log", "--d-prior-redshift",
    "--declination-cosine-sampler", "--inclination-cosine-sampler",
    "--internal-rotate-phase", "--internal-sky-network-coordinates",
    "--internal-sky-network-coordinates-raw", "--auto-logarithm-offset",
    "--pin-distance-to-sim", "--export-eos-index",
    "--export-marginal-distance-grid", "--adapt-intrinsic",
}
# ILE options taking repeated values (action=append).
_ILE_APPEND_OPTS = {"--channel-name", "--psd-file"}
# The full ILE option set (so anything not implemented is still accepted).
_ILE_ALL_OPTS = {
    "--adapt-adapt", "--adapt-floor-level", "--adapt-intrinsic", "--adapt-log",
    "--adapt-weight-exponent", "--amp-order", "--approximant",
    "--auto-logarithm-offset", "--cache-file", "--calibration-burn-in-neff",
    "--calibration-burn-in-nmax", "--calibration-dump-responsibilities",
    "--calibration-envelope-directory", "--calibration-export-posterior",
    "--calibration-fused-kernel", "--calibration-n-realizations",
    "--calibration-pilot-extrinsic", "--calibration-proposal-breadcrumb",
    "--calibration-spline-count", "--channel-name", "--check-good-enough",
    "--coinc-xml", "--convergence-tests-on", "--data-end-time",
    "--data-integration-window-half", "--data-start-time",
    "--declination-cosine-sampler", "--deff-lambda", "--distance-marginalization",
    "--distance-marginalization-lookup-table", "--distance-slice-skip-threshold",
    "--distance-slice-wing-neff", "--distance-slice-wing-nmax", "--d-max",
    "--d-min", "--d-prior", "--d-prior-redshift", "--eff-lambda", "--e-freq",
    "--event", "--event-time", "--export-distance-slices", "--export-eos-index",
    "--export-marginal-distance-grid", "--extrinsic-proposal-adapt",
    "--extrinsic-proposal-breadcrumb", "--extrinsic-proposal-output",
    "--fairdraw-extrinsic-output", "--fairdraw-extrinsic-output-n-max", "--fmax",
    "--fmin-ifo", "--fmin-template", "--fmin-template-correct-for-lmax",
    "--force-adapt-all", "--force-gpu-only", "--force-reset-all", "--force-xpy",
    "--gpu", "--inclination-cosine-sampler",
    "--internal-data-storage-window-half", "--internal-hard-fail-on-error",
    "--internal-make-empty-file-on-error", "--internal-precompute-ignore-threshold",
    "--internal-rotate-phase", "--internal-sky-network-coordinates",
    "--internal-sky-network-coordinates-raw", "--internal-soft-fail-on-cuda-error",
    "--internal-use-gwpy", "--internal-waveform-extra-kwargs",
    "--internal-waveform-extra-lalsuite-args", "--internal-waveform-fd-no-condition",
    "--internal-waveform-taper", "--interpolate-time", "--inv-spec-trunc-time",
    "--l-max", "--manual-logarithm-offset", "--mass1", "--mass2",
    "--maximize-only", "--n-chunk", "--n-distance-slice-core", "--n-eff",
    "--n-events-to-analyze", "--n-fairdraw-extrinsic-samples", "--n-max",
    "--n-distance-slice-wing", "--no-adapt", "--no-adapt-after-first",
    "--no-adapt-distance", "--no-memory", "--nr-hybrid-method", "--nr-hybrid-use",
    "--nr-index", "--nr-lookup", "--nr-lookup-group", "--nr-params",
    "--nr-perturbative-extraction", "--nr-perturbative-extraction-full",
    "--nr-use-provided-strain", "--output-file", "--output-format",
    "--parameter", "--parameter-range", "--pin-distance-to-sim", "--psd-file",
    "--psd-window-shape", "--random-event", "--reference-freq",
    "--resample-time-marginalization", "--restricted-mode-list-file",
    "--rom-group", "--rom-integrate-intrinsic", "--rom-limit-basis-size-to",
    "--rom-param", "--rom-use-basis", "--sampler-method", "--sampler-portfolio",
    "--sampler-portfolio-args", "--sampler-xpy", "--save-eccentricity",
    "--save-samples", "--save-samples-process-params", "--seed", "--sim-grid",
    "--sim-xml", "--skymap-file", "--soft-fail-event-range", "--spin1z",
    "--spin2z", "--srate", "--srate-internal",
    "--srate-resample-time-marginalization",
    "--supplementary-likelihood-factor-code",
    "--supplementary-likelihood-factor-function",
    "--supplementary-likelihood-factor-ini", "--time-marginalization",
    "--use-gwsignal", "--use-gwsignal-lmax-nyquist", "--vectorized", "--verbose",
    "--window-shape", "--zero-likelihood",
}


def _register_compat_options(optp):
    """Add every ILE option we do not already define, as accepted-but-ignored."""
    import optparse
    known = set()
    for opt in optp._get_all_options():
        known.update(opt._long_opts)
    g = OptionGroup(optp, "ILE compatibility (accepted; ignored unless implemented)")
    for name in sorted(_ILE_ALL_OPTS):
        if name in known:
            continue
        if name in _ILE_BOOL_OPTS:
            g.add_option(name, action="store_true", default=False,
                         help=optparse.SUPPRESS_HELP)
        elif name in _ILE_APPEND_OPTS:
            g.add_option(name, action="append", default=[],
                         help=optparse.SUPPRESS_HELP)
        else:
            g.add_option(name, default=None, help=optparse.SUPPRESS_HELP)
    optp.add_option_group(g)


def _dest(name):
    return name.lstrip("-").replace("-", "_")


def check_critical_and_report(opts, optp):
    """Fail on science-changing options we do not implement; report the rest."""
    def is_set(name):
        v = getattr(opts, _dest(name), None)
        return bool(v) if name in _ILE_BOOL_OPTS else (v is not None and v != [])

    fatal = []
    if getattr(opts, "calibration_n_realizations", None) not in (None, 1) \
            or is_set("--calibration-export-posterior") \
            or is_set("--calibration-envelope-directory"):
        fatal.append("calibration marginalization (--calibration-*) is not implemented")
    if is_set("--rom-use-basis") or is_set("--rom-group") or is_set("--rom-param"):
        fatal.append("ROM-basis waveforms (--rom-*) are not implemented")
    if any(is_set(x) for x in ("--supplementary-likelihood-factor-code",
                               "--supplementary-likelihood-factor-function",
                               "--supplementary-likelihood-factor-ini")):
        fatal.append("supplementary likelihood factors are not implemented")
    if is_set("--nr-lookup") or is_set("--nr-params") or is_set("--nr-index"):
        fatal.append("NR-waveform templates (--nr-*) are not implemented")
    if is_set("--zero-likelihood"):
        fatal.append("--zero-likelihood is not implemented")
    if is_set("--maximize-only"):
        fatal.append("--maximize-only is not implemented (this driver integrates)")
    if fatal:
        optp.error("Cannot run as a faithful drop-in: " + "; ".join(fatal)
                   + ".  (These would silently change the result if ignored.)")

    # Report accepted-but-ignored options that were actually passed.
    ignored = []
    implemented = {"--cache-file", "--channel-name", "--psd-file",
                   "--data-start-time", "--data-end-time", "--window-shape",
                   "--psd-window-shape", "--mass1", "--mass2", "--spin1z",
                   "--spin2z", "--eff-lambda", "--deff-lambda", "--approximant",
                   "--l-max", "--event-time", "--fmin-template",
                   "--reference-freq", "--fmax", "--srate",
                   "--data-integration-window-half",
                   "--internal-data-storage-window-half", "--d-min", "--d-max",
                   "--d-prior", "--n-max", "--n-chunk", "--output-file",
                   "--event", "--save-samples", "--verbose", "--seed",
                   "--sim-xml", "--sim-grid", "--n-events-to-analyze",
                   "--random-event", "--distance-marginalization",
                   "--time-marginalization", "--vectorized", "--use-gwsignal"}
    for name in sorted(_ILE_ALL_OPTS - implemented):
        if is_set(name):
            ignored.append(name)
    if ignored:
        print("Note: the following ILE options are accepted but IGNORED by the "
              "JAX driver (not yet implemented; behavior may differ from ILE):")
        print("   " + " ".join(ignored))


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def build_parser():
    optp = OptionParser(usage="%prog [options]", description=__doc__)

    g = OptionGroup(optp, "Data input (frame mode)")
    g.add_option("--cache-file", default=None)
    g.add_option("--channel-name", action="append", default=[],
                 help="instrument=channel. Repeatable.")
    g.add_option("--psd-file", action="append", default=[],
                 help="instrument=psd.xml.gz. Repeatable.")
    g.add_option("--data-start-time", type=float, default=None)
    g.add_option("--data-end-time", type=float, default=None)
    g.add_option("--window-shape", type=float, default=0.0)
    optp.add_option_group(g)

    g = OptionGroup(optp, "Data input (injection / self-test mode)")
    g.add_option("--inj-mode", action="store_true", default=False,
                 help="Synthesize zero-noise data from the template at the "
                      "truth sky location (no frames; for testing).")
    g.add_option("--inj-ra", type=float, default=1.2)
    g.add_option("--inj-dec", type=float, default=-0.4)
    g.add_option("--inj-psi", type=float, default=0.7)
    g.add_option("--inj-incl", type=float, default=0.9)
    g.add_option("--inj-phiref", type=float, default=2.1)
    g.add_option("--inj-distance", type=float, default=600.0)
    g.add_option("--inj-detectors", default="H1,L1,V1")
    g.add_option("--inj-deltaF", type=float, default=0.25)
    optp.add_option_group(g)

    g = OptionGroup(optp, "Intrinsic template parameters")
    g.add_option("--mass1", type=float, default=None, help="m1 (solar masses)")
    g.add_option("--mass2", type=float, default=None, help="m2 (solar masses)")
    g.add_option("--spin1z", type=float, default=0.0)
    g.add_option("--spin2z", type=float, default=0.0)
    g.add_option("--eff-lambda", type=float, default=None,
                 help="Tidal Lambda-tilde; converted to lambda1,lambda2.")
    g.add_option("--deff-lambda", type=float, default=None,
                 help="Tidal delta-Lambda-tilde.")
    g.add_option("--approximant", default="IMRPhenomD")
    g.add_option("--l-max", type=int, default=2)
    g.add_option("--sim-xml", default=None,
                 help="ligolw XML of intrinsic params (ChooseWaveformParams array). "
                      "Selects event(s) via --event / --n-events-to-analyze.")
    g.add_option("--sim-grid", default=None,
                 help="ASCII grid file (names=True header) of intrinsic params.")
    g.add_option("--event", type=int, default=0,
                 help="Index of the first event to analyze from --sim-xml/--sim-grid.")
    g.add_option("--n-events-to-analyze", type=int, default=1,
                 help="Number of consecutive events to analyze (batch loop).")
    g.add_option("--random-event", action="store_true", default=False,
                 help="Draw events at random from the table instead of in order.")
    optp.add_option_group(g)

    g = OptionGroup(optp, "Conditioning / windows / frequencies")
    g.add_option("--event-time", type=float, default=None,
                 help="GPS geocenter event time (fiducial epoch).")
    g.add_option("--fmin-template", type=float, default=30.0)
    g.add_option("--reference-freq", type=float, default=30.0)
    g.add_option("--fmax", type=float, default=1000.0)
    g.add_option("--srate", type=float, default=4096.0)
    g.add_option("--data-integration-window-half", type=float, default=75e-3)
    g.add_option("--internal-data-storage-window-half", type=float, default=0.15)
    optp.add_option_group(g)

    g = OptionGroup(optp, "Extrinsic exploration / sampling")
    g.add_option("--mode", default="laplace-is",
                 choices=["prior-mc", "laplace-is", "map", "nuts",
                          "multistart-nuts", "flowmc", "flowmc-phimarg",
                          "flowmc-phipsimarg", "flowmc-dpsimarg", "nuts-phimarg"],
                 help="Extrinsic exploration: prior-mc | laplace-is | map | "
                      "nuts (single-chain) | multistart-nuts (mode-covering) | "
                      "flowmc (normalizing-flow, 5-D) | "
                      "flowmc-phimarg (phi_ref-marginalised flowMC, 4-D — removes "
                      "the psi/phi_ref degeneracy ridge) | "
                      "flowmc-dpsimarg (distance+psi-marginalised flowMC, 4-D "
                      "(ra,dec,phiref,incl) — psi-only scan, cheaper than phimarg) | "
                      "nuts-phimarg (Fisher-whitened multi-start NUTS on the 4-D "
                      "phi_ref-marginalised posterior — whitening keeps the NUTS "
                      "step size O(1) at any SNR; recommended for high SNR). "
                      "The last five need --distance-marginalization.")
    g.add_option("--num-warmup", type=int, default=500,
                 help="NUTS warmup iterations (--mode nuts).")
    g.add_option("--num-samples", type=int, default=2000,
                 help="NUTS posterior samples per chain (--mode nuts).")
    g.add_option("--num-chains", type=int, default=1)
    g.add_option("--no-flow-reuse", action="store_true", default=False,
                 help="Disable bootstrapping the trained flow across "
                      "--n-events-to-analyze (re-train from scratch each event).")
    # ILE-compatible semantics: store_true, default OFF.  When ON, distance is
    # marginalized analytically (5-D angular problem, well conditioned); when
    # OFF, distance is sampled explicitly (6-D).  Required for --mode nuts.
    g.add_option("--distance-marginalization", action="store_true", default=False,
                 help="Marginalize distance analytically (regulates the "
                      "factored-likelihood amplitude degeneracy; required for nuts).")
    g.add_option("--n-max", type=int, default=300000)
    g.add_option("--n-chunk", type=int, default=8000)
    g.add_option("--d-min", type=float, default=1.0, help="Min distance (Mpc).")
    g.add_option("--d-max", type=float, default=10000.0, help="Max distance (Mpc).")
    g.add_option("--distance-grid-points", type=int, default=256)
    g.add_option("--phase-marginalization", action="store_true", default=False)
    g.add_option("--n-phi", type=int, default=32,
                 help="phi_ref grid size for --mode flowmc-phimarg (default 32; "
                      "use 64-128 for l-max>=4 or production quality).")
    g.add_option("--n-psi", type=int, default=8,
                 help="psi (polarization) grid size for --mode flowmc-phipsimarg "
                      "(default 8; spin-2 -> exponential convergence, ~8 exact for "
                      "l_max=2).  Keep nphi*npsi small (~64) -- the (phi,psi) scan "
                      "cost scales with it.")
    # flowMC tuning (modes flowmc / flowmc-phimarg).  Defaults match
    # samplers.flowmc_sample*; exposed so pipeline Makefiles can tune them.
    g.add_option("--n-training-loops", type=int, default=4,
                 help="flowMC training loops (flow fit refinements).")
    g.add_option("--n-production-loops", type=int, default=4,
                 help="flowMC production loops (post-training sampling).")
    g.add_option("--n-epochs", type=int, default=10,
                 help="flowMC NF training epochs per training loop.")
    g.add_option("--n-local-steps", type=int, default=20,
                 help="flowMC local (MALA) steps per loop.")
    g.add_option("--n-global-steps", type=int, default=20,
                 help="flowMC global (flow-proposal) steps per loop.")
    g.add_option("--n-prior-pilot", type=int, default=8000,
                 help="Prior pilot draws used to initialize chains (and to "
                      "seed multi-start optimizers).")
    g.add_option("--adapt-weight-exponent", type=float, default=1.0,
                 help="Static likelihood-tempering exponent beta (a la old-RIFT): "
                      "flowMC samples the broadened target L^beta * pi and "
                      "reweights to the true posterior. beta=1 (default) is OFF; "
                      "beta<1 broadens (helps the flow find sharp high-SNR peaks). "
                      "Maps to sampler temper = 1/beta (modes flowmc, flowmc-phimarg).")
    g.add_option("--adapt-adapt", action="store_true", default=False,
                 help="Adaptive likelihood tempering (flowmc-phimarg): anneal the "
                      "tempering exponent inv_T from --temper-init up to 1.0, "
                      "warm-starting the flow each stage, choosing each step so the "
                      "tempered-reweight ESS stays >= --temper-ess-frac. The "
                      "high-SNR path that avoids the static-tempering neff->1 "
                      "collapse. Overrides --adapt-weight-exponent when set.")
    g.add_option("--temper-init", type=float, default=0.02,
                 help="Initial (widest) tempering exponent inv_T for --adapt-adapt.")
    g.add_option("--temper-ess-frac", type=float, default=0.5,
                 help="Target tempered-reweight ESS fraction per anneal step (--adapt-adapt).")
    g.add_option("--temper-max-stages", type=int, default=16,
                 help="Max anneal stages before forcing the final inv_T=1 pass (--adapt-adapt).")
    g.add_option("--temper-max-dbeta", type=float, default=0.15,
                 help="Cap on the inv_T step per anneal stage, so the thermodynamic-"
                      "integration evidence trapezoid stays accurate (--adapt-adapt).")
    g.add_option("--fisher-precondition", action="store_true", default=False,
                 help="flowmc-phi(psi)marg: whiten the sample space around the MAP "
                      "using the observed Fisher (theta=MAP+F^{-1/2} y) so the "
                      "(1/SNR)-narrow high-SNR posterior is O(1) for the flow. "
                      "Restarts chains at the MAP. Recommended at SNR>=320 to avoid "
                      "the sky-sample collapse; harmless (~no-op) at low SNR.")
    g.add_option("--smc-puffball", action="store_true", default=False,
                 help="flowmc-phi(psi)marg: use adaptive SMC with a cloud-covariance "
                      "'puffball' random-walk move instead of the flow. Robust at any "
                      "SNR (sample->puffball->sample, like RIFT-AV/nested sampling); "
                      "does not collapse on sharp high-SNR peaks. Recommended at "
                      "SNR>=320 for reliable sky samples.")
    g.add_option("--smc-walkers", type=int, default=2000,
                 help="SMC puffball: number of walkers in the cloud (default 2000).")
    g.add_option("--smc-move-steps", type=int, default=10,
                 help="SMC puffball: random-walk Metropolis moves per temperature rung.")
    g.add_option("--smc-puff-scale", type=float, default=1.0,
                 help="SMC puffball: proposal scale x sqrt(cloud covariance) (~1.0).")
    g.add_option("--fisher-is-samples", type=int, default=0,
                 help="flowmc-phi(psi)marg high-SNR FALLBACK: after annealing, draw "
                      "this many sky samples by Fisher-whitened importance sampling "
                      "about the (Newton-polished) MAP and reweight by the true lnL "
                      "(no flow training -> immune to the SNR>=640 NF-collapse). "
                      "Overrides the exported samples; TI evidence is unchanged. "
                      "0=off; ~40000 recommended at SNR>=640. Implies --fisher-precondition.")
    g.add_option("--interp", default="linear", choices=["linear", "nearest"])
    g.add_option("--sky-coordinates", default="equatorial",
                 choices=["equatorial", "network"],
                 help="Optional: 'network' samples the sky in the two-detector "
                      "baseline frame (folds the time-delay ring); multistart-nuts "
                      "only, falls back to equatorial if <2 detectors.")
    g.add_option("--proposal-inflate", type=float, default=3.0,
                 help="Scale on the adaptive Gaussian proposal covariance.")
    g.add_option("--seed", type=int, default=0)
    optp.add_option_group(g)

    g = OptionGroup(optp, "Output")
    g.add_option("--output-file", default=None)
    g.add_option("--save-samples", action="store_true", default=False)
    g.add_option("--verbose", action="store_true", default=False)
    optp.add_option_group(g)

    # Register every remaining ILE option for drop-in compatibility (accepted,
    # silently ignored unless implemented; science-critical ones fail below).
    _register_compat_options(optp)
    return optp


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------
def make_template(opts, fiducial_epoch, deltaF, deltaT):
    P = lalsimutils.ChooseWaveformParams()
    P.m1 = opts.mass1 * MSUN
    P.m2 = opts.mass2 * MSUN
    P.s1z = opts.spin1z
    P.s2z = opts.spin2z
    if opts.eff_lambda is not None:
        P.lambda1, P.lambda2 = lalsimutils.tidal_lambda_from_tilde(
            opts.mass1, opts.mass2, opts.eff_lambda, opts.deff_lambda or 0.0)
    P.fmin = opts.fmin_template
    P.fref = opts.reference_freq
    P.deltaT = deltaT
    P.deltaF = deltaF
    P.fmax = 0.0
    P.approx = lalsim.GetApproximantFromString(opts.approximant)
    P.radec = True
    P.tref = fiducial_epoch
    P.dist = 1000.0 * 1e6 * PC
    return P


def load_templates(opts, fiducial_epoch, deltaF, deltaT):
    """Return a list of intrinsic templates (ChooseWaveformParams).

    Mirrors the production driver's --sim-xml / --sim-grid / --mass1,--mass2
    handling, including the [event : event+n_events] slice and the
    extrinsic-zeroing / fiducial-distance conventions.  Extrinsic parameters
    (sky, distance, inclination, phase, psi) are placeholders -- they are the
    quantities this code samples/marginalizes.
    """
    def _finalize(P):
        P.radec = True
        P.fref = opts.reference_freq
        P.fmin = opts.fmin_template
        P.deltaT = deltaT
        P.deltaF = deltaF
        P.fmax = 0.0
        P.tref = fiducial_epoch
        P.dist = 1000.0 * 1e6 * PC  # fiducial template distance (= distMpcRef)
        if opts.approximant:
            P.approx = lalsim.GetApproximantFromString(opts.approximant)
        return P

    if opts.sim_xml:
        P_all = lalsimutils.xml_to_ChooseWaveformParams_array(str(opts.sim_xml))
        if opts.random_event:
            idx = np.random.choice(len(P_all),
                                   size=min(opts.n_events_to_analyze, len(P_all)),
                                   replace=False)
            P_list = [P_all[i] for i in idx]
        else:
            if len(P_all) <= opts.event:
                print(" Event index out of range; soft exit"); sys.exit(0)
            hi = min(len(P_all), opts.event + opts.n_events_to_analyze)
            P_list = list(P_all[opts.event:hi])
        return [_finalize(P) for P in P_list]

    if opts.sim_grid:
        grid = np.genfromtxt(opts.sim_grid, names=True)
        grid = np.atleast_1d(grid)
        names = list(grid.dtype.names)
        params = [n for n in names if n in set(lalsimutils.valid_params)]
        if len(grid) <= opts.event:
            print(" Event index out of range; soft exit"); sys.exit(0)
        hi = min(len(grid), opts.event + opts.n_events_to_analyze)
        P_list = []
        for row in grid[opts.event:hi]:
            P = lalsimutils.ChooseWaveformParams()
            for name in params:
                val = float(row[name])
                if hasattr(P, name):
                    setattr(P, name, val)
                else:
                    P.assign_param(name, val)
            if P.m1 < 1e15:           # masses given in solar masses -> SI
                P.m1 *= MSUN; P.m2 *= MSUN
            P_list.append(_finalize(P))
        return P_list

    # single template from --mass1/--mass2
    if opts.mass1 is None or opts.mass2 is None:
        raise SystemExit("Provide --mass1/--mass2, or --sim-xml / --sim-grid.")
    return [make_template(opts, fiducial_epoch, deltaF, deltaT)]


def load_injection(opts, fiducial_epoch):
    detectors = opts.inj_detectors.split(",")
    deltaT = 1.0 / opts.srate
    P = make_template(opts, fiducial_epoch, opts.inj_deltaF, deltaT)
    P.phi, P.theta = opts.inj_ra, opts.inj_dec
    P.psi, P.incl, P.phiref = opts.inj_psi, opts.inj_incl, opts.inj_phiref
    P.dist = opts.inj_distance * 1e6 * PC
    data_dict, psd_dict = {}, {}
    for det in detectors:
        Pdet = P.copy(); Pdet.detector = det
        data_dict[det] = lalsimutils.non_herm_hoff(Pdet)
        psd_dict[det] = lalsim.SimNoisePSDaLIGOZeroDetHighPower
    P_template = make_template(opts, fiducial_epoch, opts.inj_deltaF, deltaT)
    return P_template, data_dict, psd_dict, detectors, True


def load_frames(opts, fiducial_epoch):
    deltaT = 1.0 / opts.srate
    data_dict, psd_dict = {}, {}
    for inst, chan in map(lambda c: c.split("="), opts.channel_name):
        if opts.verbose:
            print("Reading channel %s:%s from %s" % (inst, chan, opts.cache_file))
        data_dict[inst] = lalsimutils.frame_data_to_non_herm_hoff(
            opts.cache_file, inst + ":" + chan,
            start=opts.data_start_time, stop=opts.data_end_time,
            window_shape=opts.window_shape, deltaT=deltaT)
    for inst, psdf in map(lambda c: c.split("="), opts.psd_file):
        if opts.verbose:
            print("Reading PSD for %s from %s" % (inst, psdf))
        psd_dict[inst] = lalsimutils.get_psd_series_from_xmldoc(psdf, inst)
        psd_dict[inst] = lalsimutils.resample_psd_series(
            psd_dict[inst], data_dict[inst].deltaF)
    detectors = list(data_dict.keys())
    return data_dict, psd_dict, detectors, False


# ---------------------------------------------------------------------------
# Priors (the angular block is shared; distance added only when NOT marginalized)
# ---------------------------------------------------------------------------
def sample_prior(n, opts, rng, with_distance):
    ra = rng.uniform(0.0, 2 * np.pi, n)
    dec = np.arcsin(rng.uniform(-1.0, 1.0, n))
    psi = rng.uniform(0.0, np.pi, n)
    incl = np.arccos(rng.uniform(-1.0, 1.0, n))
    phiref = rng.uniform(0.0, 2 * np.pi, n)
    cols = [ra, dec, psi, incl, phiref]
    if with_distance:
        u = rng.uniform(0.0, 1.0, n)
        dmin3, dmax3 = opts.d_min ** 3, opts.d_max ** 3
        cols.append((dmin3 + u * (dmax3 - dmin3)) ** (1.0 / 3.0))
    theta = np.stack(cols, axis=-1)
    return theta, log_prior(theta, opts, with_distance)


def log_prior(theta, opts, with_distance):
    ra, dec, psi, incl, phiref = [theta[..., i] for i in range(5)]
    inb = ((ra >= 0) & (ra <= 2 * np.pi) & (dec >= -np.pi / 2) & (dec <= np.pi / 2)
           & (psi >= 0) & (psi <= np.pi) & (incl >= 0) & (incl <= np.pi)
           & (phiref >= 0) & (phiref <= 2 * np.pi))
    with np.errstate(divide="ignore", invalid="ignore"):
        logp = (np.log(np.cos(dec)) - np.log(2.0)
                + np.log(np.sin(incl)) - np.log(2.0)
                - np.log(2 * np.pi) - np.log(np.pi) - np.log(2 * np.pi))
        if with_distance:
            dist = theta[..., 5]
            inb = inb & (dist >= opts.d_min) & (dist <= opts.d_max)
            dmin3, dmax3 = opts.d_min ** 3, opts.d_max ** 3
            logp = logp + np.log(3.0) + 2 * np.log(dist) - np.log(dmax3 - dmin3)
    return np.where(inb, logp, -np.inf)


# ---------------------------------------------------------------------------
# Batched lnL
# ---------------------------------------------------------------------------
def eval_lnL(like, theta, opts, with_distance):
    N = theta.shape[0]
    out = np.empty(N)
    for i in range(0, N, opts.n_chunk):
        sl = slice(i, min(i + opts.n_chunk, N))
        cols = [theta[sl, j] for j in range(theta.shape[1])]
        out[sl] = np.asarray(like.log_likelihood(*cols))
    return out


# ---------------------------------------------------------------------------
# Evidence helpers
# ---------------------------------------------------------------------------
def evidence_from_logweights(logw):
    """(logZ, sigma/Z, neff) for Z = E[w] from log importance weights."""
    fin = np.isfinite(logw)
    logw = logw[fin]
    if logw.size == 0:
        return -np.inf, np.inf, 0.0
    m = np.max(logw)
    w = np.exp(logw - m)
    n = len(w)
    Zhat = np.mean(w)
    logZ = m + np.log(Zhat)
    sigma_over_Z = np.sqrt(np.var(w) / n) / Zhat
    neff = (np.sum(w) ** 2) / np.sum(w ** 2)
    return logZ, sigma_over_Z, neff


def _gaussian_logq(theta, mu, cov):
    diff = theta - mu[None, :]
    sol = np.linalg.solve(cov, diff.T).T
    return (-0.5 * np.einsum("ij,ij->i", diff, sol)
            - 0.5 * theta.shape[1] * np.log(2 * np.pi)
            - 0.5 * np.linalg.slogdet(cov)[1])


def _moment_match(theta, logL):
    w = np.exp(logL - np.max(logL))
    w = w / np.sum(w)
    mu = np.sum(w[:, None] * theta, axis=0)
    d = theta - mu[None, :]
    cov = (w[:, None, None] * d[:, :, None] * d[:, None, :]).sum(axis=0)
    cov += 1e-9 * np.eye(theta.shape[1]) * (np.trace(cov) / theta.shape[1] + 1e-12)
    return mu, cov


# ---------------------------------------------------------------------------
# Estimators
# ---------------------------------------------------------------------------
def run_prior_mc(like, opts, rng, dim, with_distance):
    theta, _ = sample_prior(opts.n_max, opts, rng, with_distance)
    lnL = eval_lnL(like, theta, opts, with_distance)
    logZ, sig, neff = evidence_from_logweights(lnL)  # draw from prior -> w = L
    return logZ, sig, neff, opts.n_max, theta, lnL


def run_laplace_is(like, opts, rng, dim, with_distance, n_adapt=2):
    """Prior-seeded adaptive Gaussian importance sampling (Z = E_q[L p / q])."""
    n_pilot = min(opts.n_max // 4, 40000)
    theta_p, _ = sample_prior(n_pilot, opts, rng, with_distance)
    lnL_p = eval_lnL(like, theta_p, opts, with_distance)
    logL_p = lnL_p + log_prior(theta_p, opts, with_distance)
    mu, cov = _moment_match(theta_p, logL_p)

    per_round = max(opts.n_max // (n_adapt + 1), 1)
    all_theta, all_logw, all_lnL = [], [], []
    for r in range(n_adapt + 1):
        cov_use = cov * opts.proposal_inflate
        Lc = np.linalg.cholesky(cov_use + 1e-12 * np.eye(dim))
        z = rng.standard_normal((per_round, dim))
        theta = mu[None, :] + z @ Lc.T
        logq = _gaussian_logq(theta, mu, cov_use)
        logp = log_prior(theta, opts, with_distance)
        valid = np.isfinite(logp)
        lnL = np.full(per_round, -np.inf)
        if valid.any():
            lnL[valid] = eval_lnL(like, theta[valid], opts, with_distance)
        logw = np.where(valid, lnL + logp - logq, -np.inf)
        all_theta.append(theta); all_logw.append(logw); all_lnL.append(lnL)
        good = np.isfinite(logw)
        if r < n_adapt and good.sum() > 50:
            mu, cov = _moment_match(theta[good], logw[good])

    theta = np.concatenate(all_theta); logw = np.concatenate(all_logw)
    lnL = np.concatenate(all_lnL)
    logZ, sig, neff = evidence_from_logweights(logw)
    return logZ, sig, neff, len(theta), theta, lnL


def run_nuts(like, opts, rng, with_distance):
    """Gradient-based NUTS over the distance-marginalized angular posterior.

    Demonstrates the AD payoff: NUTS uses the exact JAX gradient of the
    (smooth, distance-marginalized) log-likelihood to climb and sample the
    narrow extrinsic peak efficiently -- where brute-force MC stalls.  The
    posterior samples then seed a high-neff importance estimate of the evidence
    (a Gaussian proposal moment-matched to the NUTS draws), and are written out
    as the per-sample product (the natural seed for a flowMC / AV closeout).
    """
    if with_distance:
        raise SystemExit("--mode nuts requires distance marginalization (the "
                         "bare 6-D likelihood is degenerate); pass "
                         "--distance-marginalization.")
    import numpyro
    import numpyro.distributions as dist
    from numpyro.infer import MCMC, NUTS
    import jax.numpy as jnp

    two_pi = float(2 * np.pi)

    def model():
        ra = numpyro.sample("ra", dist.Uniform(0.0, two_pi))
        sin_dec = numpyro.sample("sin_dec", dist.Uniform(-1.0, 1.0))
        psi = numpyro.sample("psi", dist.Uniform(0.0, float(np.pi)))
        cos_incl = numpyro.sample("cos_incl", dist.Uniform(-1.0, 1.0))
        phiref = numpyro.sample("phiref", dist.Uniform(0.0, two_pi))
        dec = jnp.arcsin(sin_dec)
        incl = jnp.arccos(cos_incl)
        lnL = like._scalar(jnp.stack([ra, dec, psi, incl, phiref]))
        numpyro.factor("loglike", lnL)

    # Seed NUTS at the best prior draw so it climbs the GLOBAL mode rather than
    # one of the (many) secondary sky modes from the detector time-delay ring.
    seed_theta, _ = sample_prior(4000, opts, rng, with_distance=False)
    seed_lnL = eval_lnL(like, seed_theta, opts, with_distance=False)
    th0 = seed_theta[np.argmax(seed_lnL)]
    init_vals = {
        "ra": float(th0[0]), "sin_dec": float(np.sin(th0[1])),
        "psi": float(th0[2]), "cos_incl": float(np.cos(th0[3])),
        "phiref": float(th0[4]),
    }
    print("  NUTS seeded at best-of-%d prior draw (lnL=%.2f)"
          % (len(seed_theta), seed_lnL.max()))
    from numpyro.infer import init_to_value
    key = jax.random.PRNGKey(opts.seed)
    kernel = NUTS(model, target_accept_prob=0.8,
                  init_strategy=init_to_value(values=init_vals))
    mcmc = MCMC(kernel, num_warmup=opts.num_warmup,
                num_samples=opts.num_samples, num_chains=opts.num_chains,
                progress_bar=opts.verbose)
    mcmc.run(key)
    s = mcmc.get_samples()
    ra = np.asarray(s["ra"]); dec = np.arcsin(np.asarray(s["sin_dec"]))
    psi = np.asarray(s["psi"]); incl = np.arccos(np.asarray(s["cos_incl"]))
    phiref = np.asarray(s["phiref"])
    theta = np.stack([ra, dec, psi, incl, phiref], axis=-1)
    lnL = eval_lnL(like, theta, opts, with_distance=False)
    print("  NUTS drew %d samples; lnL in [%.2f, %.2f], peak at:"
          % (len(theta), lnL.min(), lnL.max()))
    for nm, v in zip(ANG_NAMES, theta[np.argmax(lnL)]):
        print("     %-8s = %.5f" % (nm, v))

    # NUTS-seeded importance estimate of the evidence (Gaussian proposal
    # moment-matched to the posterior draws -> high neff vs prior seeding).
    mu, cov = _moment_match(theta, np.zeros(len(theta)))  # posterior moments
    n_is = min(opts.n_max, 40000)
    Lc = np.linalg.cholesky(cov * opts.proposal_inflate + 1e-12 * np.eye(5))
    z = rng.standard_normal((n_is, 5))
    th_is = mu[None, :] + z @ Lc.T
    logq = _gaussian_logq(th_is, mu, cov * opts.proposal_inflate)
    logp = log_prior(th_is, opts, with_distance=False)
    valid = np.isfinite(logp)
    lnL_is = np.full(n_is, -np.inf)
    if valid.any():
        lnL_is[valid] = eval_lnL(like, th_is[valid], opts, with_distance=False)
    logw = np.where(valid, lnL_is + logp - logq, -np.inf)
    logZ, sig, neff = evidence_from_logweights(logw)
    return logZ, sig, neff, n_is, theta, lnL


def run_map(like, opts, rng, dim, with_distance):
    """Gradient-ascend the (well-conditioned, distance-marginalized) peak."""
    from scipy.optimize import minimize
    ang_bounds = [(0, 2 * np.pi), (-np.pi / 2 + 1e-3, np.pi / 2 - 1e-3),
                  (0, np.pi), (1e-3, np.pi - 1e-3), (0, 2 * np.pi)]
    bounds = ang_bounds + ([(opts.d_min, opts.d_max)] if with_distance else [])
    theta_seed, _ = sample_prior(4000, opts, rng, with_distance)
    lnL_seed = eval_lnL(like, theta_seed, opts, with_distance)
    x0 = theta_seed[np.argmax(lnL_seed)]

    def negf(x):
        v, g = like.value_and_grad(x)
        return -float(v), -np.asarray(g)

    res = minimize(negf, x0, jac=True, method="L-BFGS-B", bounds=bounds,
                   options={"maxiter": 400})
    return res.x, -res.fun


# ---------------------------------------------------------------------------
# Output
# ---------------------------------------------------------------------------
def write_dat(opts, P, out_index, event_id, logZ, sigma_over_Z, ntotal, neff):
    if not opts.output_file:
        return
    m1, m2 = P.m1 / MSUN, P.m2 / MSUN
    fname = opts.output_file + "_" + str(out_index) + "_" + ".dat"
    row = np.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y,
                     P.s2z, logZ, sigma_over_Z, ntotal, neff]])
    np.savetxt(fname, row,
               header="event_id m1 m2 s1x s1y s1z s2x s2y s2z lnL sigma_lnL ntotal neff")
    print("Wrote %s" % fname)


def write_samples(opts, out_index, theta, lnL, with_distance):
    if not (opts.output_file and opts.save_samples) or theta is None:
        return
    good = np.isfinite(lnL)
    ndim = theta.shape[1] if theta.ndim == 2 else len(theta)
    if with_distance:
        # 6-D: ra, dec, psi, incl, phiref, dist
        cols = np.column_stack([theta[good, 0], theta[good, 1], theta[good, 5],
                                theta[good, 3], theta[good, 2], theta[good, 4],
                                lnL[good]])
        hdr = "right_ascension declination distance inclination psi phi_orb loglikelihood"
    elif ndim == 4 and opts.mode == "flowmc-dpsimarg":
        # 4-D (flowmc-dpsimarg): theta = ra, dec, phiref, incl (psi marginalised,
        # phi_ref sampled).  Write ra, dec, incl, phi_orb.
        cols = np.column_stack([theta[good, 0], theta[good, 1],
                                theta[good, 3], theta[good, 2], lnL[good]])
        hdr = "right_ascension declination inclination phi_orb loglikelihood"
    elif ndim == 4:
        # 4-D (flowmc-phimarg): ra, dec, psi, incl (phi_ref marginalised out)
        cols = np.column_stack([theta[good, 0], theta[good, 1],
                                theta[good, 3], theta[good, 2], lnL[good]])
        hdr = "right_ascension declination inclination psi loglikelihood"
    elif ndim == 3:
        # 3-D (flowmc-phipsimarg): ra, dec, incl (phi_ref AND psi marginalised out)
        cols = np.column_stack([theta[good, 0], theta[good, 1],
                                theta[good, 2], lnL[good]])
        hdr = "right_ascension declination inclination loglikelihood"
    else:
        # 5-D: ra, dec, psi, incl, phiref
        cols = np.column_stack([theta[good, 0], theta[good, 1], theta[good, 3],
                                theta[good, 2], theta[good, 4], lnL[good]])
        hdr = "right_ascension declination inclination psi phi_orb loglikelihood"
    sname = opts.output_file + "_" + str(out_index) + "_samples.dat"
    np.savetxt(sname, cols, header=hdr)
    print("Wrote %s (%d samples)" % (sname, int(good.sum())))


def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch,
                rng, out_index, event_id, flow_state=None):
    """Build the JAX likelihood for one intrinsic template and run --mode.

    Returns ``(value, flow_state)`` where ``flow_state`` is the trained-flow
    state to bootstrap the next event (``--mode flowmc`` only; ``None`` else).
    """
    print("Building JAX likelihood (PrecomputeLikelihoodTerms + pack)...")
    like_data, extras = build_data_from_precompute(
        P.copy(), data_dict, psd_dict, fiducial_epoch,
        opts.internal_data_storage_window_half, opts.data_integration_window_half,
        opts.l_max, opts.fmax,
        analyticPSD_Q=analyticPSD_Q, verbose=opts.verbose,
        use_gwsignal=bool(getattr(opts, "use_gwsignal", False)),
        use_gwsignal_approx=(opts.approximant if getattr(opts, "use_gwsignal", False) else None))
    print("  modes:", like_data.lms, "  guessed SNR:", extras["guess_snr"])

    with_distance = not opts.distance_marginalization
    if opts.mode in ("flowmc-phimarg", "nuts-phimarg"):
        # phi_ref-marginalised: requires distance marginalisation (baked in);
        # produces a 4-D (ra, dec, psi, incl) posterior.
        if not opts.distance_marginalization:
            raise SystemExit("--mode %s requires --distance-marginalization."
                             % opts.mode)
        from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiMargLikelihood
        nphi = getattr(opts, "n_phi", 32)
        print("Distance + phi_ref marginalization: ON (grid=%d, nphi=%d, d in [%g,%g] Mpc)"
              % (opts.distance_grid_points, nphi, opts.d_min, opts.d_max))
        like = JAXDistPhiMargLikelihood(
            like_data, opts.d_min, opts.d_max,
            nphi=nphi, n_grid=opts.distance_grid_points,
            interp=opts.interp, guess_snr=extras["guess_snr"])
        if getattr(like, "dist_grid_info", {}).get("mode") == "adaptive":
            gi = like.dist_grid_info
            print("  distance grid: ADAPTIVE  d_peak=%.3g Mpc  sigma_d=%.3g Mpc  npts=%d"
                  % (gi["d_peak"], gi["sigma_d"], gi["n"]))
        with_distance = False   # distance is already marginalised inside like
        dim = 4
    elif opts.mode == "flowmc-phipsimarg":
        # distance + phi_ref + psi marginalised: 3-D (ra, dec, incl) posterior.
        if not opts.distance_marginalization:
            raise SystemExit("--mode flowmc-phipsimarg requires --distance-marginalization.")
        from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood
        nphi = getattr(opts, "n_phi", 32)
        npsi = getattr(opts, "n_psi", 16)
        print("Distance + phi_ref + psi marginalization: ON (grid=%d, nphi=%d, npsi=%d, d in [%g,%g] Mpc)"
              % (opts.distance_grid_points, nphi, npsi, opts.d_min, opts.d_max))
        like = JAXDistPhiPsiMargLikelihood(
            like_data, opts.d_min, opts.d_max, nphi=nphi, npsi=npsi,
            n_grid=opts.distance_grid_points, interp=opts.interp,
            guess_snr=extras["guess_snr"])
        if getattr(like, "dist_grid_info", {}).get("mode") == "adaptive":
            gi = like.dist_grid_info
            print("  distance grid: ADAPTIVE  d_peak=%.3g Mpc  sigma_d=%.3g Mpc  npts=%d"
                  % (gi["d_peak"], gi["sigma_d"], gi["n"]))
        with_distance = False
        dim = 3
    elif opts.mode == "flowmc-dpsimarg":
        # distance + psi marginalised, phi_ref SAMPLED: 4-D (ra,dec,phiref,incl).
        if not opts.distance_marginalization:
            raise SystemExit("--mode flowmc-dpsimarg requires --distance-marginalization.")
        from RIFT.likelihood.jax_ile.wrapper import JAXDistPsiMargLikelihood
        npsi = getattr(opts, "n_psi", 8)
        print("Distance + psi marginalization (phi_ref sampled): ON "
              "(grid=%d, npsi=%d, d in [%g,%g] Mpc)"
              % (opts.distance_grid_points, npsi, opts.d_min, opts.d_max))
        like = JAXDistPsiMargLikelihood(
            like_data, opts.d_min, opts.d_max, npsi=npsi,
            n_grid=opts.distance_grid_points, interp=opts.interp,
            guess_snr=extras["guess_snr"])
        if getattr(like, "dist_grid_info", {}).get("mode") == "adaptive":
            gi = like.dist_grid_info
            print("  distance grid: ADAPTIVE  d_peak=%.3g Mpc  sigma_d=%.3g Mpc  npts=%d"
                  % (gi["d_peak"], gi["sigma_d"], gi["n"]))
        with_distance = False
        dim = 4
    elif opts.distance_marginalization:
        print("Distance marginalization: ON (grid=%d, d in [%g,%g] Mpc)"
              % (opts.distance_grid_points, opts.d_min, opts.d_max))
        like = JAXDistanceMarginalizedLikelihood(
            like_data, opts.d_min, opts.d_max, n_grid=opts.distance_grid_points,
            interp=opts.interp, phase_marginalization=opts.phase_marginalization)
        dim = 5
    else:
        like = JAXExtrinsicLikelihood(
            like_data, interp=opts.interp,
            phase_marginalization=opts.phase_marginalization)
        dim = 6

    if opts.mode == "map":
        theta_map, lnL_map = run_map(like, opts, rng, dim, with_distance)
        fish = like.fisher(theta_map)
        names = FULL_NAMES if with_distance else ANG_NAMES
        print("peak lnL = %.5f" % lnL_map)
        for nm, v in zip(names, theta_map):
            print("   %-8s = %.6f" % (nm, v))
        print("Fisher diag:", np.array2string(np.diag(fish), precision=4))
        write_dat(opts, P, out_index, event_id, lnL_map, float("nan"), 0, float("nan"))
        return lnL_map, None

    out_flow_state = None
    if opts.mode in ("multistart-nuts", "flowmc", "flowmc-phimarg",
                     "flowmc-phipsimarg", "flowmc-dpsimarg", "nuts-phimarg"):
        if opts.mode not in ("flowmc-phimarg", "flowmc-phipsimarg",
                             "flowmc-dpsimarg", "nuts-phimarg") and with_distance:
            raise SystemExit("--mode %s requires --distance-marginalization "
                             "(it samples the 5-D angular posterior)." % opts.mode)
        from RIFT.likelihood.jax_ile import samplers as _samplers
        n_starts = opts.num_chains if opts.num_chains and opts.num_chains > 1 else 8
        if opts.mode == "nuts-phimarg":
            # Fisher-whitened multi-start NUTS on the 4-D phimarg posterior.
            res = _samplers.fisher_nuts_sample_phimarg(
                like, num_warmup=opts.num_warmup,
                num_samples=opts.num_samples,
                n_starts=max(n_starts, 8),
                n_prior_pilot=opts.n_prior_pilot,
                seed=opts.seed, verbose=opts.verbose)
        elif opts.mode == "multistart-nuts":
            res = _samplers.multistart_nuts(
                like, opts.d_min, opts.d_max, n_starts=n_starts,
                num_warmup=opts.num_warmup, num_samples=opts.num_samples,
                seed=opts.seed, sky_coords=opts.sky_coordinates,
                verbose=opts.verbose)
        elif (opts.mode in ("flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg")
              and getattr(opts, "smc_puffball", False)):
            # Robust high-SNR alternative to the flow: adaptive SMC with a
            # cloud-covariance "puffball" random-walk move (sample->puffball->sample,
            # a la RIFT-AV / nested sampling).  Does not collapse on sharp peaks.
            res = _samplers.smc_puffball_sample(
                like, opts.d_min, opts.d_max,
                n_walkers=opts.smc_walkers, n_move=opts.smc_move_steps,
                ess_frac=opts.temper_ess_frac, max_dbeta=opts.temper_max_dbeta,
                max_stages=max(opts.temper_max_stages, 80),
                puff_scale=opts.smc_puff_scale,
                seed=opts.seed, verbose=opts.verbose)
        elif opts.mode in ("flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg"):
            # 4-D phi-marg (ra,dec,psi,incl), 3-D phi+psi-marg (ra,dec,incl), or
            # 4-D d+psi-marg (ra,dec,phiref,incl); flowmc_sample_phimarg is
            # dimension-agnostic (helpers chosen from ANGULAR_PARAM_ORDER).
            # Static tempering: beta = --adapt-weight-exponent -> temper = 1/beta.
            _beta = float(opts.adapt_weight_exponent)
            _temper = 1.0 / _beta if _beta > 0 else 1.0
            res = _samplers.flowmc_sample_phimarg(
                like, opts.d_min, opts.d_max,
                n_chains=max(n_starts, 20),
                n_local_steps=opts.n_local_steps,
                n_global_steps=opts.n_global_steps,
                n_training_loops=opts.n_training_loops,
                n_production_loops=opts.n_production_loops,
                n_epochs=opts.n_epochs,
                n_prior_pilot=opts.n_prior_pilot, temper=_temper,
                temper_adapt=opts.adapt_adapt, temper_init=opts.temper_init,
                temper_ess_frac=opts.temper_ess_frac,
                temper_max_stages=opts.temper_max_stages,
                temper_max_dbeta=opts.temper_max_dbeta,
                fisher_precondition=(opts.fisher_precondition
                                     or opts.fisher_is_samples > 0),
                fisher_is_samples=opts.fisher_is_samples,
                seed=opts.seed, reuse_state=flow_state, verbose=opts.verbose)
            out_flow_state = res.get("flow_state")
        else:
            _beta = float(opts.adapt_weight_exponent)
            _temper = 1.0 / _beta if _beta > 0 else 1.0
            res = _samplers.flowmc_sample(
                like, opts.d_min, opts.d_max, n_chains=max(n_starts, 20),
                n_local_steps=opts.n_local_steps,
                n_global_steps=opts.n_global_steps,
                n_training_loops=opts.n_training_loops,
                n_production_loops=opts.n_production_loops,
                n_epochs=opts.n_epochs,
                n_prior_pilot=opts.n_prior_pilot, temper=_temper,
                seed=opts.seed, reuse_state=flow_state, verbose=opts.verbose)
            out_flow_state = res.get("flow_state")  # bootstrap the next event
        theta, lnL = res["theta"], res["lnL"]
        logZ, sig, neff, ntot = res["logZ"], res["sigma_over_Z"], res["neff"], len(theta)
    elif opts.mode == "prior-mc":
        logZ, sig, neff, ntot, theta, lnL = run_prior_mc(like, opts, rng, dim, with_distance)
    elif opts.mode == "nuts":
        logZ, sig, neff, ntot, theta, lnL = run_nuts(like, opts, rng, with_distance)
    else:
        logZ, sig, neff, ntot, theta, lnL = run_laplace_is(like, opts, rng, dim, with_distance)

    print("\n==== Result (event %d) ====" % event_id)
    print("  log evidence (lnL marginal over extrinsic) = %.5f" % logZ)
    print("  sigma_lnL = %.4g    neff = %.1f    ntotal = %d" % (sig, neff, ntot))
    write_dat(opts, P, out_index, event_id, logZ, sig, ntot, neff)
    write_samples(opts, out_index, theta, lnL, with_distance)
    return logZ, out_flow_state


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main(argv=None):
    optp = build_parser()
    opts, _ = optp.parse_args(argv)
    check_critical_and_report(opts, optp)

    if opts.event_time is None:
        if opts.inj_mode:
            opts.event_time = 1126259462.0
        else:
            optp.error("--event-time is required (frame mode)")
    fiducial_epoch = opts.event_time
    rng = np.random.default_rng(opts.seed)
    deltaT = 1.0 / opts.srate

    # --- data (loaded once; shared across the intrinsic batch) ---
    if opts.inj_mode:
        if opts.mass1 is None or opts.mass2 is None:
            optp.error("--inj-mode requires --mass1 and --mass2")
        P_inj, data_dict, psd_dict, detectors, analyticPSD_Q = load_injection(opts, fiducial_epoch)
        deltaF = data_dict[detectors[0]].deltaF
        P_list = [P_inj]            # injection mode is single-event
    else:
        if not opts.cache_file or not opts.channel_name:
            optp.error("frame mode requires --cache-file and --channel-name (or --inj-mode)")
        data_dict, psd_dict, detectors, analyticPSD_Q = load_frames(opts, fiducial_epoch)
        deltaF = data_dict[detectors[0]].deltaF
        P_list = load_templates(opts, fiducial_epoch, deltaF, deltaT)

    print("Detectors:", detectors)
    print("Analyzing %d event(s) [event index base = %d]" % (len(P_list), opts.event))

    # --- batch loop over intrinsic templates ---
    # For --mode flowmc the trained normalizing flow is threaded from one event
    # to the next (partial-flow re-use), unless --no-flow-reuse is given.
    flow_state = None
    for out_index, P in enumerate(P_list):
        P.deltaT = deltaT
        P.deltaF = deltaF
        event_id = opts.event if (opts.sim_xml or opts.sim_grid or opts.inj_mode) else -1
        print("\n======== event %d (output index %d) : m1=%.3f m2=%.3f ========"
              % (event_id, out_index, P.m1 / MSUN, P.m2 / MSUN))
        try:
            _, new_flow_state = analyze_one(
                opts, P, data_dict, psd_dict, analyticPSD_Q,
                fiducial_epoch, rng, out_index, event_id,
                flow_state=(None if opts.no_flow_reuse else flow_state))
            if not opts.no_flow_reuse and new_flow_state is not None:
                flow_state = new_flow_state
        except Exception as e:
            if getattr(opts, "soft_fail_event_range", False):
                print("  event %d failed (soft-fail): %s" % (event_id, e))
                continue
            raise


if __name__ == "__main__":
    main()
