load_chains.py

"""Per-probe chain/posterior loaders.

Each loader returns a `ChainStats` (mean + covariance over the params it
constrains, on the standard ΛCDM-6 sub-block). For probes where we have
full chain files we compute mean/cov directly; for probes where we only
have a published summary we synthesize a Gaussian posterior consistent
with the published mean ± σ ± correlation.

Conventions:
- All probe loaders return parameters using the canonical names from
  src.extract_fisher.PARAM_ORDER.
- Loaders raise FileNotFoundError if their input files are missing,
  so the orchestrator can decide whether to fall back to a literature
  summary or skip the probe.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from .extract_fisher import PARAM_ORDER, ChainStats, compute_chain_stats

ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"


# ---------------------------------------------------------------------------
#  CosmoMC chain reader (Planck format)
# ---------------------------------------------------------------------------


def parse_paramnames(path: Path) -> list[str]:
    """Read a CosmoMC .paramnames file. Returns parameter names in order.

    Format is tab-separated: `<name>\\t<latex_label>`, one parameter per line.
    Derived parameters end in '*'; the star is stripped here so callers can
    look up by clean name.
    """
    names = []
    with open(path) as f:
        for line in f:
            line = line.rstrip("\n")
            if not line.strip():
                continue
            # Tab-separated; the first field is the parameter name.
            parts = line.split("\t")
            name = parts[0].strip()
            if name.endswith("*"):
                name = name[:-1]
            names.append(name)
    return names


# CosmoMC name -> our PARAM_ORDER name
COSMOMC_TO_STANDARD = {
    "omegabh2": "Omega_b_h2",
    "omegam": "Omega_m",
    "H0": "H0",
    "sigma8": "sigma_8",
    "ns": "n_s",
    "tau": "tau",
}


def load_cosmomc_chain(
    chain_dir: Path,
    chain_root: str,
    *,
    burnin_frac: float = 0.3,
    max_samples_per_chain: int | None = None,
) -> tuple[np.ndarray, np.ndarray, list[str]]:
    """Load CosmoMC chain files matching `<chain_root>_<i>.txt`.

    Returns (samples, weights, names) where samples has columns in
    PARAM_ORDER for any standard parameters present, and `names` lists
    those that were found.

    burnin_frac : fraction of samples per chain to discard at the start.
    """
    paramnames_file = chain_dir / f"{chain_root}.paramnames"
    if not paramnames_file.exists():
        raise FileNotFoundError(paramnames_file)

    raw_names = parse_paramnames(paramnames_file)
    # Find columns we care about. CosmoMC chain layout:
    #   col 0 = weight, col 1 = -logL, col 2+ = parameters in raw_names order.
    standard_columns = []
    found_names = []
    for cosmomc_name, standard_name in COSMOMC_TO_STANDARD.items():
        if cosmomc_name in raw_names:
            idx = raw_names.index(cosmomc_name)
            standard_columns.append(2 + idx)
            found_names.append(standard_name)
        # else: probe doesn't constrain this parameter, skip silently.

    # Order found_names according to PARAM_ORDER
    order_idx = sorted(range(len(found_names)),
                       key=lambda i: PARAM_ORDER.index(found_names[i]))
    standard_columns = [standard_columns[i] for i in order_idx]
    found_names = [found_names[i] for i in order_idx]

    # Find the primary chain files (numbered _1.txt, _2.txt, ...; skip _post_*).
    chain_files = sorted(
        p for p in chain_dir.glob(f"{chain_root}_*.txt")
        if "_post_" not in p.name
    )
    if not chain_files:
        raise FileNotFoundError(f"no chains matching {chain_root}_*.txt in {chain_dir}")

    samples_chunks = []
    weights_chunks = []
    for cf in chain_files:
        arr = np.loadtxt(cf)
        if arr.ndim == 1:
            arr = arr[None, :]
        n = arr.shape[0]
        burn = int(n * burnin_frac)
        arr = arr[burn:]
        if max_samples_per_chain is not None and arr.shape[0] > max_samples_per_chain:
            # Subsample uniformly; preserves marginal but may distort eff. sample size
            idx = np.linspace(0, arr.shape[0] - 1, max_samples_per_chain, dtype=int)
            arr = arr[idx]
        weights_chunks.append(arr[:, 0])
        samples_chunks.append(arr[:, standard_columns])

    weights = np.concatenate(weights_chunks)
    samples = np.concatenate(samples_chunks)
    return samples, weights, found_names


# ---------------------------------------------------------------------------
#  Gaussian-summary loader (for probes available only as mean ± σ + corr)
# ---------------------------------------------------------------------------


def load_gaussian_summary(
    summary_path: Path,
    name: str,
    *,
    n_samples: int = 100_000,
    seed: int = 0,
) -> ChainStats:
    """Build a ChainStats from a JSON summary file.

    The JSON must contain:
      - "constrained": list of parameter names (subset of PARAM_ORDER)
      - "mean": dict name -> value
      - "sigma": dict name -> value
      - optional "correlation": full correlation matrix (matching constrained order)

    We synthesize a Gaussian sample to populate ChainStats.cov via the
    standard pipeline; the result's mean & cov match the literature input
    exactly (because we compute them from the synthesized samples that
    were drawn from a Gaussian with that mean and cov).
    """
    if not summary_path.exists():
        raise FileNotFoundError(summary_path)
    summary = json.loads(summary_path.read_text())
    constrained = summary["constrained"]
    means = np.array([summary["mean"][p] for p in constrained])
    sigmas = np.array([summary["sigma"][p] for p in constrained])
    corr = summary.get("correlation")
    if corr is None:
        corr = np.eye(len(constrained))
    else:
        corr = np.asarray(corr)
        if corr.shape != (len(constrained), len(constrained)):
            raise ValueError(f"correlation shape mismatch in {summary_path}")
    cov = np.outer(sigmas, sigmas) * corr

    # Synthesize Gaussian samples and run them through compute_chain_stats so
    # the output uses the same code path as real chains.
    rng = np.random.default_rng(seed)
    samples = rng.multivariate_normal(means, cov, size=n_samples)
    weights = np.ones(n_samples)
    stats = compute_chain_stats(name, samples, constrained, weights)
    stats.notes = f"Gaussian summary from {summary_path.name}: source={summary.get('source', '?')}"
    return stats


# ---------------------------------------------------------------------------
#  Per-probe loaders
# ---------------------------------------------------------------------------


def load_planck() -> ChainStats:
    chain_dir = DATA / "planck2018_chains" / "base" / "plikHM_TTTEEE_lowl_lowE"
    chain_root = "base_plikHM_TTTEEE_lowl_lowE"
    samples, weights, names = load_cosmomc_chain(
        chain_dir, chain_root,
        burnin_frac=0.3,
        max_samples_per_chain=None,
    )
    stats = compute_chain_stats("planck", samples, names, weights)
    stats.notes = (
        f"Planck 2018 PR3 base ΛCDM TT,TE,EE+lowE chain "
        f"({len(np.unique(weights))} weight values, {samples.shape[0]} samples after 30% burn-in)"
    )
    return stats


def load_pantheon() -> ChainStats:
    """Pantheon+ posterior on Omega_m (literature summary).

    Pantheon+ alone constrains Omega_m well; H0 is degenerate with M0 nuisance
    in the SNe-only fit and is not reported here. To include H0 information
    one would use the Pantheon+SH0ES joint analysis (which would double-count
    SH0ES — explicitly out of scope for this battery per pre-reg §1).
    """
    return load_gaussian_summary(
        DATA / "pantheon_plus" / "summary_lcdm.json",
        "pantheon",
    )


def load_desi_bao() -> ChainStats:
    """DESI DR1 BAO + BBN posterior on (H0, Omega_m)."""
    return load_gaussian_summary(
        DATA / "desi_bao_dr1" / "summary_lcdm.json",
        "desi_bao",
    )


def load_kids() -> ChainStats:
    """KiDS-1000 COSEBIs cosmic-shear MultiNest chain on (sigma_8, Omega_m).

    Loaded directly from the public KiDS-1000 cosmic-shear data release
    (``data/kids1000_cosebis/.../output_multinest_C.txt``). The previous
    Gaussian-summary path is preserved as ``load_kids_synth_legacy`` for
    audit comparability; it is deprecated and should not be used in
    production batteries.
    """
    chain_path = (
        DATA / "kids1000" / "KiDS1000_cosmis_shear_data_release"
        / "chains_and_config_files" / "main_chains_iterative_covariance"
        / "cosebis" / "chain" / "output_multinest_C.txt"
    )
    if not chain_path.exists():
        return load_gaussian_summary(DATA / "kids1000" / "summary_lcdm.json", "kids")
    header_cols = None
    with open(chain_path) as f:
        for line in f:
            if line.startswith("#"):
                low = line[1:].strip()
                if "\t" in low and header_cols is None:
                    header_cols = [c.strip().lower() for c in low.split("\t")]
                continue
            break
    if header_cols is None:
        return load_gaussian_summary(DATA / "kids1000" / "summary_lcdm.json", "kids")
    arr = np.loadtxt(chain_path, comments="#")
    s8_idx = header_cols.index("cosmological_parameters--sigma_8")
    om_idx = header_cols.index("cosmological_parameters--omega_m")
    w_idx = header_cols.index("weight")
    samples = np.column_stack([arr[:, om_idx], arr[:, s8_idx]])
    weights = arr[:, w_idx]
    constrained = ["Omega_m", "sigma_8"]
    stats = compute_chain_stats("kids", samples, constrained, weights=weights)
    stats.notes = (
        f"KiDS-1000 COSEBIs MultiNest cosmic-shear chain "
        f"(public data release), {samples.shape[0]} samples, "
        f"n_eff={stats.n_eff:,.0f}"
    )
    return stats


def load_kids_synth_legacy() -> ChainStats:
    """DEPRECATED: pre-audit synthesized KiDS Gaussian summary.

    The synthesis JSON encoded ``corr(sigma_8, Omega_m) = +0.7`` (from a
    misread of the published ``corr(S_8, Omega_m)`` table) and marginal
    sigmas 1.8–3.2x smaller than the real chain. Preserved for audit
    comparability only; do not use in production batteries.
    """
    return load_gaussian_summary(DATA / "kids1000" / "summary_lcdm.json", "kids_synth_legacy")


def load_shoes() -> ChainStats:
    """SH0ES local distance-ladder H0 posterior (1-D)."""
    return load_gaussian_summary(
        DATA / "shoes_h0" / "summary_lcdm.json",
        "shoes",
    )


# ---------------------------------------------------------------------------
#  Driver
# ---------------------------------------------------------------------------


def load_planck_lensing() -> ChainStats:
    return load_gaussian_summary(DATA / "planck_lensing" / "summary_lcdm.json", "planck_lensing")


def load_pantheon_shoes_chain() -> ChainStats:
    """Pantheon+SH0ES FlatLambdaCDM REAL MCMC chain (official release).

    Source: github.com/PantheonPlusSH0ES/DataRelease,
    Pantheon+_Data/5_COSMOLOGY/chains/Pantheon+SH0ES/Pantheon+SH0ES_FlatLambdaCDM.txt
    (CosmoSIS/emcee; columns: omega_m, h0 [little h], m, prior, post; unweighted).
    Cite: Brout et al. 2022 (Pantheon+), Riess et al. 2022 (SH0ES).

    X29-P3 upgrade of the Gaussian-summary loader ``load_pantheon_shoes``:
    same probe, chain-derived (H0, Omega_m). h0 is dimensionless and is
    converted to H0 in km/s/Mpc (x100).
    """
    f = DATA / "pantheon_plus_chains" / "Pantheon+SH0ES_FlatLambdaCDM.txt"
    a = np.loadtxt(f, comments="#")
    samples = np.column_stack([100.0 * a[:, 1], a[:, 0]])   # (H0, Omega_m)
    stats = compute_chain_stats("pantheon_shoes", samples, ["H0", "Omega_m"])
    stats.notes = (f"Pantheon+SH0ES FlatLambdaCDM official chain, "
                   f"{samples.shape[0]} samples (emcee, unweighted)")
    return stats


def load_desi_bao_chain() -> ChainStats:
    """DESI DR1 BAO-all + Schoeneberg+24 BBN REAL MCMC chain (official VAC).

    Source: data.desi.lbl.gov/public/dr1/vac/dr1/bao-cosmo-params/v1.0/
    cobaya/base/desi-bao-all_schoneberg2024-bbn/chain.{1-4}.txt
    (Cobaya MH; H0 sampled freely with the BBN ombh2 prior).
    Cite: DESI Collaboration 2024 (2404.03002).

    X29-P4 upgrade of the Gaussian-summary loader ``load_desi_bao``: same
    probe, chain-derived (H0, Omega_m), 30% per-chain burn-in, multiplicity
    weights. NB: the older top-level data/desi_y1_chains/chain.*.txt are a
    DIFFERENT fixed-rdrag variant in which H0 is pinned; do not use them here.
    """
    chain_dir = DATA / "desi_y1_chains" / "desi-bao-all_bbn"
    chunks, ws = [], []
    for i in range(1, 5):
        hdr = (chain_dir / f"chain.{i}.txt").open().readline().lstrip("#").split()
        cols = {n: j for j, n in enumerate(hdr)}
        a = np.loadtxt(chain_dir / f"chain.{i}.txt", comments="#")
        a = a[int(0.3 * len(a)):]
        chunks.append(np.column_stack([a[:, cols["H0"]], a[:, cols["omegam"]]]))
        ws.append(a[:, cols["weight"]])
    samples, w = np.vstack(chunks), np.concatenate(ws)
    stats = compute_chain_stats("desi_bao", samples, ["H0", "Omega_m"], weights=w)
    stats.notes = (f"DESI DR1 BAO-all + BBN official Cobaya chain, "
                   f"{samples.shape[0]} samples post 30% burn-in")
    return stats


def load_planck_lensing_chain() -> ChainStats:
    """Planck 2018 PR3 lensing-only (base_lensing_lenspriors, R3.01) REAL MCMC chain.

    Used by the orientation battery (v68) for the decisive CMB-lensing placement in
    (Omega_m, sigma_8). The earlier ``load_planck_lensing`` Gaussian summary encoded a
    wrong-sign Omega_m-sigma_8 correlation; the real chain gives the correct strongly
    negative correlation (CMB lensing groups with cosmic shear, not primary CMB). See
    paper/preregs/cmb_lensing_chain_orientation_*. The summary loader is retained only
    for the other batteries' non-headline descriptive use.
    """
    chain_dir = DATA / "cmb_lensing_chains" / "planck2018_lensonly"
    chain_root = "base_lensing_lenspriors"
    samples, weights, names = load_cosmomc_chain(
        chain_dir, chain_root, burnin_frac=0.3, max_samples_per_chain=None,
    )
    stats = compute_chain_stats("planck_lensing", samples, names, weights)
    stats.notes = (
        "Planck 2018 PR3 lensing-only (base_lensing_lenspriors, R3.01) real chain; "
        f"{samples.shape[0]} samples after 30% burn-in"
    )
    return stats


def load_act_dr6() -> ChainStats:
    """ACT DR6 actbase ΛCDM chain (Cobaya format, 4 walkers, 30% burn-in).

    Loaded directly from the public ACT-DR6 likelihood-pipeline chains in
    ``data/act_dr6_chains/actbase_lcdm_camb/``. Column indices (0-based):
    weight=0; sampled ombh2=2, omch2=3, ns=6, tau=7; derived H0=39,
    sigma8=41, omegam=51 (verified against V25/V43 means).
    """
    chain_dir = DATA / "act_dr6_chains" / "actbase_lcdm_camb" / "actbase_lcdm_camb"
    files = sorted(chain_dir.glob("actbase_lcdm_camb.[1-4].txt"))
    if not files:
        # Fall back to published Gaussian summary if chain missing.
        return load_gaussian_summary(DATA / "act_dr6" / "summary_lcdm.json", "act_dr6")
    sample_chunks, weight_chunks = [], []
    for f in files:
        a = np.loadtxt(f, comments="#")
        n = a.shape[0]
        burn = int(0.3 * n)
        a = a[burn:]
        # cols: H0, Omega_m, Omega_b_h2, sigma_8, n_s, tau (PARAM_ORDER)
        sample_chunks.append(np.column_stack([
            a[:, 39], a[:, 51], a[:, 2], a[:, 41], a[:, 6], a[:, 7],
        ]))
        weight_chunks.append(a[:, 0])
    samples = np.vstack(sample_chunks)
    weights = np.concatenate(weight_chunks)
    stats = compute_chain_stats(
        "act_dr6", samples, list(PARAM_ORDER), weights=weights,
    )
    stats.notes = (
        "ACT DR6 actbase ΛCDM chain (Cobaya), 4 walkers, 30% burn-in, "
        f"{samples.shape[0]} post-burn samples, n_eff={stats.n_eff:,.0f}"
    )
    return stats


def load_spt3g() -> ChainStats:
    return load_gaussian_summary(DATA / "spt3g" / "summary_lcdm.json", "spt3g")


def load_pantheon_shoes() -> ChainStats:
    return load_gaussian_summary(DATA / "pantheon_shoes" / "summary_lcdm.json", "pantheon_shoes")


def load_eboss_bao_rsd() -> ChainStats:
    return load_gaussian_summary(DATA / "eboss_bao_rsd" / "summary_lcdm.json", "eboss_bao_rsd")


def load_hsc_y3() -> ChainStats:
    return load_gaussian_summary(DATA / "hsc_y3" / "summary_lcdm.json", "hsc_y3")


def load_des_y3_synth_legacy() -> ChainStats:
    """DEPRECATED: pre-audit synthesized DES Y3 Gaussian summary.

    The synthesis JSON encoded ``corr(sigma_8, Omega_m) = +0.6`` (from a
    misread of the published ``corr(S_8, Omega_m)`` table) and a marginal
    sigma assignment with sigma(sigma_8) and sigma(Omega_m) inverted
    relative to the chain. Preserved for audit comparability only.
    """
    return load_gaussian_summary(DATA / "des_y3" / "summary_lcdm.json", "des_y3_synth_legacy")


def load_des_y3() -> ChainStats:
    """DES Y3 3x2pt LCDM chain (CosmoSIS PolyChord, public release).

    Loaded directly from ``data/des_y3/chain_3x2pt_lcdm_SR_maglim.txt``.
    Constrained 5-D sub-block: (H0, Omega_m, Omega_b_h2, sigma_8, n_s).
    DES does not constrain tau. Column indices (0-based) per the chain
    header: Omega_m=0, h=1 (H0=h*100), Omega_b=2 (Omega_b_h2=Omega_b*h^2),
    n_s=3, sigma_8=31, weight=37.
    """
    chain_path = DATA / "des_y3" / "chain_3x2pt_lcdm_SR_maglim.txt"
    if not chain_path.exists():
        return load_gaussian_summary(DATA / "des_y3" / "summary_lcdm.json", "des_y3")
    a = np.loadtxt(chain_path, comments="#")
    # Burn 30% (PolyChord weights already account for nested-sampling structure;
    # we follow the same convention as other chain loaders for consistency).
    n = a.shape[0]
    burn = int(0.3 * n)
    a = a[burn:]
    h = a[:, 1]
    H0 = h * 100.0
    Omega_m = a[:, 0]
    Omega_b_h2 = a[:, 2] * (h ** 2)
    n_s = a[:, 3]
    sigma_8 = a[:, 31]
    weights = a[:, 37]
    samples = np.column_stack([H0, Omega_m, Omega_b_h2, sigma_8, n_s])
    constrained = ["H0", "Omega_m", "Omega_b_h2", "sigma_8", "n_s"]
    stats = compute_chain_stats("des_y3", samples, constrained, weights=weights)
    stats.notes = (
        f"DES Y3 3x2pt LCDM PolyChord chain (public, SR maglim), "
        f"{samples.shape[0]} post-burn samples, n_eff={stats.n_eff:,.0f}; "
        f"tau not sampled by DES."
    )
    return stats


def load_cosmic_chronometers() -> ChainStats:
    return load_gaussian_summary(DATA / "cosmic_chronometers" / "summary_lcdm.json", "cosmic_chronometers")


def load_gw_sirens() -> ChainStats:
    return load_gaussian_summary(DATA / "gw_sirens" / "summary_lcdm.json", "gw_sirens")


def load_clusters() -> ChainStats:
    return load_gaussian_summary(DATA / "clusters" / "summary_lcdm.json", "clusters")


# V3 held-out probes
def load_wmap9() -> ChainStats:
    return load_gaussian_summary(DATA / "wmap9" / "summary_lcdm.json", "wmap9")


def load_eboss_lya_bao() -> ChainStats:
    return load_gaussian_summary(DATA / "eboss_lya_bao" / "summary_lcdm.json", "eboss_lya_bao")


def load_h0licow_tdcosmo() -> ChainStats:
    return load_gaussian_summary(DATA / "h0licow_tdcosmo" / "summary_lcdm.json", "h0licow_tdcosmo")


def load_spt_sz() -> ChainStats:
    return load_gaussian_summary(DATA / "spt_sz" / "summary_lcdm.json", "spt_sz")


def load_erosita() -> ChainStats:
    return load_gaussian_summary(DATA / "erosita" / "summary_lcdm.json", "erosita")


# V3.5 multi-release CMB-primary
def load_planck_pr2() -> ChainStats:
    return load_gaussian_summary(DATA / "planck_pr2" / "summary_lcdm.json", "planck_pr2")


def load_planck_pr4_primary() -> ChainStats:
    return load_gaussian_summary(DATA / "planck_pr4_primary" / "summary_lcdm.json", "planck_pr4_primary")


def load_act_dr4() -> ChainStats:
    return load_gaussian_summary(DATA / "act_dr4" / "summary_lcdm.json", "act_dr4")


# V7-D blind predictive test
def load_euclid_q1() -> ChainStats:
    return load_gaussian_summary(DATA / "euclid_q1" / "summary_lcdm.json", "euclid_q1")


# V4 age/redshift extension probes
def load_bbn() -> ChainStats:
    return load_gaussian_summary(DATA / "bbn" / "summary_lcdm.json", "bbn")


def load_gc_ages_valcin() -> ChainStats:
    return load_gaussian_summary(DATA / "gc_ages_valcin" / "summary_lcdm.json", "gc_ages_valcin")


def load_wd_cooling_ngc6791() -> ChainStats:
    return load_gaussian_summary(DATA / "wd_cooling_ngc6791" / "summary_lcdm.json", "wd_cooling_ngc6791")


def load_jwst_high_z() -> ChainStats:
    return load_gaussian_summary(DATA / "jwst_high_z" / "summary_lcdm.json", "jwst_high_z")


def load_stellar_pop_halo() -> ChainStats:
    return load_gaussian_summary(DATA / "stellar_pop_halo" / "summary_lcdm.json", "stellar_pop_halo")


# E4b H_0 expansion: chain-derived (H_0, Omega_m) probes added for V60-extended.
def load_h0licow() -> ChainStats:
    """H0LiCOW Wong+2020 6-lens FLCDM joint chain.

    Source: github.com/shsuyu/H0LiCOW-public/cosmo_parameter_chains/FLCDM
    Constrained: (H_0, Omega_m). 232k unweighted samples.
    Independent of CMB, BAO, SN.
    """
    chain_path = DATA / "h0licow_chains" / "FLCDM_joint.dat"
    if not chain_path.exists():
        return load_gaussian_summary(
            DATA / "h0licow_tdcosmo" / "summary_lcdm.json", "h0licow",
        )
    samples = np.loadtxt(chain_path, delimiter=",", comments="#")
    constrained = ["H0", "Omega_m"]
    stats = compute_chain_stats("h0licow", samples, constrained)
    stats.notes = (
        f"H0LiCOW Wong+2020 6-lens FLCDM joint chain, "
        f"{samples.shape[0]} samples (unweighted)"
    )
    return stats


def load_tdcosmo2025() -> ChainStats:
    """TDCOSMO 2025 ULambdaCDM1: time-delay alone (RXJ1131).

    Source: github.com/TDCOSMO/TDCOSMO2025_public ULambdaCDM1.h5
    arXiv:2506.03023; dataset = "TDCOSMO" (no SN, no SLACS, no CMB).
    Cosmology subset: (h0, om) -> (H_0, Omega_m). 500k samples.
    Other 6 params (lambda_mst*, alpha_lambda, gamma_pl_0, a_ani*) are
    nuisance and dropped from the cosmology subspace.
    """
    import h5py
    chain_path = DATA / "tdcosmo2025_chains" / "ULambdaCDM1.h5"
    if not chain_path.exists():
        return load_gaussian_summary(
            DATA / "h0licow_tdcosmo" / "summary_lcdm.json", "tdcosmo2025",
        )
    with h5py.File(chain_path, "r") as f:
        s = f["samples"][:]
        params = [p.decode() for p in f["parameters"][:]]
    h0_idx = params.index("h0")
    om_idx = params.index("om")
    samples = np.column_stack([s[:, h0_idx], s[:, om_idx]])
    constrained = ["H0", "Omega_m"]
    stats = compute_chain_stats("tdcosmo2025", samples, constrained)
    stats.notes = (
        f"TDCOSMO 2025 ULambdaCDM1 (TDCOSMO time-delay alone, RXJ1131), "
        f"{samples.shape[0]} samples (unweighted), nuisance params dropped"
    )
    return stats


PROBE_LOADERS = {
    # Original 5
    "planck": load_planck,
    "pantheon": load_pantheon,
    "desi_bao": load_desi_bao,
    "kids": load_kids,
    "shoes": load_shoes,
    # Extended 10
    "planck_lensing": load_planck_lensing,
    "act_dr6": load_act_dr6,
    "spt3g": load_spt3g,
    "pantheon_shoes": load_pantheon_shoes,
    "eboss_bao_rsd": load_eboss_bao_rsd,
    "hsc_y3": load_hsc_y3,
    "des_y3": load_des_y3,
    "cosmic_chronometers": load_cosmic_chronometers,
    "gw_sirens": load_gw_sirens,
    "clusters": load_clusters,
    # V3 held-out
    "wmap9": load_wmap9,
    "eboss_lya_bao": load_eboss_lya_bao,
    "h0licow_tdcosmo": load_h0licow_tdcosmo,
    "spt_sz": load_spt_sz,
    "erosita": load_erosita,
    # V3.5 multi-release CMB-primary
    "planck_pr2": load_planck_pr2,
    "planck_pr4_primary": load_planck_pr4_primary,
    "act_dr4": load_act_dr4,
    # V7-D blind predictive test
    "euclid_q1": load_euclid_q1,
    # V4 age/redshift extension
    "bbn": load_bbn,
    "gc_ages_valcin": load_gc_ages_valcin,
    "wd_cooling_ngc6791": load_wd_cooling_ngc6791,
    "jwst_high_z": load_jwst_high_z,
    "stellar_pop_halo": load_stellar_pop_halo,
    # E4b H_0 expansion (chain-derived (H_0, Omega_m) probes)
    "h0licow": load_h0licow,
    "tdcosmo2025": load_tdcosmo2025,
}


def load_all() -> dict[str, ChainStats]:
    out = {}
    for name, fn in PROBE_LOADERS.items():
        try:
            stats = fn()
            out[name] = stats
            print(
                f"[{name}] {stats.n_samples} samples, "
                f"params={stats.constrained_params}, "
                f"H0={'%.2f' % stats.mean[stats.constrained_params.index('H0')] if 'H0' in stats.constrained_params else '—'}"
            )
        except FileNotFoundError as e:
            print(f"[{name}] missing data: {e}")
    return out