qra/organized_uncertainty_analysis/bayesian/core.py

# bayesian.py
# ============================================================
# CONDITIONAL BAYESIAN RISK FIT + CONSISTENT CI
# ============================================================

# bayesian_final_clean.py
# ============================================================
# FINAL CONSTRAINED CONDITIONAL RISK MODEL + UNCERTAINTY
# ============================================================
#
# Model fitted by MAP:
#   Y | X=x ~ Bernoulli(P(AE | X=x))
#
# Constraints are enforced through an unconstrained parameter vector phi.
# The risk function P(AE | X=x) is derived from Bayes' rule.
#
# Uncertainty:
#   - local Wald/delta approximation in phi-space
#   - nonparametric pairs bootstrap (NPBS)
#   - generative constrained-Bayesian parametric bootstrap (PBS)
#
# Every clean and bootstrap dataset uses the same conditional MAP estimator.
# Bootstrap refits begin at the clean fit and use seeded fallback starts.
# ============================================================

import os
import time
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

from scipy import optimize, stats
from scipy.ndimage import gaussian_filter1d
from scipy.io import loadmat
from scipy.optimize import brentq
from scipy.special import betaln, gammaln

from ..data import get_data


# ============================================================
# 1) DATA
# ============================================================

def load_xy(
    perc=95,
    suv_path="suv_percentilesSLOthenUWM.mat",
    flags_path="flags_combined.mat",
):
    here = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    suv_full = os.path.join(here, suv_path)
    flags_full = os.path.join(here, flags_path)

    print("Loading SUV from:", suv_full)
    print("Loading FLAGS from:", flags_full)

    suv_dict = loadmat(suv_full)
    flags_dict = loadmat(flags_full)

    X, y = get_data(perc, suv_dict, flags_dict)
    X = np.asarray(X, float).ravel()
    y = np.asarray(y, int).ravel()

    mask = np.isfinite(X)
    X, y = X[mask], y[mask]
    X = np.clip(X, 1e-12, None)

    return X, y


def make_trimmed_dataset(X, y, value_to_drop=2.48122597, tol=1e-3):
    X = np.asarray(X, float)
    y = np.asarray(y, int)

    mask_keep = np.abs(X - value_to_drop) > tol

    return {
        "X_orig": X.copy(),
        "y_orig": y.copy(),
        "X_trim": X[mask_keep],
        "y_trim": y[mask_keep],
        "removed_idx": np.where(~mask_keep)[0],
    }


# ============================================================
# 2) PARAMETERIZATION AND MODEL
# ============================================================

def sigmoid(t):
    return 1.0 / (1.0 + np.exp(-np.clip(t, -60.0, 60.0)))


def softplus(t):
    t = np.asarray(t, float)
    return np.log1p(np.exp(-np.abs(t))) + np.maximum(t, 0.0)


def theta_max(a, b, k, s, eps=1e-12):
    """Monotonicity cap for vartheta."""
    A = a - k
    if A <= 0.0:
        return np.inf

    r = np.sqrt(a + b) - np.sqrt(max(A, eps))
    return np.inf if r <= 1e-12 else s / (r * r)


def unpack(phi):
    """
    Map unconstrained phi to scientific parameters.

    phi = [omega_raw, b_raw, s_raw, k_raw, delta_raw, r_raw]

    omega in (0,1)
    b,s,k > 0
    a = k + delta, delta > 0
    vartheta = vartheta_cap * sigmoid(r_raw)
    """
    omega_raw, b_raw, s_raw, k_raw, delta_raw, r_raw = np.asarray(phi, float)

    omega = sigmoid(omega_raw)

    b = float(softplus(b_raw) + 1e-6)
    s = float(softplus(s_raw) + 1e-6)
    k = float(softplus(k_raw) + 1e-6)

    delta = float(softplus(delta_raw) + 1e-6)
    a = k + delta

    vartheta_cap = theta_max(a, b, k, s)
    vartheta = vartheta_cap * sigmoid(r_raw)

    return omega, a, b, s, k, vartheta, vartheta_cap


def dE_full(x, a, b, s, k, vartheta):
    """log f_AE(x) - log f_NC(x), including normalizing constants."""
    x = np.asarray(x, float)

    return (
        (a - k) * np.log(x)
        - (a + b) * np.log1p(x / s)
        + x / vartheta
        - a * np.log(s)
        - betaln(a, b)
        + k * np.log(vartheta)
        + gammaln(k)
    )


def P_with(theta_hat, x):
    """Derived risk function P(AE | X=x)."""
    omega, a, b, s, k, vartheta, _ = theta_hat
    x = np.asarray(x, float)

    logit_omega = np.log(omega) - np.log1p(-omega)
    eta = logit_omega + dE_full(x, a, b, s, k, vartheta)

    return sigmoid(eta)


def make_priors(y, tau=25.0):
    """Beta(tau*p_emp, tau*(1-p_emp)) prior on prevalence omega."""
    p_emp = float(np.mean(y))
    alpha = max(tau * p_emp, 1e-6)
    beta = max(tau * (1.0 - p_emp), 1e-6)
    return alpha, beta


def init_phi(X, y):
    """Stable generic start in unconstrained coordinates."""
    X = np.asarray(X, float)
    y = np.asarray(y, int)

    X0 = X[y == 0]
    m0 = X0.mean() if X0.size else X.mean()
    v0 = X0.var() if X0.size else X.var()
    k0 = 2.0 if v0 <= 0 else max((m0 * m0) / (v0 + 1e-9), 1.5)

    X1 = X[y == 1]
    m1 = np.median(X1) if X1.size else np.median(X)

    omega0 = np.clip(float(np.mean(y)), 1e-3, 1.0 - 1e-3)
    b0 = 1.5
    s0 = max(m1, 0.5)

    return np.array(
        [
            np.log(omega0 / (1.0 - omega0)),
            np.log(np.expm1(b0) + 1e-9),
            np.log(np.expm1(s0) + 1e-9),
            np.log(np.expm1(k0) + 1e-9),
            np.log(np.expm1(1.0) + 1e-9),
            -0.2,
        ],
        dtype=float,
    )


# ============================================================
# 3) CONDITIONAL NEGATIVE LOG-POSTERIOR
# ============================================================

def neg_post(
    phi,
    X,
    y,
    alpha,
    beta,
    use_prior_p=True,
    prior_r=(1.05, 1.05),
):
    """Conditional Bernoulli negative log-posterior for P(Y=1 | X)."""
    phi = np.asarray(phi, float).reshape(-1)
    X = np.asarray(X, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)

    if (
        phi.size != 6
        or X.size != y.size
        or not np.all(np.isfinite(phi))
        or not np.all(np.isfinite(X))
        or np.any(X <= 0.0)
    ):
        return 1e100

    try:
        omega, a, b, s, k, vartheta, vartheta_cap = unpack(phi)
    except Exception:
        return 1e100

    pars = np.asarray(
        [omega, a, b, s, k, vartheta, vartheta_cap],
        float,
    )

    if not np.all(np.isfinite(pars)):
        return 1e100

    if not (
        0.0 < omega < 1.0
        and a > k > 0.0
        and b > 0.0
        and s > 0.0
        and vartheta > 0.0
        and vartheta_cap > 0.0
        and vartheta < vartheta_cap
    ):
        return 1e100

    eps = 1e-12
    omega_safe = float(np.clip(omega, eps, 1.0 - eps))
    eta = (
        np.log(omega_safe)
        - np.log1p(-omega_safe)
        + dE_full(X, a, b, s, k, vartheta)
    )
    probability = np.clip(sigmoid(eta), eps, 1.0 - eps)
    loglik = float(
        np.sum(y * np.log(probability) + (1 - y) * np.log1p(-probability))
    )

    logprior = 0.0

    if use_prior_p:
        logprior += float(
            (alpha - 1.0) * np.log(omega_safe)
            + (beta - 1.0) * np.log1p(-omega_safe)
        )

    if prior_r is not None:
        r = vartheta / vartheta_cap

        if not np.isfinite(r) or not (0.0 < r < 1.0):
            return 1e100

        r = float(np.clip(r, 1e-12, 1.0 - 1e-12))

        logprior += float(
            (prior_r[0] - 1.0) * np.log(r)
            + (prior_r[1] - 1.0) * np.log1p(-r)
        )

    value = -(loglik + logprior)
    return float(value) if np.isfinite(value) else 1e100


# ============================================================
# 4) MAP FITTING
# ============================================================

def fit_bayes(
    X,
    y,
    seed=0,
    use_prior_p=True,
    prior_r=(1.05, 1.05),
    tau=25.0,
    phi_start=None,
    n_starts=1,
    maxiter=6000,
    ftol=1e-9,
    gtol=1e-8,
    maxls=50,
    refine=False,
    prior_alpha_beta=None,
):
    """Fit the conditional MAP estimator with deterministic multi-starts.

    ``phi_start`` is used first when supplied (notably for bootstrap refits).
    Additional starts are seeded perturbations, and the best converged fit is
    returned. This keeps the estimator consistent across all datasets while
    making the previously exposed fitting arguments effective.
    """
    X = np.asarray(X, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)

    valid = np.isfinite(X) & (X > 0.0)
    X, y = X[valid], y[valid]

    if X.size == 0 or X.size != y.size:
        raise ValueError("Invalid or empty dataset.")
    if np.unique(y).size < 2:
        raise ValueError("Both outcome classes are required.")

    if prior_alpha_beta is None:
        alpha, beta = make_priors(y, tau=tau)
    else:
        alpha, beta = map(float, prior_alpha_beta)

    objective = lambda w: neg_post(
        w,
        X,
        y,
        alpha,
        beta,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
    )

    del refine
    rng = np.random.default_rng(seed)
    default_start = init_phi(X, y)
    first_start = (
        np.asarray(phi_start, float).reshape(-1)
        if phi_start is not None else default_start
    )
    if first_start.size != 6 or not np.all(np.isfinite(first_start)):
        raise ValueError("phi_start must contain six finite values.")

    starts = [first_start]
    if phi_start is not None and int(n_starts) > 1:
        starts.append(default_start)
    while len(starts) < max(1, int(n_starts)):
        starts.append(first_start + rng.normal(0.0, 0.2, size=6))

    results = []
    options = {
        "maxiter": int(maxiter),
        "ftol": float(ftol),
        "gtol": float(gtol),
        "maxls": int(maxls),
    }
    for start in starts:
        candidate = optimize.minimize(
            objective, start, method="L-BFGS-B", options=options
        )
        if candidate.success and np.isfinite(candidate.fun):
            results.append(candidate)

    if not results:
        fallback = first_start + rng.normal(0.0, 0.2, size=6)
        candidate = optimize.minimize(
            objective, fallback, method="L-BFGS-B", options=options
        )
        if candidate.success and np.isfinite(candidate.fun):
            results.append(candidate)

    if not results:
        raise RuntimeError("Conditional MAP fit did not converge.")

    res = min(results, key=lambda item: float(item.fun))
    return unpack(res.x), res


def _fit_bootstrap(
    Xb,
    yb,
    phi_clean,
    seed,
    use_prior_p,
    prior_r,
    tau,
    n_starts_boot=5,
    prior_alpha_beta=None,
):
    """Refit a bootstrap dataset using the clean fit as the first start."""
    return fit_bayes(
        Xb,
        yb,
        seed=seed,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
        phi_start=phi_clean,
        n_starts=n_starts_boot,
        maxiter=6000,
        ftol=1e-9,
        gtol=1e-7,
        maxls=50,
        refine=False,
        prior_alpha_beta=prior_alpha_beta,
    )


# ============================================================
# 5) DERIVED QUANTITIES
# ============================================================

def x_at_p(
    theta_hat,
    p_target=0.5,
    lo=1e-8,
    hi=6.0,
    hi_max=100.0,
    expansion_factor=2.0,
):
    lo = max(float(lo), 1e-12)
    hi = max(float(hi), lo * 1.01)
    hi_max = max(float(hi_max), hi)

    def f(x):
        return float(np.asarray(P_with(theta_hat, x))) - float(p_target)

    try:
        f_lo = f(lo)
        f_hi = f(hi)

        if not (np.isfinite(f_lo) and np.isfinite(f_hi)):
            return np.nan

        while f_lo > 0.0 and lo > 1e-12:
            lo_new = max(lo / expansion_factor, 1e-12)
            if lo_new == lo:
                break
            lo = lo_new
            f_lo = f(lo)

        while f_hi < 0.0 and hi < hi_max:
            hi_new = min(hi * expansion_factor, hi_max)
            if hi_new == hi:
                break
            hi = hi_new
            f_hi = f(hi)

        if f_lo == 0.0:
            return lo
        if f_hi == 0.0:
            return hi
        if f_lo * f_hi > 0.0:
            return np.nan

        return float(brentq(f, lo, hi))

    except Exception:
        return np.nan


def slope_at_x(theta_hat, x0):
    if not np.isfinite(x0) or x0 <= 0.0:
        return np.nan

    h = 1e-4 * max(1.0, abs(float(x0)))
    xm = max(float(x0) - h, 1e-12)
    xp = float(x0) + h

    try:
        pm = float(P_with(theta_hat, xm))
        pp = float(P_with(theta_hat, xp))
        return float((pp - pm) / (xp - xm))
    except Exception:
        return np.nan


# ============================================================
# 6) NUMERICAL DIFFERENTIATION / WALD
# ============================================================

def hess_fd(F, x):
    x = np.asarray(x, float).reshape(-1)
    n = x.size
    H = np.zeros((n, n), float)
    h = 1e-4 * (1.0 + np.abs(x))

    def grad(G, z):
        g = np.full(n, np.nan)
        for j in range(n):
            zp, zm = z.copy(), z.copy()
            zp[j] += h[j]
            zm[j] -= h[j]
            fp, fm = G(zp), G(zm)
            if np.isfinite(fp) and np.isfinite(fm):
                g[j] = (fp - fm) / (2.0 * h[j])
        return g

    for i in range(n):
        xp, xm = x.copy(), x.copy()
        xp[i] += h[i]
        xm[i] -= h[i]
        H[:, i] = (grad(F, xp) - grad(F, xm)) / (2.0 * h[i])

    return 0.5 * (H + H.T)


def jac_fd(Fvec, w):
    w = np.asarray(w, float).reshape(-1)
    f0 = np.asarray(Fvec(w), float).reshape(-1)
    J = np.zeros((f0.size, w.size), float)
    h = 1e-4 * (1.0 + np.abs(w))

    for j in range(w.size):
        wp, wm = w.copy(), w.copy()
        wp[j] += h[j]
        wm[j] -= h[j]

        fp = np.asarray(Fvec(wp), float).reshape(-1)
        fm = np.asarray(Fvec(wm), float).reshape(-1)

        J[:, j] = (fp - fm) / (2.0 * h[j])

    return J


def grad_scalar_fd(F, w):
    w = np.asarray(w, float).reshape(-1)
    g = np.full(w.size, np.nan)
    h = 1e-4 * (1.0 + np.abs(w))

    for j in range(w.size):
        wp, wm = w.copy(), w.copy()
        wp[j] += h[j]
        wm[j] -= h[j]
        fp, fm = F(wp), F(wm)

        if np.isfinite(fp) and np.isfinite(fm):
            g[j] = (fp - fm) / (2.0 * h[j])

    return g


def covariance_from_hessian(
    H,
    max_condition=1e6,
    condition_warn=1e8,
):
    """Positive spectral stabilization followed by covariance inversion.

    Eigenvalues smaller than ``lambda_max / max_condition`` (including
    negative eigenvalues) are replaced by that positive floor. The selected
    condition limit is recorded so regularized Wald/MCA results remain fully
    auditable and can be subjected to a sensitivity analysis.
    """
    H = np.asarray(H, float)
    H = 0.5 * (H + H.T)

    eigvals, eigvecs = np.linalg.eigh(H)

    if not np.all(np.isfinite(eigvals)):
        raise RuntimeError("Non-finite Hessian eigenvalues.")

    max_abs = float(np.max(np.abs(eigvals)))
    if not np.isfinite(max_abs) or max_abs <= 0.0:
        raise RuntimeError("Hessian has no usable curvature.")

    positive = eigvals[eigvals > 0.0]
    original_condition = (
        float(eigvals.max() / positive.min())
        if positive.size == eigvals.size else np.inf
    )
    max_condition = float(max_condition)
    if not np.isfinite(max_condition) or max_condition <= 1.0:
        raise ValueError("max_condition must be a finite number greater than 1.")
    eig_floor = max(max_abs / max_condition, 1e-12)
    eigvals_stable = np.maximum(eigvals, eig_floor)
    stabilized = bool(np.any(eigvals < eig_floor))
    condition = float(eigvals_stable.max() / eigvals_stable.min())

    reliable = bool(
        np.all(eigvals > 0.0)
        and np.isfinite(original_condition)
        and original_condition <= max_condition
        and not stabilized
    )
    if not reliable:
        warnings.warn(
            f"Hessian is weak or indefinite and required spectral "
            f"stabilization (original condition={original_condition:.3e}; "
            f"stabilized condition={condition:.3e}). "
            "Interpret Wald/MCA as local sensitivity approximations.",
            RuntimeWarning,
        )

    Sigma = eigvecs @ np.diag(1.0 / eigvals_stable) @ eigvecs.T
    Sigma = 0.5 * (Sigma + Sigma.T)

    diagnostics = {
        "eigenvalues": eigvals,
        "minimum_eigenvalue": float(eigvals.min()),
        "maximum_eigenvalue": float(eigvals.max()),
        "condition_number": condition,
        "original_condition_number": original_condition,
        "positive_definite": bool(np.all(eigvals > 0.0)),
        "original_positive_definite": bool(np.all(eigvals > 0.0)),
        "stabilized": stabilized,
        "eigenvalue_floor": eig_floor,
        "stabilized_eigenvalues": eigvals_stable,
        "reliable": reliable,
        "maximum_allowed_condition": max_condition,
    }

    return Sigma, diagnostics


def reported_theta_from_phi(phi):
    omega, a, b, s, k, vartheta, _ = unpack(phi)
    return np.array([omega, a, b, s, k, vartheta], float)


def reported_theta_from_hat(theta_hat):
    omega, a, b, s, k, vartheta, _ = theta_hat
    return np.array([omega, a, b, s, k, vartheta], float)


def _wald_scalar(estimate, gradient, covariance, z):
    gradient = np.asarray(gradient, float)

    if not np.isfinite(estimate) or not np.all(np.isfinite(gradient)):
        return np.nan, np.nan

    var = float(gradient @ covariance @ gradient)

    if not np.isfinite(var) or var < -1e-12:
        return np.nan, np.nan

    se = np.sqrt(max(var, 0.0))
    return float(estimate - z * se), float(estimate + z * se)


# ============================================================
# 7) CI FOR ONE DATASET
# ============================================================

def estimate_ci_bundle(
    X,
    y,
    label,
    x_grid,
    B_nonpar=400,
    B_param=400,
    M_mca=10000,
    mca_seed=None,
    seed=123,
    use_prior_p=True,
    prior_r=(1.05, 1.05),
    tau=25.0,
    alpha_ci=0.05,
    n_starts_clean=1,
    n_starts_boot=1,
    progress_every=25,
    max_attempt_multiplier=20,
    max_hessian_condition=1e6,
):
    X = np.asarray(X, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)
    x_grid = np.asarray(x_grid, float).reshape(-1)

    rng = np.random.default_rng(seed)
    n = X.size

    z = float(stats.norm.ppf(1.0 - alpha_ci / 2.0))
    qlo, qhi = 100.0 * alpha_ci / 2.0, 100.0 * (1.0 - alpha_ci / 2.0)

    # ---------------- CLEAN MAP ----------------
    t0 = time.time()
    theta_hat, res = fit_bayes(
        X,
        y,
        seed=seed,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
        n_starts=n_starts_clean,
        refine=False,
    )

    phi_hat = np.asarray(res.x, float)
    theta_hat_vec = reported_theta_from_phi(phi_hat)

    pmap = np.asarray(P_with(theta_hat, x_grid), float)
    # Freeze the empirical-Bayes prior calibrated on the original dataset.
    # Bootstrap samples vary through resampling/simulation, not by redefining
    # the prior from each replicate's random event count.
    prior_alpha_beta = make_priors(y, tau=tau)
    x50_hat = x_at_p(theta_hat, 0.5, hi=max(6.0, float(x_grid.max())))
    s50_hat = slope_at_x(theta_hat, x50_hat)

    print(f"{label}: clean MAP completed in {(time.time()-t0):.1f} s")

    # ---------------- WALD ----------------
    alpha_p, beta_p = make_priors(y, tau=tau)

    objective = lambda w: neg_post(
        w, X, y, alpha_p, beta_p,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
    )

    H = hess_fd(objective, phi_hat)

    wald_available = True
    try:
        Sigma_phi, hdiag = covariance_from_hessian(
            H, max_condition=max_hessian_condition
        )

        J_curve = jac_fd(lambda w: P_with(unpack(w), x_grid), phi_hat)
        var_curve = np.einsum("ij,jk,ik->i", J_curve, Sigma_phi, J_curve)
        se_curve = np.sqrt(np.maximum(var_curve, 0.0))
        wald_lo = np.clip(pmap - z * se_curve, 0.0, 1.0)
        wald_hi = np.clip(pmap + z * se_curve, 0.0, 1.0)

        J_theta = jac_fd(reported_theta_from_phi, phi_hat)
        Sigma_theta = J_theta @ Sigma_phi @ J_theta.T
        Sigma_theta = 0.5 * (Sigma_theta + Sigma_theta.T)

        se_theta = np.sqrt(np.maximum(np.diag(Sigma_theta), 0.0))
        wald_param_lo = theta_hat_vec - z * se_theta
        wald_param_hi = theta_hat_vec + z * se_theta

        def x50_phi(w):
            return x_at_p(
                unpack(w),
                0.5,
                hi=max(6.0, float(x_grid.max())),
            )

        gx = grad_scalar_fd(x50_phi, phi_hat)
        wald_x50_lo, wald_x50_hi = _wald_scalar(
            x50_hat, gx, Sigma_phi, z
        )

        def s50_phi(w):
            th = unpack(w)
            xx = x_at_p(
                th,
                0.5,
                hi=max(6.0, float(x_grid.max())),
            )
            return slope_at_x(th, xx)

        gs = grad_scalar_fd(s50_phi, phi_hat)
        wald_s50_lo, wald_s50_hi = _wald_scalar(
            s50_hat, gs, Sigma_phi, z
        )

    except RuntimeError as exc:
        warnings.warn(f"{label}: Wald unavailable: {exc}", RuntimeWarning)
        wald_available = False
        Sigma_phi = np.full((6, 6), np.nan)
        Sigma_theta = np.full((6, 6), np.nan)
        hdiag = {}
        wald_lo = wald_hi = np.full_like(pmap, np.nan)
        wald_param_lo = wald_param_hi = np.full(6, np.nan)
        wald_x50_lo = wald_x50_hi = np.nan
        wald_s50_lo = wald_s50_hi = np.nan

    # ---------------- MONTE CARLO APPROXIMATION ----------------
    # Laplace/normal propagation using the same covariance as Wald. The
    # diagnostics state explicitly when historical spectral stabilization was
    # required before inversion.
    mca_curves, mca_theta, mca_x50, mca_s50 = [], [], [], []
    mca_attempted = int(M_mca)
    if mca_seed is None:
        mca_seed = int(seed) + 10_000

    if wald_available and int(M_mca) > 0:
        rng_mca = np.random.default_rng(mca_seed)
        try:
            phi_draws = rng_mca.multivariate_normal(
                mean=phi_hat,
                cov=Sigma_phi,
                size=int(M_mca),
                check_valid="raise",
            )
            for draw in np.atleast_2d(phi_draws):
                try:
                    thm = unpack(draw)
                    curve = np.asarray(P_with(thm, x_grid), float)
                    xx = x_at_p(
                        thm, 0.5, hi=max(6.0, float(x_grid.max()))
                    )
                    ss = slope_at_x(thm, xx)
                    theta_vec = reported_theta_from_hat(thm)
                    if (
                        np.all(np.isfinite(curve))
                        and np.all(np.isfinite(theta_vec))
                        and np.isfinite(xx)
                        and np.isfinite(ss)
                    ):
                        mca_curves.append(curve)
                        mca_theta.append(theta_vec)
                        mca_x50.append(xx)
                        mca_s50.append(ss)
                except Exception:
                    continue
        except Exception as exc:
            warnings.warn(f"{label}: MCA unavailable: {exc}", RuntimeWarning)

    mca_curves = np.asarray(mca_curves, float)
    mca_theta = np.asarray(mca_theta, float)
    mca_x50 = np.asarray(mca_x50, float)
    mca_s50 = np.asarray(mca_s50, float)
    mca_lo = (
        np.percentile(mca_curves, qlo, axis=0) if len(mca_curves) else None
    )
    mca_hi = (
        np.percentile(mca_curves, qhi, axis=0) if len(mca_curves) else None
    )
    diagnostics_mca = {
        "scheme": "local_Gaussian_phi_with_spectral_stabilization",
        "target": int(M_mca),
        "attempted": mca_attempted,
        "successful_curve": len(mca_curves),
        "successful_x50": len(mca_x50),
        "successful_s50": len(mca_s50),
        "available": bool(len(mca_curves) > 0),
        "reason_if_unavailable": (
            None if len(mca_curves) else
            "No usable covariance or no finite Monte Carlo draws"
        ),
    }

    # ---------------- NPBS ----------------
    np_curves, np_theta, np_x50, np_s50 = [], [], [], []
    attempted_np = rejected_single_np = rejected_fit_np = 0
    t_np = time.time()

    max_attempts_np = max(int(B_nonpar), int(max_attempt_multiplier * B_nonpar))
    while len(np_curves) < int(B_nonpar) and attempted_np < max_attempts_np:
        attempted_np += 1

        idx = rng.integers(0, n, size=n)
        Xb, yb = X[idx], y[idx]

        if np.unique(yb).size < 2:
            rejected_single_np += 1
            continue

        try:
            thb, rb = _fit_bootstrap(
                Xb, yb, phi_hat,
                int(rng.integers(0, 10_000_000)),
                use_prior_p, prior_r, tau,
                n_starts_boot=n_starts_boot,
                prior_alpha_beta=prior_alpha_beta,
            )

            curve = np.asarray(P_with(thb, x_grid), float)
            xx = x_at_p(thb, 0.5, hi=max(6.0, float(x_grid.max())))
            ss = slope_at_x(thb, xx)

            if not (
                np.all(np.isfinite(curve))
                and np.isfinite(xx)
                and np.isfinite(ss)
            ):
                rejected_fit_np += 1
                continue

            np_curves.append(curve)
            np_theta.append(reported_theta_from_hat(thb))
            np_x50.append(xx)
            np_s50.append(ss)

            if progress_every and len(np_curves) % progress_every == 0:
                elapsed = (time.time() - t_np) / 60.0
                print(
                    f"{label} NPBS: {len(np_curves)}/{B_nonpar} "
                    f"({elapsed:.1f} min)"
                )

        except Exception:
            rejected_fit_np += 1

    # ---------------- PBS ----------------
    pb_curves, pb_theta, pb_x50, pb_s50 = [], [], [], []
    attempted_pb = rejected_single_pb = rejected_fit_pb = 0
    omega_hat, a_hat, b_hat, s_hat, k_hat, vartheta_hat, _ = theta_hat
    t_pb = time.time()

    max_attempts_pb = max(int(B_param), int(max_attempt_multiplier * B_param))
    while len(pb_curves) < int(B_param) and attempted_pb < max_attempts_pb:
        attempted_pb += 1

        # Generative CB-PBS from the fitted prevalence and class-conditional
        # biomarker distributions, as specified in the methodology.
        yb = rng.binomial(1, omega_hat, size=n).astype(int)

        if np.unique(yb).size < 2:
            rejected_single_pb += 1
            continue

        Xb = np.empty(n, float)
        is_nc = yb == 0
        is_ae = ~is_nc
        Xb[is_nc] = rng.gamma(
            shape=k_hat, scale=vartheta_hat, size=int(is_nc.sum())
        )
        numerator = rng.gamma(shape=a_hat, scale=1.0, size=int(is_ae.sum()))
        denominator = rng.gamma(shape=b_hat, scale=1.0, size=int(is_ae.sum()))
        Xb[is_ae] = s_hat * numerator / denominator

        try:
            thb, rb = _fit_bootstrap(
                Xb, yb, phi_hat,
                int(rng.integers(0, 10_000_000)),
                use_prior_p, prior_r, tau,
                n_starts_boot=n_starts_boot,
                prior_alpha_beta=prior_alpha_beta,
            )

            curve = np.asarray(P_with(thb, x_grid), float)
            xx = x_at_p(thb, 0.5, hi=max(6.0, float(x_grid.max())))
            ss = slope_at_x(thb, xx)

            if not (
                np.all(np.isfinite(curve))
                and np.isfinite(xx)
                and np.isfinite(ss)
            ):
                rejected_fit_pb += 1
                continue

            pb_curves.append(curve)
            pb_theta.append(reported_theta_from_hat(thb))
            pb_x50.append(xx)
            pb_s50.append(ss)

            if progress_every and len(pb_curves) % progress_every == 0:
                elapsed = (time.time() - t_pb) / 60.0
                print(
                    f"{label} PBS: {len(pb_curves)}/{B_param} "
                    f"({elapsed:.1f} min)"
                )

        except Exception:
            rejected_fit_pb += 1

    np_curves = np.asarray(np_curves, float)
    np_theta = np.asarray(np_theta, float)
    np_x50 = np.asarray(np_x50, float)
    np_s50 = np.asarray(np_s50, float)

    pb_curves = np.asarray(pb_curves, float)
    pb_theta = np.asarray(pb_theta, float)
    pb_x50 = np.asarray(pb_x50, float)
    pb_s50 = np.asarray(pb_s50, float)

    np_lo = np.percentile(np_curves, qlo, axis=0) if len(np_curves) else None
    np_hi = np.percentile(np_curves, qhi, axis=0) if len(np_curves) else None

    pb_lo = np.percentile(pb_curves, qlo, axis=0) if len(pb_curves) else None
    pb_hi = np.percentile(pb_curves, qhi, axis=0) if len(pb_curves) else None

    diagnostics_np = {
        "scheme": "ordinary_pairs",
        "target": B_nonpar,
        "attempted": attempted_np,
        "successful_curve": len(np_curves),
        "successful_x50": len(np_x50),
        "successful_s50": len(np_s50),
        "rejected_single_class": rejected_single_np,
        "rejected_fit": rejected_fit_np,
        "completed_target": bool(len(np_curves) == int(B_nonpar)),
        "max_attempts": max_attempts_np,
    }

    diagnostics_pb = {
        "scheme": "generative_CB_prevalence_and_class_conditional_X",
        "target": B_param,
        "attempted": attempted_pb,
        "successful_curve": len(pb_curves),
        "successful_x50": len(pb_x50),
        "successful_s50": len(pb_s50),
        "rejected_single_class": rejected_single_pb,
        "rejected_fit": rejected_fit_pb,
        "completed_target": bool(len(pb_curves) == int(B_param)),
        "max_attempts": max_attempts_pb,
    }

    return {
        "label": label,
        "theta_hat": theta_hat,
        "theta_hat_vec": theta_hat_vec,
        "phi_hat": phi_hat,
        "res": res,
        "objective": float(res.fun),

        "x50": x50_hat,
        "s50": s50_hat,
        "pmap": pmap,

        "wald_available": wald_available,
        "wald_lo": wald_lo,
        "wald_hi": wald_hi,
        "wald_param_lo": wald_param_lo,
        "wald_param_hi": wald_param_hi,
        "wald_x50_lo": wald_x50_lo,
        "wald_x50_hi": wald_x50_hi,
        "wald_s50_lo": wald_s50_lo,
        "wald_s50_hi": wald_s50_hi,

        "np_lo": np_lo,
        "np_hi": np_hi,
        "pb_lo": pb_lo,
        "pb_hi": pb_hi,
        "mca_lo": mca_lo,
        "mca_hi": mca_hi,

        "theta_np": np_theta,
        "theta_pb": pb_theta,
        "x50_np": np_x50,
        "x50_pb": pb_x50,
        "s50_np": np_s50,
        "s50_pb": pb_s50,
        "theta_mca": mca_theta,
        "x50_mca": mca_x50,
        "s50_mca": mca_s50,

        "used_np": len(np_curves),
        "used_pb": len(pb_curves),
        "used_mca": len(mca_curves),

        "diagnostics_np": diagnostics_np,
        "diagnostics_pb": diagnostics_pb,
        "diagnostics_mca": diagnostics_mca,
        "hessian_diagnostics": hdiag,
        "prior_alpha_beta": prior_alpha_beta,
        "alpha_ci": float(alpha_ci),
        "max_hessian_condition": float(max_hessian_condition),

        "H_phi": H,
        "Sigma_phi": Sigma_phi,
        "Sigma_theta": Sigma_theta,
    }


# ============================================================
# 8) FULL + TRIM CI
# ============================================================

def run_bayesian_ci(
    perc=95,
    suv_path="suv_percentilesSLOthenUWM.mat",
    flags_path="flags_combined.mat",
    value_to_drop=2.48122597,
    tol=1e-3,
    xmax=4.5,
    n_grid=600,
    B_nonpar=400,
    B_param=400,
    M_mca=10000,
    mca_seed=None,
    seed=123,
    use_prior_p=True,
    prior_r=(1.05, 1.05),
    tau=25.0,
    alpha_ci=0.05,
    grid_lower_fraction=0.75,
    n_starts_clean=1,
    n_starts_boot=1,
    progress_every=25,
    max_attempt_multiplier=20,
    max_hessian_condition=1e6,
):
    X_all, y_all = load_xy(
        perc=perc,
        suv_path=suv_path,
        flags_path=flags_path,
    )

    ds = make_trimmed_dataset(
        X_all,
        y_all,
        value_to_drop=value_to_drop,
        tol=tol,
    )

    x_min = max(
        1e-8,
        grid_lower_fraction
        * float(min(ds["X_orig"].min(), ds["X_trim"].min())),
    )

    x_max = max(
        float(xmax),
        float(ds["X_orig"].max()),
        float(ds["X_trim"].max()),
    )

    x_grid = np.linspace(x_min, x_max, int(n_grid))

    full = estimate_ci_bundle(
        ds["X_orig"],
        ds["y_orig"],
        "FULL",
        x_grid,
        B_nonpar=B_nonpar,
        B_param=B_param,
        M_mca=M_mca,
        mca_seed=(int(seed) + 10_000 if mca_seed is None else int(mca_seed)),
        seed=seed,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
        alpha_ci=alpha_ci,
        n_starts_clean=n_starts_clean,
        n_starts_boot=n_starts_boot,
        progress_every=progress_every,
        max_attempt_multiplier=max_attempt_multiplier,
        max_hessian_condition=max_hessian_condition,
    )

    trim = estimate_ci_bundle(
        ds["X_trim"],
        ds["y_trim"],
        "TRIM",
        x_grid,
        B_nonpar=B_nonpar,
        B_param=B_param,
        M_mca=M_mca,
        mca_seed=(int(seed) + 10_001 if mca_seed is None else int(mca_seed) + 1),
        seed=seed + 1,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
        alpha_ci=alpha_ci,
        n_starts_clean=n_starts_clean,
        n_starts_boot=n_starts_boot,
        progress_every=progress_every,
        max_attempt_multiplier=max_attempt_multiplier,
        max_hessian_condition=max_hessian_condition,
    )

    return {
        **ds,
        "x_grid": x_grid,
        "orig": full,
        "trim": trim,
    }


# ============================================================
# 9) TABLES / DIAGNOSTICS
# ============================================================

PARAMETER_NAMES = ["omega", "a", "b", "s", "k", "vartheta"]


HISTORICAL_REFERENCE = {
    "FULL": {
        "objective": 9.612766,
        "theta": np.array([
            0.05013031639340312,
            835.7394119971459,
            812.0599466844761,
            1.5849829783592195,
            260.69110931079496,
            0.005737139709007281,
        ]),
        "x50": 1.7923,
    },
    "TRIM": {
        "objective": 7.423491,
        "theta": np.array([
            0.050915239967693926,
            547.8229343628883,
            297.1652388267381,
            0.9319626391393511,
            116.9852193858385,
            0.012177886089862443,
        ]),
        "x50": 1.7541,
    },
}


def check_historical_reference(ci_res, rtol=5e-3, atol=5e-4):
    """Compare the clean fit with the saved approved notebook results."""
    rows = []
    for key, dataset in (("orig", "FULL"), ("trim", "TRIM")):
        out = ci_res[key]
        ref = HISTORICAL_REFERENCE[dataset]
        theta = np.asarray(out["theta_hat_vec"], float)
        rows.append({
            "Dataset": dataset,
            "Objective": float(out["objective"]),
            "Reference_objective": ref["objective"],
            "x50": float(out["x50"]),
            "Reference_x50": ref["x50"],
            "Objective_match": bool(np.isclose(
                out["objective"], ref["objective"], rtol=rtol, atol=atol
            )),
            "Parameters_match": bool(np.allclose(
                theta, ref["theta"], rtol=rtol, atol=atol
            )),
            "x50_match": bool(np.isclose(
                out["x50"], ref["x50"], rtol=rtol, atol=atol
            )),
        })
    return pd.DataFrame(rows)


def print_ci_diagnostics(ci_res):
    for key, name in (("orig", "FULL"), ("trim", "TRIM")):
        out = ci_res[key]

        print(f"\n{name}")
        print("-" * len(name))
        print(f"Conditional negative log-posterior: {out['objective']:.10g}")

        hd = out["hessian_diagnostics"]
        if hd:
            print("Original Hessian positive definite:", hd["positive_definite"])
            print("Hessian minimum eigenvalue:", hd["minimum_eigenvalue"])
            print("Hessian condition number:", hd["condition_number"])
            print(
                "Original Hessian condition number:",
                hd["original_condition_number"],
            )
            print("Hessian spectrally stabilized:", hd["stabilized"])
            print("Wald/MCA covariance reliable:", hd["reliable"])
            print("Eigenvalue floor:", hd["eigenvalue_floor"])
            print(
                "Maximum allowed stabilized condition:",
                hd["maximum_allowed_condition"],
            )

        print(
            "MAP parameters:",
            dict(zip(PARAMETER_NAMES, out["theta_hat_vec"])),
        )
        print(f"x50={out['x50']:.8g}, s50={out['s50']:.8g}")

        print("NPBS:", out["diagnostics_np"])
        print("PBS:", out["diagnostics_pb"])
        print("MCA:", out["diagnostics_mca"])


def _percentile_or_nan(values, q):
    values = np.asarray(values, float)
    return float(np.percentile(values, q)) if values.size else np.nan


def _ci_percentiles(out):
    alpha = float(out.get("alpha_ci", 0.05))
    return 100.0 * alpha / 2.0, 100.0 * (1.0 - alpha / 2.0)


def _column_or_empty(values, j):
    values = np.asarray(values, float)
    return values[:, j] if values.ndim == 2 and values.shape[1] > j else []


def make_parameter_ci_table(ci_res):
    rows = []

    for key, dataset in (("orig", "FULL"), ("trim", "TRIM")):
        out = ci_res[key]
        estimate = np.asarray(out["theta_hat_vec"], float)
        qlo, qhi = _ci_percentiles(out)

        for j, name in enumerate(PARAMETER_NAMES):
            rows.append({
                "Dataset": dataset,
                "Parameter": name,
                "Estimate": estimate[j],
                "Wald_LL": out["wald_param_lo"][j],
                "Wald_UL": out["wald_param_hi"][j],
                "NPBS_LL": _percentile_or_nan(_column_or_empty(out["theta_np"], j), qlo),
                "NPBS_UL": _percentile_or_nan(_column_or_empty(out["theta_np"], j), qhi),
                "PBS_LL": _percentile_or_nan(_column_or_empty(out["theta_pb"], j), qlo),
                "PBS_UL": _percentile_or_nan(_column_or_empty(out["theta_pb"], j), qhi),
                "MCA_LL": _percentile_or_nan(
                    out["theta_mca"][:, j] if out["theta_mca"].ndim == 2
                    else [], qlo
                ),
                "MCA_UL": _percentile_or_nan(
                    out["theta_mca"][:, j] if out["theta_mca"].ndim == 2
                    else [], qhi
                ),
            })

    return pd.DataFrame(rows)


def make_derived_ci_table(ci_res):
    rows = []

    for key, dataset in (("orig", "FULL"), ("trim", "TRIM")):
        out = ci_res[key]
        qlo, qhi = _ci_percentiles(out)

        for char, est, wlo, whi, npv, pbv in [
            (
                "x50",
                out["x50"],
                out["wald_x50_lo"],
                out["wald_x50_hi"],
                out["x50_np"],
                out["x50_pb"],
            ),
            (
                "s50",
                out["s50"],
                out["wald_s50_lo"],
                out["wald_s50_hi"],
                out["s50_np"],
                out["s50_pb"],
            ),
        ]:
            rows.append({
                "Dataset": dataset,
                "Characteristic": char,
                "Estimate": est,
                "Wald_LL": wlo,
                "Wald_UL": whi,
                "NPBS_LL": _percentile_or_nan(npv, qlo),
                "NPBS_UL": _percentile_or_nan(npv, qhi),
                "PBS_LL": _percentile_or_nan(pbv, qlo),
                "PBS_UL": _percentile_or_nan(pbv, qhi),
                "MCA_LL": _percentile_or_nan(
                    out[f"{char}_mca"], qlo
                ),
                "MCA_UL": _percentile_or_nan(
                    out[f"{char}_mca"], qhi
                ),
            })

    return pd.DataFrame(rows)

def smooth_ci_bounds(lower, upper, sigma=0.0):
    """Optionally smooth CI boundaries for presentation only.

    The empirical confidence limits stored in ``ci_res`` are not modified.
    Numerical summaries and tables therefore continue to use the original
    unsmoothed bootstrap distributions and confidence limits.
    """
    lower = np.asarray(lower, dtype=float)
    upper = np.asarray(upper, dtype=float)

    if lower.shape != upper.shape:
        raise ValueError("Lower and upper CI boundaries must have equal shape.")
    if lower.ndim != 1:
        raise ValueError("CI boundaries must be one-dimensional arrays.")
    if sigma <= 0:
        return lower.copy(), upper.copy()

    lower_smooth = gaussian_filter1d(lower, sigma=sigma, mode="nearest")
    upper_smooth = gaussian_filter1d(upper, sigma=sigma, mode="nearest")

    lower_smooth = np.clip(lower_smooth, 0.0, 1.0)
    upper_smooth = np.clip(upper_smooth, 0.0, 1.0)

    # Maintain a valid ordered confidence band after numerical smoothing.
    lower_final = np.minimum(lower_smooth, upper_smooth)
    upper_final = np.maximum(lower_smooth, upper_smooth)

    return lower_final, upper_final


def plot_bayesian_ci(
    ci_res,
    xmax=4.5,
    figsize=(10, 4),
    dpi=140,
    smooth_sigma=0.0,
    jitter=0.018,
    jitter_seed=123,
    band_support="observed",
):
    """Plot Bayesian risk functions and confidence bands.

    Empirical boundaries are shown without smoothing by default. They are
    smoothed for visualization only when
    ``smooth_sigma`` is positive. The original empirical boundaries in
    ``ci_res`` remain unchanged and continue to support all calculations.
    Set ``smooth_sigma=0`` to display the original unsmoothed boundaries.
    With ``band_support='observed'`` (default), all bands are displayed only
    between the smallest and largest observed biomarker values in each panel;
    use ``band_support='grid'`` to display bands over the entire model grid.
    """
    fig, axes = plt.subplots(
        1,
        2,
        figsize=figsize,
        dpi=dpi,
        sharey=True,
    )
    alpha_plot = float(ci_res["orig"].get("alpha_ci", 0.05))
    ci_level = 100.0 * (1.0 - alpha_plot)
    ci_text = f"{ci_level:g}%"

    # Colours corresponding to the previous figure
    colors = {
        "wald": "#138A24",   # green
        "npbs": "#173BFF",   # blue
        "pbs":  "#00CFE3",   # cyan
        "mca":  "#B23AEE",   # purple
        "nc":   "#7479FF",   # periwinkle
        "ae":   "#FFBE63",   # orange
    }

    rng = np.random.default_rng(jitter_seed)

    for ax, key, x_key, y_key, panel_label in zip(
        axes,
        ("orig", "trim"),
        ("X_orig", "X_trim"),
        ("y_orig", "y_trim"),
        ("A", "B"),
    ):
        out = ci_res[key]
        xg = ci_res["x_grid"]

        # Observed samples with small vertical jitter to prevent overlap.
        x_obs = np.asarray(ci_res[x_key])
        y_obs = np.asarray(ci_res[y_key])
        if band_support == "observed":
            band_mask = (xg >= float(x_obs.min())) & (xg <= float(x_obs.max()))
        elif band_support == "grid":
            band_mask = np.ones(xg.shape, dtype=bool)
        else:
            raise ValueError("band_support must be 'observed' or 'grid'.")
        x_band = xg[band_mask]
        y_jittered = y_obs + rng.uniform(-jitter, jitter, size=y_obs.size)
        ax.scatter(
            x_obs[y_obs == 0],
            y_jittered[y_obs == 0],
            s=22,
            color=colors["nc"],
            alpha=0.72,
            edgecolors="none",
            zorder=7,
        )
        ax.scatter(
            x_obs[y_obs == 1],
            y_jittered[y_obs == 1],
            s=25,
            color=colors["ae"],
            alpha=0.90,
            edgecolors="none",
            zorder=8,
        )

        # Fitted risk function
        ax.plot(
            xg,
            out["pmap"],
            color="black",
            lw=2.2,
            label="Fitted risk function",
            zorder=5,
        )

        # Wald confidence band
        if out["wald_available"]:
            ax.fill_between(
                x_band,
                out["wald_lo"][band_mask],
                out["wald_hi"][band_mask],
                color=colors["wald"],
                alpha=0.16,
                label=f"Wald {ci_text} CI",
                zorder=1,
            )
            ax.plot(
                x_band,
                out["wald_lo"][band_mask],
                color=colors["wald"],
                ls="-.",
                lw=1.4,
                zorder=3,
            )
            ax.plot(
                x_band,
                out["wald_hi"][band_mask],
                color=colors["wald"],
                ls="-.",
                lw=1.4,
                zorder=3,
            )

        # Nonparametric pairs bootstrap band
        if out["np_lo"] is not None:
            np_lo_plot, np_hi_plot = smooth_ci_bounds(
                out["np_lo"],
                out["np_hi"],
                sigma=smooth_sigma,
            )
            ax.fill_between(
                x_band,
                np_lo_plot[band_mask],
                np_hi_plot[band_mask],
                color=colors["npbs"],
                alpha=0.14,
                label=f"NPBS {ci_text} CI",
                zorder=1,
            )
            ax.plot(
                x_band,
                np_lo_plot[band_mask],
                color=colors["npbs"],
                ls=":",
                lw=1.4,
                zorder=3,
            )
            ax.plot(
                x_band,
                np_hi_plot[band_mask],
                color=colors["npbs"],
                ls=":",
                lw=1.4,
                zorder=3,
            )

        # Parametric bootstrap band
        if out["pb_lo"] is not None:
            pb_lo_plot, pb_hi_plot = smooth_ci_bounds(
                out["pb_lo"],
                out["pb_hi"],
                sigma=smooth_sigma,
            )
            ax.fill_between(
                x_band,
                pb_lo_plot[band_mask],
                pb_hi_plot[band_mask],
                color=colors["pbs"],
                alpha=0.14,
                label=f"PBS {ci_text} CI",
                zorder=1,
            )
            ax.plot(
                x_band,
                pb_lo_plot[band_mask],
                color=colors["pbs"],
                ls="--",
                lw=1.6,
                zorder=3,
            )
            ax.plot(
                x_band,
                pb_hi_plot[band_mask],
                color=colors["pbs"],
                ls="--",
                lw=1.6,
                zorder=3,
            )

        # Strict local-Gaussian Monte Carlo approximation
        if out["mca_lo"] is not None:
            mca_lo_plot, mca_hi_plot = smooth_ci_bounds(
                out["mca_lo"], out["mca_hi"], sigma=smooth_sigma
            )
            ax.fill_between(
                x_band, mca_lo_plot[band_mask], mca_hi_plot[band_mask],
                color=colors["mca"], alpha=0.10, zorder=1,
            )
            ax.plot(
                x_band, mca_lo_plot[band_mask], color=colors["mca"],
                ls=(0, (5, 2, 1, 2)), lw=1.5, zorder=3,
            )
            ax.plot(
                x_band, mca_hi_plot[band_mask], color=colors["mca"],
                ls=(0, (5, 2, 1, 2)), lw=1.5, zorder=3,
            )

        ax.set_xlim(xg.min(), xmax)
        ax.set_ylim(-0.05, 1.05)
        ax.set_xlabel(r"$X$")
        ax.grid(False)
        ax.text(
            0.045,
            0.955,
            panel_label,
            transform=ax.transAxes,
            ha="left",
            va="top",
            fontsize=14,
            fontweight="normal",
            zorder=10,
        )

    axes[0].set_ylabel(r"$P(\mathrm{AE}\mid X=x)$")

    # Complete legend inside the lower-right corner of panel B.
    legend_handles = [
        Line2D(
            [0], [0],
            marker="o",
            linestyle="none",
            markerfacecolor=colors["nc"],
            markeredgecolor="none",
            markersize=7,
            label="data: NC",
        ),
        Line2D(
            [0], [0],
            marker="o",
            linestyle="none",
            markerfacecolor=colors["ae"],
            markeredgecolor="none",
            markersize=7,
            label="data: AE",
        ),
        Line2D(
            [0], [0],
            color="black",
            ls="-",
            lw=2.2,
            label="fit",
        ),
    ]

    if any(ci_res[key]["wald_available"] for key in ("orig", "trim")):
        legend_handles.append(Line2D(
            [0], [0], color=colors["wald"], ls="-.", lw=1.8,
            label=f"CI: Wald {ci_text}",
        ))
    if any(ci_res[key]["np_lo"] is not None for key in ("orig", "trim")):
        legend_handles.append(Line2D(
            [0], [0], color=colors["npbs"], ls=":", lw=1.8,
            label=f"CI: NPBS {ci_text}",
        ))
    if any(ci_res[key]["pb_lo"] is not None for key in ("orig", "trim")):
        legend_handles.append(Line2D(
            [0], [0], color=colors["pbs"], ls="--", lw=2.0,
            label=f"CI: PBS {ci_text}",
        ))
    if any(ci_res[key]["mca_lo"] is not None for key in ("orig", "trim")):
        legend_handles.append(Line2D(
            [0], [0], color=colors["mca"],
            ls=(0, (5, 2, 1, 2)), lw=1.8,
            label=f"CI: MCA {ci_text}",
        ))

    axes[1].legend(
        handles=legend_handles,
        loc="lower right",
        frameon=True,
        framealpha=0.95,
        fontsize=9,
    )

    plt.tight_layout()

    return fig, axes

# ============================================================