# bayesian.py
# ============================================================
# BAYESIAN FIT + CI + ELASTICITY
# ============================================================

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

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

from data_utils import get_data


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

def load_xy(
    perc=95,
    suv_path="suv_percentilesSLOthenUWM.mat",
    flags_path="flags_combined.mat",
):
    """
    Load feature x and binary label y.
    """
    here = 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()

    m = np.isfinite(x)
    x, y = x[m], y[m]

    x = np.clip(x, 1e-12, None)
    return x, y


# ============================================================
# 2) CORE MODEL FUNCTIONS
# ============================================================

def logistic(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -60, 60)))


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


def dE_full(x, a, b, s, k, th):
    """
    log f_BP(x|a,b,s) - log f_Gamma(x|k,th), including constants.
    """
    x = np.asarray(x, float)
    return (
        (a - k) * np.log(x)
        - (a + b) * np.log1p(x / s)
        + x / th
        - a * np.log(s)
        - betaln(a, b)
        + k * np.log(th)
        + gammaln(k)
    )


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 theta.
    """
    A = a - k
    if A <= 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):
    """
    Reparameterisation:
      phi = [p_raw, b_raw, s_raw, k_raw, d_raw, u_raw]

      p in (0,1)
      b,s,k > 0
      a = k + delta with delta > 0
      theta = theta_cap * sigmoid(u_raw)
    """
    p_raw, b_raw, s_raw, k_raw, d_raw, u_raw = phi

    p = sigmoid(p_raw)

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

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

    thcap = theta_max(a, b, k, s)
    th = thcap * sigmoid(u_raw)

    return p, a, b, s, k, th, thcap


def make_priors(y, tau=25.0):
    """
    Beta(TAU*p_emp, TAU*(1-p_emp)) prior on prevalence p.
    """
    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 neg_post(phi, X, y, alpha, beta, use_prior_p=True, prior_r=(1.05, 1.05)):
    """
    Negative log-posterior = NLL + optional priors.

    Priors used here:
      - Beta prior on prevalence p
      - Beta prior on r = theta/theta_cap

    No extra priors on a, b, s, k.
    """
    p, a, b, s, k, th, thcap = unpack(phi)
    eps = 1e-12

    logit_val = (np.log(p) - np.log(1.0 - p)) + dE_full(X, a, b, s, k, th)
    px = logistic(logit_val)

    nll = -np.sum(y * np.log(px + eps) + (1 - y) * np.log(1 - px + eps))

    if use_prior_p:
        nll += -((alpha - 1) * np.log(p + eps) + (beta - 1) * np.log(1 - p + eps))

    if prior_r is not None and np.isfinite(thcap) and thcap > 0:
        r = np.clip(th / thcap, 1e-9, 1 - 1e-9)
        nll += -((prior_r[0] - 1) * np.log(r) + (prior_r[1] - 1) * np.log(1 - r))

    return float(nll)


def init_phi(X, y):
    """
    Stable initial values.
    """
    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)

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

    return np.array(
        [
            np.log(p0 / (1 - p0)),         # p_raw
            np.log(np.expm1(b0) + 1e-9),   # b_raw
            np.log(np.expm1(s0) + 1e-9),   # s_raw
            np.log(np.expm1(k0) + 1e-9),   # k_raw
            np.log(np.expm1(1.0) + 1e-9),  # d_raw
            -0.2,                          # u_raw
        ],
        dtype=float,
    )


def fit_bayes(X, y, seed=0, use_prior_p=True, prior_r=(1.05, 1.05), tau=25.0):
    """
    MAP fit using L-BFGS-B.
    """
    X = np.asarray(X, float)
    y = np.asarray(y, int)

    alpha, beta = make_priors(y, tau=tau)

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

    w0 = init_phi(X, y)
    res = optimize.minimize(
        obj,
        w0,
        method="L-BFGS-B",
        options={"maxiter": 6000, "ftol": 1e-9},
    )

    if not (res.success and np.isfinite(res.fun)):
        rng = np.random.default_rng(seed)
        w1 = w0 + rng.normal(0, 0.2, size=w0.shape)
        res = optimize.minimize(
            obj,
            w1,
            method="L-BFGS-B",
            options={"maxiter": 6000, "ftol": 1e-9},
        )

    theta_hat = unpack(res.x)
    return theta_hat, res


def P_with(theta_hat, x):
    """
    Posterior risk curve P(AE|x) under fitted model.
    """
    p, a, b, s, k, th, _ = theta_hat
    x = np.asarray(x, float)
    logit_val = (np.log(p) - np.log(1 - p)) + dE_full(x, a, b, s, k, th)
    return logistic(logit_val)


# ============================================================
# 3) ORIGINAL / TRIM DATASETS
# ============================================================

def make_trimmed_dataset(X, y, value_to_drop=2.48122597, tol=1e-3):
    """
    Remove point(s) with x approximately equal to value_to_drop.
    """
    X = np.asarray(X, float)
    y = np.asarray(y, int)

    mask_keep = np.abs(X - value_to_drop) > tol
    removed_idx = np.where(~mask_keep)[0]

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


def run_bayesian_group_fit(
    perc=95,
    suv_path="suv_percentilesSLOthenUWM.mat",
    flags_path="flags_combined.mat",
    value_to_drop=2.48122597,
    tol=1e-3,
    use_prior_p=True,
    prior_r=(1.05, 1.05),
    tau=25.0,
):
    """
    Load data, create ORIGINAL/TRIM datasets,
    and fit constrained Bayesian group model on both.
    """
    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)

    theta_orig, res_orig = fit_bayes(
        ds["X_orig"],
        ds["y_orig"],
        seed=0,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
    )

    theta_trim, res_trim = fit_bayes(
        ds["X_trim"],
        ds["y_trim"],
        seed=1,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
    )

    return {
        **ds,
        "theta_orig": theta_orig,
        "theta_trim": theta_trim,
        "res_orig": res_orig,
        "res_trim": res_trim,
    }


def summarize_theta(theta_hat):
    p, a, b, s, k, th, thcap = theta_hat
    return {
        "p": p,
        "a": a,
        "b": b,
        "s": s,
        "k": k,
        "theta": th,
        "theta_cap": thcap,
        "r": th / thcap if np.isfinite(thcap) and thcap > 0 else np.nan,
    }


def plot_bayesian_orig_trim_raw(
    fit_results,
    xmax=6.0,
    suptitle="Conditional Probability of AE",
    figsize=(12, 7),
    dpi=140,
):
    """
    One raw-x plot:
      - ORIGINAL curve
      - TRIM curve
      - ORIGINAL data dots
      - highlight removed point(s)
    """
    X_orig = fit_results["X_orig"]
    y_orig = fit_results["y_orig"]
    removed_idx = fit_results["removed_idx"]
    theta_orig = fit_results["theta_orig"]
    theta_trim = fit_results["theta_trim"]

    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    fig.text(0.02, 0.5, suptitle, va="center", rotation="vertical", fontsize=14)

    x_min = max(float(np.min(X_orig)), 1e-8)
    x_max = float(xmax)

    ax.set_xlim(x_min, x_max)
    ax.set_xlabel("x")
    ax.set_ylabel("P(AE | x)")
    ax.set_ylim(-0.10, 1.10)
    ax.grid(alpha=0.35)

    x_grid = np.exp(np.linspace(np.log(x_min), np.log(x_max), 900))
    p_curve_orig = P_with(theta_orig, x_grid)
    p_curve_trim = P_with(theta_trim, x_grid)

    l1, = ax.plot(x_grid, p_curve_orig, lw=2.2, color="C0", label="ORIGINAL (Bayesian fit)")
    l2, = ax.plot(x_grid, p_curve_trim, lw=2.2, color="C1", label="TRIM (Bayesian fit)")

    rng = np.random.default_rng(999)
    jit = (rng.random(len(y_orig)) - 0.5) * 0.06

    d_nc = ax.scatter(
        X_orig[y_orig == 0],
        (y_orig + jit)[y_orig == 0],
        s=22,
        alpha=0.65,
        edgecolors="none",
        color="C0",
        label="NC samples (ORIGINAL)",
    )
    d_ae = ax.scatter(
        X_orig[y_orig == 1],
        (y_orig + jit)[y_orig == 1],
        s=26,
        alpha=0.85,
        edgecolors="none",
        color="C1",
        label="AE samples (ORIGINAL)",
    )

    dout = None
    if removed_idx.size > 0:
        for j, i in enumerate(removed_idx):
            jit_out = (rng.random() - 0.5) * 0.06
            label = "Removed point" if j == 0 else None
            dout = ax.scatter(
                [float(X_orig[i])],
                [float(y_orig[i] + jit_out)],
                marker="x",
                s=90,
                linewidths=2,
                color="k",
                label=label,
            )

    handles = [l1, l2, d_nc, d_ae]
    if dout is not None:
        handles.append(dout)
    labels = [h.get_label() for h in handles]
    ax.legend(handles, labels, frameon=False, ncol=2, loc="lower right")

    plt.tight_layout(rect=(0.06, 0.0, 1.0, 1.0))
    return fig, ax


# ============================================================
# 4) CI ESTIMATION
# ============================================================

def x_at_p(theta_hat, p_target=0.5, lo=1e-6, hi=10.0):
    """
    Solve P(AE|x) = p_target for x.
    """
    f = lambda x: P_with(theta_hat, x) - p_target
    try:
        if f(lo) * f(hi) > 0:
            return np.nan
        return float(brentq(f, lo, hi))
    except Exception:
        return np.nan


def slope_at_x(theta_hat, x0):
    """
    Numerical derivative of P(AE|x) at x0.
    """
    if not np.isfinite(x0):
        return np.nan
    h = 1e-3 * (1 + abs(x0))
    return float((P_with(theta_hat, x0 + h) - P_with(theta_hat, x0 - h)) / (2 * h))


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

    def grad_fd(G, z):
        g = np.zeros_like(z)
        for j in range(n):
            ej = np.zeros_like(z)
            ej[j] = 1.0
            g[j] = (G(z + h[j] * ej) - G(z - h[j] * ej)) / (2 * h[j])
        return g

    for i in range(n):
        ei = np.zeros_like(x)
        ei[i] = 1.0
        g_plus = grad_fd(F, x + h[i] * ei)
        g_minus = grad_fd(F, x - h[i] * ei)
        H[:, i] = (g_plus - g_minus) / (2 * h[i])

    return 0.5 * (H + H.T)


def jac_fd(Fvec, w):
    """
    Finite-difference Jacobian for vector-valued function.
    """
    f0 = Fvec(w)
    m = f0.size
    n = w.size
    J = np.zeros((m, n))
    h = 1e-4 * (1 + np.abs(w))

    for j in range(n):
        ej = np.zeros_like(w)
        ej[j] = 1.0
        J[:, j] = (Fvec(w + h[j] * ej) - Fvec(w - h[j] * ej)) / (2 * h[j])

    return J


def estimate_ci_bundle(
    X,
    y,
    label,
    x_grid,
    B_nonpar=400,
    B_param=400,
    seed=123,
    use_prior_p=True,
    prior_r=(1.05, 1.05),
    tau=25.0,
):
    X = np.asarray(X, float)
    y = np.asarray(y, int)
    rng = np.random.default_rng(seed)
    n = len(X)

    theta_hat, res = fit_bayes(
        X,
        y,
        seed=seed,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
    )

    pmap = P_with(theta_hat, x_grid)
    x50 = x_at_p(theta_hat, 0.5, lo=max(1e-6, x_grid.min()), hi=x_grid.max())
    s50 = slope_at_x(theta_hat, x50)

    phi_hat = res.x
    alpha, beta = make_priors(y, tau=tau)

    H = hess_fd(
        lambda w: neg_post(
            w,
            X,
            y,
            alpha,
            beta,
            use_prior_p=use_prior_p,
            prior_r=prior_r,
        ),
        phi_hat,
    )

    Jp = jac_fd(lambda w: P_with(unpack(w), x_grid), phi_hat)

    try:
        Sigma_phi = np.linalg.inv(H)
    except np.linalg.LinAlgError:
        Sigma_phi = np.linalg.pinv(H)

    var_p = np.einsum("ij,jk,ik->i", Jp, Sigma_phi, Jp)
    se_p = np.sqrt(np.maximum(var_p, 0.0))
    wald_lo = np.clip(pmap - 1.96 * se_p, 0, 1)
    wald_hi = np.clip(pmap + 1.96 * se_p, 0, 1)

    def theta_vec_from_phi(w):
        p, a, b, s, k, th, thcap = unpack(w)
        r = th / thcap if np.isfinite(thcap) and thcap > 0 else np.nan
        return np.array([p, a, b, s, k, th, r], float)

    Jtheta = jac_fd(theta_vec_from_phi, phi_hat)
    Sigma_theta = Jtheta @ Sigma_phi @ Jtheta.T
    theta_hat_vec = theta_vec_from_phi(phi_hat)
    se_theta = np.sqrt(np.maximum(np.diag(Sigma_theta), 0.0))
    wald_param_lo = theta_hat_vec - 1.96 * se_theta
    wald_param_hi = theta_hat_vec + 1.96 * se_theta

    curves_np = []
    theta_np = []
    x50_np = []
    used_np = 0

    for _ in range(B_nonpar):
        idx = rng.integers(0, n, n)
        Xb, yb = X[idx], y[idx]

        if yb.sum() == 0 or yb.sum() == len(yb):
            continue

        try:
            thb, rb = fit_bayes(
                Xb,
                yb,
                seed=int(rng.integers(0, 10_000_000)),
                use_prior_p=use_prior_p,
                prior_r=prior_r,
                tau=tau,
            )
            if not rb.success or not np.isfinite(rb.fun):
                continue

            curves_np.append(P_with(thb, x_grid))

            p_b, a_b, b_b, s_b, k_b, th_b, thcap_b = thb
            r_b = th_b / thcap_b if np.isfinite(thcap_b) and thcap_b > 0 else np.nan
            theta_np.append([p_b, a_b, b_b, s_b, k_b, th_b, r_b])

            x50_b = x_at_p(thb, 0.5, lo=max(1e-6, x_grid.min()), hi=x_grid.max())
            x50_np.append(x50_b)

            used_np += 1
        except Exception:
            continue

    curves_np = np.asarray(curves_np)
    theta_np = np.asarray(theta_np, float) if len(theta_np) else np.empty((0, 7))
    x50_np = np.asarray(x50_np, float) if len(x50_np) else np.empty((0,))

    np_lo = np.percentile(curves_np, 2.5, axis=0) if used_np else None
    np_hi = np.percentile(curves_np, 97.5, axis=0) if used_np else None

    curves_pb = []
    theta_pb = []
    x50_pb = []
    used_pb = 0

    p_hat, a_hat, b_hat, s_hat, k_hat, th_hat, _ = theta_hat

    for _ in range(B_param):
        yb = rng.binomial(1, p_hat, size=n)

        if yb.sum() == 0 or yb.sum() == n:
            continue

        Xb = np.zeros(n, dtype=float)

        idx_nc = np.where(yb == 0)[0]
        idx_ae = np.where(yb == 1)[0]

        if len(idx_nc) > 0:
            Xb[idx_nc] = stats.gamma.rvs(
                k_hat,
                scale=th_hat,
                size=len(idx_nc),
                random_state=rng,
            )

        if len(idx_ae) > 0:
            Xb[idx_ae] = stats.betaprime.rvs(
                a_hat,
                b_hat,
                scale=s_hat,
                size=len(idx_ae),
                random_state=rng,
            )

        Xb = np.clip(Xb, 1e-12, None)

        try:
            thb, rb = fit_bayes(
                Xb,
                yb,
                seed=int(rng.integers(0, 10_000_000)),
                use_prior_p=use_prior_p,
                prior_r=prior_r,
                tau=tau,
            )
            if not rb.success or not np.isfinite(rb.fun):
                continue

            curves_pb.append(P_with(thb, x_grid))

            p_b, a_b, b_b, s_b, k_b, th_b, thcap_b = thb
            r_b = th_b / thcap_b if np.isfinite(thcap_b) and thcap_b > 0 else np.nan
            theta_pb.append([p_b, a_b, b_b, s_b, k_b, th_b, r_b])

            x50_b = x_at_p(thb, 0.5, lo=max(1e-6, x_grid.min()), hi=x_grid.max())
            x50_pb.append(x50_b)

            used_pb += 1
        except Exception:
            continue

    curves_pb = np.asarray(curves_pb)
    theta_pb = np.asarray(theta_pb, float) if len(theta_pb) else np.empty((0, 7))
    x50_pb = np.asarray(x50_pb, float) if len(x50_pb) else np.empty((0,))

    pb_lo = np.percentile(curves_pb, 2.5, axis=0) if used_pb else None
    pb_hi = np.percentile(curves_pb, 97.5, axis=0) if used_pb else None

    return {
        "label": label,
        "theta_hat": theta_hat,
        "res": res,
        "x50": x50,
        "s50": s50,
        "pmap": pmap,
        "wald_lo": wald_lo,
        "wald_hi": wald_hi,
        "np_lo": np_lo,
        "np_hi": np_hi,
        "pb_lo": pb_lo,
        "pb_hi": pb_hi,
        "used_np": used_np,
        "used_pb": used_pb,
        "theta_hat_vec": theta_hat_vec,
        "wald_param_lo": wald_param_lo,
        "wald_param_hi": wald_param_hi,
        "theta_np": theta_np,
        "theta_pb": theta_pb,
        "x50_np": x50_np,
        "x50_pb": x50_pb,
    }


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=10.0,
    n_grid=600,
    B_nonpar=400,
    B_param=400,
    seed=123,
    use_prior_p=True,
    prior_r=(1.05, 1.05),
    tau=25.0,
):
    """
    Run CI estimation for both ORIGINAL and TRIM datasets.
    """
    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_grid = np.linspace(0, xmax, n_grid)

    out_orig = estimate_ci_bundle(
        ds["X_orig"],
        ds["y_orig"],
        label="ORIGINAL",
        x_grid=x_grid,
        B_nonpar=B_nonpar,
        B_param=B_param,
        seed=seed,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
    )

    out_trim = estimate_ci_bundle(
        ds["X_trim"],
        ds["y_trim"],
        label="TRIM",
        x_grid=x_grid,
        B_nonpar=B_nonpar,
        B_param=B_param,
        seed=seed + 1,
        use_prior_p=use_prior_p,
        prior_r=prior_r,
        tau=tau,
    )

    return {
        **ds,
        "x_grid": x_grid,
        "orig": out_orig,
        "trim": out_trim,
    }


def make_param_ci_table(ci_res):
    """
    Parameter CI table for FULL and TRIM, for Wald / Nonparam / Parametric.
    """
    rows = []
    names = ["p", "a", "b", "s", "k", "theta", "r"]

    for dataset_key, dataset_name in [("orig", "FULL"), ("trim", "TRIM")]:
        out = ci_res[dataset_key]
        hat = out["theta_hat_vec"]

        for i, name in enumerate(names):
            rows.append(
                {
                    "Dataset": dataset_name,
                    "Method": "Wald",
                    "Parameter": name,
                    "Estimate": hat[i],
                    "LL": out["wald_param_lo"][i],
                    "UL": out["wald_param_hi"][i],
                }
            )

        if out["theta_np"].shape[0] > 0:
            lo = np.nanpercentile(out["theta_np"], 2.5, axis=0)
            hi = np.nanpercentile(out["theta_np"], 97.5, axis=0)
            for i, name in enumerate(names):
                rows.append(
                    {
                        "Dataset": dataset_name,
                        "Method": "Nonparam",
                        "Parameter": name,
                        "Estimate": hat[i],
                        "LL": lo[i],
                        "UL": hi[i],
                    }
                )

        if out["theta_pb"].shape[0] > 0:
            lo = np.nanpercentile(out["theta_pb"], 2.5, axis=0)
            hi = np.nanpercentile(out["theta_pb"], 97.5, axis=0)
            for i, name in enumerate(names):
                rows.append(
                    {
                        "Dataset": dataset_name,
                        "Method": "Parametric",
                        "Parameter": name,
                        "Estimate": hat[i],
                        "LL": lo[i],
                        "UL": hi[i],
                    }
                )

    return pd.DataFrame(rows)


def make_x50_ci_table(ci_res):
    """
    x50 CI table for FULL and TRIM, for Wald / Nonparam / Parametric.
    """
    rows = []

    for dataset_key, dataset_name in [("orig", "FULL"), ("trim", "TRIM")]:
        out = ci_res[dataset_key]

        wald_x50_lo = np.nan
        wald_x50_hi = np.nan
        try:
            wald_x50_lo = np.interp(0.5, out["wald_lo"], ci_res["x_grid"])
            wald_x50_hi = np.interp(0.5, out["wald_hi"], ci_res["x_grid"])
        except Exception:
            pass

        rows.append(
            {
                "Dataset": dataset_name,
                "Method": "Wald",
                "Estimate": out["x50"],
                "LL": wald_x50_lo,
                "UL": wald_x50_hi,
            }
        )

        if len(out["x50_np"]) > 0:
            rows.append(
                {
                    "Dataset": dataset_name,
                    "Method": "Nonparam",
                    "Estimate": out["x50"],
                    "LL": np.nanpercentile(out["x50_np"], 2.5),
                    "UL": np.nanpercentile(out["x50_np"], 97.5),
                }
            )

        if len(out["x50_pb"]) > 0:
            rows.append(
                {
                    "Dataset": dataset_name,
                    "Method": "Parametric",
                    "Estimate": out["x50"],
                    "LL": np.nanpercentile(out["x50_pb"], 2.5),
                    "UL": np.nanpercentile(out["x50_pb"], 97.5),
                }
            )

    return pd.DataFrame(rows)


def make_curve_ci_table(ci_res, grid_every=25):
    """
    Long-format curve CI table.
    Contains LL/UL of P(AE|X) across x-grid for all methods.
    """
    rows = []
    x_grid = ci_res["x_grid"][::grid_every]

    for dataset_key, dataset_name in [("orig", "FULL"), ("trim", "TRIM")]:
        out = ci_res[dataset_key]

        for method, lo_key, hi_key in [
            ("Wald", "wald_lo", "wald_hi"),
            ("Nonparam", "np_lo", "np_hi"),
            ("Parametric", "pb_lo", "pb_hi"),
        ]:
            lo = out.get(lo_key, None)
            hi = out.get(hi_key, None)
            if lo is None or hi is None:
                continue

            lo = lo[::grid_every]
            hi = hi[::grid_every]
            est = out["pmap"][::grid_every]

            for x, e, l, u in zip(x_grid, est, lo, hi):
                rows.append(
                    {
                        "Dataset": dataset_name,
                        "Method": method,
                        "X": x,
                        "Estimate": e,
                        "LL": l,
                        "UL": u,
                    }
                )

    return pd.DataFrame(rows)


def plot_ci_original_trim(ci_res, xmax=6.0):
    """
    2-panel figure:
      left  = FULL
      right = TRIM
    """
    x_grid = ci_res["x_grid"]
    rng = np.random.default_rng(999)

    fig, axes = plt.subplots(1, 2, figsize=(16, 7.0), dpi=150, sharey=True)

    for ax, X, y, out, title, panel in [
        (axes[0], ci_res["X_orig"], ci_res["y_orig"], ci_res["orig"], "FULL", "A"),
        (axes[1], ci_res["X_trim"], ci_res["y_trim"], ci_res["trim"], "TRIM", "B"),
    ]:
        ax.fill_between(x_grid, out["wald_lo"], out["wald_hi"], color="#2ca02c", alpha=0.10)

        if out["used_np"]:
            ax.fill_between(x_grid, out["np_lo"], out["np_hi"], color="#17becf", alpha=0.10)

        if out["used_pb"]:
            ax.fill_between(x_grid, out["pb_lo"], out["pb_hi"], color="#e91e63", alpha=0.10)

        ax.plot(x_grid, out["wald_lo"], color="#2ca02c", lw=1.6, ls="--")
        h_wald, = ax.plot(
            x_grid,
            out["wald_hi"],
            color="#2ca02c",
            lw=1.6,
            ls="--",
            label="CI: Wald (delta) 95%",
        )

        h_np = None
        if out["used_np"]:
            ax.plot(x_grid, out["np_lo"], color="#17becf", lw=1.6, ls=(0, (1, 2)))
            h_np, = ax.plot(
                x_grid,
                out["np_hi"],
                color="#17becf",
                lw=1.6,
                ls=(0, (1, 2)),
                label="CI: Nonparam bootstrap 95%",
            )

        h_pb = None
        if out["used_pb"]:
            ax.plot(x_grid, out["pb_lo"], color="#e91e63", lw=1.6, ls="-.")
            h_pb, = ax.plot(
                x_grid,
                out["pb_hi"],
                color="#e91e63",
                lw=1.6,
                ls="-.",
                label="CI: Parametric bootstrap 95%",
            )

        h_fit, = ax.plot(x_grid, out["pmap"], color="k", lw=2.4, label="Bayesian fit")

        jit = (rng.random(len(y)) - 0.5) * 0.035

        h_nc = ax.scatter(
            X[y == 0],
            (y + jit)[y == 0],
            s=18,
            alpha=0.55,
            color="#5dade2",
            edgecolors="none",
            label="NC",
        )

        h_ae = ax.scatter(
            X[y == 1],
            (y + jit)[y == 1],
            s=24,
            alpha=0.80,
            color="#f39c3d",
            edgecolors="none",
            label="AE",
        )

        ax.text(
            0.02,
            0.98,
            panel,
            transform=ax.transAxes,
            ha="left",
            va="top",
            fontsize=16,
            fontweight="bold",
        )

        ax.set_xlim(0, xmax)
        ax.set_ylim(-0.05, 1.05)
        ax.set_xlabel("X", fontsize=13, fontweight="bold")
        ax.set_title(title, fontsize=13, fontweight="bold")
        ax.grid(alpha=0.25)

        handles = [h_nc, h_ae, h_fit, h_wald]
        if h_np is not None:
            handles.append(h_np)
        if h_pb is not None:
            handles.append(h_pb)

        labels = [h.get_label() for h in handles]

        ax.legend(
            handles,
            labels,
            loc="upper center",
            bbox_to_anchor=(0.5, -0.20),
            ncol=2,
            frameon=False,
            fontsize=10,
        )

    axes[0].set_ylabel("P(AE | X)", fontsize=13, fontweight="bold")
    plt.tight_layout(rect=(0, 0.08, 1, 1))
    return fig, axes


# ============================================================
# 5) X50 ELASTICITY ANALYSIS
# ============================================================

PARAM_NAMES_X50_ELAS = [r"$\pi$", r"$a$", r"$b$", r"$s$", r"$k$", r"$\vartheta$"]


def theta6_from_hat(theta_hat):
    """
    Extract the first 6 raw model parameters from theta_hat:
      (p, a, b, s, k, th)
    """
    th = np.asarray(theta_hat, float).ravel()
    if th.size < 6:
        raise ValueError(f"Expected at least 6 parameters, got {th.size}")
    return th[:6].copy()


def step_vec_theta(theta, rel_step=1e-6, abs_min=1e-10):
    """
    Relative finite-difference step on raw theta scale.
    """
    theta = np.asarray(theta, float).ravel()
    return np.maximum(abs_min, rel_step * np.maximum(1.0, np.abs(theta)))


def grad_central_theta(F_theta, theta0, rel_step=1e-6, abs_min=1e-10, pi_eps=1e-12):
    """
    Central differences in RAW theta.
    Keeps:
      - p in (pi_eps, 1-pi_eps)
      - positive parameters > 0
    """
    theta0 = np.asarray(theta0, float).ravel()
    h = step_vec_theta(theta0, rel_step=rel_step, abs_min=abs_min)
    g = np.zeros_like(theta0)

    for j in range(theta0.size):
        th_plus = theta0.copy()
        th_minus = theta0.copy()
        hj = h[j]

        if j == 0:
            p0 = float(np.clip(theta0[0], pi_eps, 1 - pi_eps))
            hj = min(hj, p0 - pi_eps, (1 - pi_eps) - p0)
            hj = max(hj, abs_min)
            th_plus[0] = np.clip(p0 + hj, pi_eps, 1 - pi_eps)
            th_minus[0] = np.clip(p0 - hj, pi_eps, 1 - pi_eps)
        else:
            q0 = float(max(theta0[j], 1e-15))
            hj = min(hj, 0.5 * q0)
            hj = max(hj, abs_min)
            th_plus[j] = q0 + hj
            th_minus[j] = max(q0 - hj, 1e-15)

        g[j] = (F_theta(th_plus) - F_theta(th_minus)) / (2.0 * hj)

    return g


def P_with_theta6(theta6, x):
    """
    Same posterior risk curve as P_with(), but accepts only the 6 raw parameters:
    (p, a, b, s, k, th)
    """
    p, a, b, s, k, th = np.asarray(theta6, float).ravel()[:6]
    x = np.asarray(x, float)
    eps = 1e-12

    logit_val = (
        np.log(np.clip(p, eps, 1 - eps))
        - np.log(np.clip(1 - p, eps, 1.0))
        + dE_full(x, a, b, s, k, th)
    )
    return logistic(logit_val)


def x50_theta6(theta6, lo=1e-6, hi=6.0, hi_max=100.0):
    """
    Solve P(AE|x) = 0.5 using theta6 = (p, a, b, s, k, th),
    with adaptive bracketing.
    """
    f = lambda x: P_with_theta6(theta6, x) - 0.5

    fa = f(lo)
    fb = f(hi)

    while np.isfinite(fa) and np.isfinite(fb) and fa * fb > 0 and hi < hi_max:
        hi *= 2.0
        fb = f(hi)

    if (not np.isfinite(fa)) or (not np.isfinite(fb)) or fa * fb > 0:
        return np.nan

    try:
        return float(brentq(f, lo, hi))
    except Exception:
        return np.nan


def compute_x50_elasticity_rawtheta(
    theta0,
    x50_fun=x50_theta6,
    rel_step=1e-6,
    abs_min=1e-10,
    pi_eps=1e-12,
):
    """
    x50 elasticity on RAW theta scale:

      E_x50_j = | (theta_j / x50) * d x50 / d theta_j |

    where theta0 = [p, a, b, s, k, th]
    """
    theta0 = np.asarray(theta0, float).ravel()
    if theta0.size != 6:
        raise ValueError(f"Expected theta0 of length 6, got {theta0.size}")

    x50_0 = float(x50_fun(theta0))

    d_x_dth = grad_central_theta(
        x50_fun, theta0, rel_step=rel_step, abs_min=abs_min, pi_eps=pi_eps
    )

    eps = 1e-12
    x_safe = max(abs(x50_0), eps)

    th_safe = theta0.copy()
    th_safe[0] = float(np.clip(th_safe[0], pi_eps, 1 - pi_eps))
    th_safe[1:] = np.maximum(th_safe[1:], 1e-15)

    elas_x = np.abs(d_x_dth) * np.abs(th_safe) / x_safe

    return {
        "theta0": theta0,
        "x50": x50_0,
        "d_x_dtheta": d_x_dth,
        "elas_x50": elas_x,
        "param_names": PARAM_NAMES_X50_ELAS,
    }


def run_x50_elasticity_from_fit_results(
    fit_results,
    rel_step=1e-6,
    abs_min=1e-10,
    pi_eps=1e-12,
):
    """
    Compute x50 elasticity for FULL (orig) and TRIM directly from fit_results.
    """
    elas_orig = compute_x50_elasticity_rawtheta(
        theta6_from_hat(fit_results["theta_orig"]),
        x50_fun=x50_theta6,
        rel_step=rel_step,
        abs_min=abs_min,
        pi_eps=pi_eps,
    )

    elas_trim = compute_x50_elasticity_rawtheta(
        theta6_from_hat(fit_results["theta_trim"]),
        x50_fun=x50_theta6,
        rel_step=rel_step,
        abs_min=abs_min,
        pi_eps=pi_eps,
    )

    return {
        "orig": elas_orig,
        "trim": elas_trim,
    }


def make_x50_elasticity_table(elas_res, dataset_name="FULL"):
    """
    Tidy x50 elasticity table.
    """
    rows = []
    for name, ex in zip(
        elas_res["param_names"],
        elas_res["elas_x50"],
    ):
        rows.append(
            {
                "Dataset": dataset_name,
                "Parameter": name,
                "Elasticity_x50": float(ex),
            }
        )
    return pd.DataFrame(rows)


def plot_x50_elasticity_bars(elas_full, elas_trim, figsize=(7, 5), dpi=150):
    """
    One-panel bar plot for x50 elasticity.
    """
    names = elas_full["param_names"]
    x = np.arange(len(names))
    width = 0.36

    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)

    ax.bar(
        x - width / 2,
        elas_full["elas_x50"],
        width=width,
        label="FULL",
        alpha=0.85,
    )
    ax.bar(
        x + width / 2,
        elas_trim["elas_x50"],
        width=width,
        label="TRIM",
        alpha=0.85,
    )

    ax.set_xticks(x)
    ax.set_xticklabels(names)
    ax.set_ylabel("Elasticity", fontweight="bold")
    ax.set_title("Elasticity of x50", fontweight="bold")
    ax.grid(axis="y", alpha=0.25)
    ax.legend(frameon=False)

    plt.tight_layout()
    return fig, ax