qra/logit.py

import numpy as np


# ============================================================
# Stable sigmoid
# ============================================================

def _sigmoid_stable(z):
    z = np.asarray(z, float)
    z = np.clip(z, -50.0, 50.0)
    return 1.0 / (1.0 + np.exp(-z))


# ============================================================
# 1) Model
# ============================================================

def model_p(x, b):
    """p(x|b) = sigmoid(b0 + b1*x)."""
    x = np.asarray(x, float).reshape(-1)
    b0, b1 = np.asarray(b, float).reshape(2)
    return _sigmoid_stable(b0 + b1 * x)


def design_matrix(x):
    """Design matrix X = [1, x]."""
    x = np.asarray(x, float).reshape(-1)
    return np.column_stack([np.ones_like(x), x])


# ============================================================
# 2) Likelihood
# ============================================================

def nll(x, y, b, l2=0.0):
    """
    Penalized negative log-likelihood:
        NLL(b) = -sum[y log p + (1-y) log(1-p)] + 0.5*l2*||b||^2
    """
    x = np.asarray(x, float).reshape(-1)
    y = np.asarray(y, float).reshape(-1)
    b = np.asarray(b, float).reshape(2)

    p = model_p(x, b)
    eps = 1e-12
    p = np.clip(p, eps, 1 - eps)

    base = -np.sum(y * np.log(p) + (1 - y) * np.log(1 - p))
    pen = 0.5 * l2 * float(np.dot(b, b))
    return base + pen


def llf(x, y, b):
    """
    Ordinary (unpenalized) log-likelihood at fitted parameters.
    """
    x = np.asarray(x, float).reshape(-1)
    y = np.asarray(y, float).reshape(-1)
    b = np.asarray(b, float).reshape(2)

    p = model_p(x, b)
    eps = 1e-12
    p = np.clip(p, eps, 1 - eps)
    return float(np.sum(y * np.log(p) + (1 - y) * np.log(1 - p)))


# ============================================================
# 3) Gradient / Hessian / Covariance
# ============================================================

def grad_nll(x, y, b, l2=0.0):
    """
    Gradient of penalized NLL:
        g(b) = X^T (p - y) + l2*b
    """
    X = design_matrix(x)
    y = np.asarray(y, float).reshape(-1)
    b = np.asarray(b, float).reshape(2)

    p = model_p(x, b)
    return X.T @ (p - y) + l2 * b


def hess_nll(x, b, l2=0.0):
    """
    Hessian of penalized NLL:
        H(b) = X^T W X + l2*I
        W = diag(p*(1-p))
    """
    X = design_matrix(x)
    b = np.asarray(b, float).reshape(2)

    p = model_p(x, b)
    w = p * (1 - p)
    return X.T @ (w[:, None] * X) + l2 * np.eye(2)


def covariance(x, b, l2=0.0):
    """
    Cov(b) ≈ H(b)^(-1). Robust to near-singular Hessians.
    """
    H = hess_nll(x, b, l2=l2)
    try:
        return np.linalg.inv(H)
    except np.linalg.LinAlgError:
        return np.linalg.pinv(H)


def standard_errors(x, b, l2=0.0):
    """
    SE = sqrt(diag(Cov)).
    """
    C = covariance(x, b, l2=l2)
    return np.sqrt(np.maximum(np.diag(C), 0.0))


# Compatibility alias
def logit_poly_cov(x, b, l2=0.0):
    return covariance(x, b, l2=l2)


# ============================================================
# 4) Fit
# ============================================================

def fit_newton(x, y, b_start=None, max_iter=50, tol=1e-8, l2=0.0):
    """
    Newton updates for penalized NLL with backtracking line-search.

    Update:
        b_new = b - alpha * H^{-1} g
    alpha shrinks until NLL decreases.
    """
    x = np.asarray(x, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)

    if b_start is None:
        b = np.array([0.0, 0.0], float)
    else:
        b = np.asarray(b_start, float).reshape(2)

    f = nll(x, y, b, l2=l2)

    for _ in range(max_iter):
        g = grad_nll(x, y, b, l2=l2)
        H = hess_nll(x, b, l2=l2)

        try:
            step = np.linalg.solve(H, g)
        except np.linalg.LinAlgError:
            step = np.linalg.pinv(H) @ g

        alpha = 1.0
        while alpha > 1e-6:
            b_new = b - alpha * step
            f_new = nll(x, y, b_new, l2=l2)
            if np.isfinite(f_new) and f_new <= f:
                break
            alpha *= 0.5

        if alpha <= 1e-6:
            break

        if np.max(np.abs(b_new - b)) < tol:
            b = b_new
            break

        b, f = b_new, f_new

    return b


# ============================================================
# 5) Goodness of fit
# ============================================================

def goodness_of_fit(x, y, b, thresh=0.5, l2=0.0):
    """
    Returns:
      LLF, NLL, AIC, BIC, Accuracy, n, k

    Note:
      Fit may use l2 > 0, but GOF metrics below are computed from the
      ordinary (unpenalized) likelihood at the fitted parameters.
    """
    x = np.asarray(x, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)
    b = np.asarray(b, float).reshape(2)

    p = model_p(x, b)
    eps = 1e-12
    p = np.clip(p, eps, 1 - eps)

    LLF = np.sum(y * np.log(p) + (1 - y) * np.log(1 - p))
    NLL = -LLF

    n = len(x)
    k = len(b)

    AIC = 2 * k - 2 * LLF
    BIC = k * np.log(n) - 2 * LLF

    yhat = (p >= thresh).astype(int)
    acc = np.mean(yhat == y)

    return {"LLF": LLF, "NLL": NLL, "AIC": AIC, "BIC": BIC, "A": acc, "n": n, "k": k}


# ============================================================
# 6) x50 / Wald helpers / compact fit
# ============================================================

def x50(b):
    """
    Model-scale midpoint:
        x50 = -b0 / b1

    For LOG panels, this is on the log(x) scale.
    Raw-scale SUV50 is exp(x50).
    """
    b0, b1 = np.asarray(b, float).reshape(2)
    return np.nan if np.abs(b1) < 1e-12 else (-b0 / b1)


def x50_wald_ci(b, cov, z=1.959963984540054):
    """
    Wald CI for x50 = -b0/b1 via delta method.
    Returned on MODEL scale.
    """
    b = np.asarray(b, float).reshape(2)
    cov = np.asarray(cov, float).reshape(2, 2)
    b0, b1 = b

    if np.abs(b1) < 1e-12:
        return np.nan, np.nan

    xhat = -b0 / b1
    grad = np.array([-1.0 / b1, b0 / (b1 ** 2)], float)
    var = float(grad.T @ cov @ grad)
    se = np.sqrt(max(var, 0.0))
    return float(xhat - z * se), float(xhat + z * se)


def wald_ci(b, cov, z=1.959963984540054):
    """
    Wald CI for parameters: b_i ± z*SE_i.
    """
    b = np.asarray(b, float).reshape(2)
    cov = np.asarray(cov, float).reshape(2, 2)
    se = np.sqrt(np.maximum(np.diag(cov), 0.0))
    return b - z * se, b + z * se


def fit_pack(x, y, name="", thresh=0.5, l2=0.0):
    """
    Fit + covariance + GOF + parameter Wald CI.
    x should already be on the MODEL scale.
    """
    x = np.asarray(x, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)

    b = fit_newton(x, y, l2=l2)
    cov = covariance(x, b, l2=l2)
    gof = goodness_of_fit(x, y, b, thresh=thresh, l2=l2)
    lcl, ucl = wald_ci(b, cov)

    return {
        "name": name,
        "x": x,
        "y": y,
        "b": b,
        "cov": cov,
        "gof": gof,
        "LCL": lcl,
        "UCL": ucl,
        "l2": float(l2),
    }


def trim_nc_by_value(x_raw, y, target=2.48, tol=0.05):
    """
    Remove ONE NC sample (y==0) with x_raw closest to target.
    """
    x_raw = np.asarray(x_raw, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)

    nc_idx = np.where(y == 0)[0]
    if len(nc_idx) == 0:
        raise ValueError("No NC samples found (y==0).")

    j = nc_idx[np.argmin(np.abs(x_raw[nc_idx] - target))]
    diff = float(np.abs(x_raw[j] - target))
    if diff > tol:
        print(f"[trim warning] closest NC to {target} is {x_raw[j]:.6f} (diff={diff:.6f}) > tol={tol}")

    mask = np.ones_like(y, dtype=bool)
    mask[j] = False
    print(f"[trim] removed index={j}, x_raw={x_raw[j]:.6f}, y={y[j]}")
    return x_raw[mask], y[mask]


# ============================================================
# 7) Analytic CI bands on a grid
# ============================================================

def eta_se_grid(x_grid, cov):
    """
    Standard error of eta(x) = b0 + b1*x on a grid.
    x_grid must be on the MODEL scale.
    """
    x_grid = np.asarray(x_grid, float).reshape(-1)
    cov = np.asarray(cov, float).reshape(2, 2)

    Xg = design_matrix(x_grid)
    var_eta = np.einsum("ij,jk,ik->i", Xg, cov, Xg)
    return np.sqrt(np.maximum(var_eta, 0.0))


def ci_band_normal(x_grid, b, cov, z=1.959963984540054):
    """
    Normal-on-eta CI band:
        eta ± z*SE(eta), then transform with sigmoid.
    Returns: lo, mid, hi
    """
    x_grid = np.asarray(x_grid, float).reshape(-1)
    b = np.asarray(b, float).reshape(2)

    eta = b[0] + b[1] * x_grid
    se_eta = eta_se_grid(x_grid, cov)

    lo = _sigmoid_stable(eta - z * se_eta)
    md = _sigmoid_stable(eta)
    hi = _sigmoid_stable(eta + z * se_eta)
    return lo, md, hi


def ci_band_delta(x_grid, b, cov, z=1.959963984540054):
    """
    Delta-method CI band on probability scale:
        p(x) ± z * SE_p(x)
    where
        SE_p = p(1-p) * SE_eta
    Returns: lo, mid, hi
    """
    x_grid = np.asarray(x_grid, float).reshape(-1)
    b = np.asarray(b, float).reshape(2)

    p = model_p(x_grid, b)
    se_eta = eta_se_grid(x_grid, cov)
    se_p = p * (1.0 - p) * se_eta

    lo = np.clip(p - z * se_p, 0.0, 1.0)
    md = p
    hi = np.clip(p + z * se_p, 0.0, 1.0)
    return lo, md, hi


# ============================================================
# 8) Bootstrap parameter generators
# ============================================================

def bootstrap_params_stratified(x, y, B=2000, seed=123, l2=0.0, b_start=None):
    """
    Stratified nonparametric bootstrap on MODEL-scale x.
    Preserves class counts exactly.
    Returns array of shape (n_ok, 2).
    """
    rng = np.random.default_rng(seed)

    x = np.asarray(x, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)

    x0 = x[y == 0]
    x1 = x[y == 1]

    n0 = len(x0)
    n1 = len(x1)

    if n0 == 0 or n1 == 0:
        return np.empty((0, 2), float)

    out = []
    for _ in range(B):
        xb0 = rng.choice(x0, size=n0, replace=True)
        xb1 = rng.choice(x1, size=n1, replace=True)

        xb = np.concatenate([xb0, xb1])
        yb = np.concatenate([np.zeros(n0, dtype=int), np.ones(n1, dtype=int)])

        try:
            bb = fit_newton(xb, yb, b_start=b_start, l2=l2)
            if np.all(np.isfinite(bb)):
                out.append(bb)
        except Exception:
            pass

    if len(out) == 0:
        return np.empty((0, 2), float)
    return np.asarray(out, float)


def bootstrap_params_parametric(x, b, B=2000, seed=123, l2=0.0, min_ae=2):
    """
    Parametric bootstrap on MODEL-scale x.
    Simulates y* ~ Bernoulli(p_hat(x)).
    Keeps only samples with at least min_ae positives and at least one negative.
    Returns array of shape (n_ok, 2).
    """
    rng = np.random.default_rng(seed)

    x = np.asarray(x, float).reshape(-1)
    b = np.asarray(b, float).reshape(2)

    p = model_p(x, b)
    n = len(x)

    out = []
    tries = 0
    max_tries = max(10 * B, 1000)

    while len(out) < B and tries < max_tries:
        tries += 1
        yb = rng.binomial(1, p, size=n).astype(int)

        n1 = int(np.sum(yb))
        n0 = n - n1
        if n1 < min_ae or n0 < 1:
            continue

        try:
            bb = fit_newton(x, yb, b_start=b, l2=l2)
            if np.all(np.isfinite(bb)):
                out.append(bb)
        except Exception:
            pass

    if len(out) == 0:
        return np.empty((0, 2), float)
    return np.asarray(out, float)


# ============================================================
# 9) Convert bootstrap parameters to curve bands
# ============================================================

def bootstrap_band_from_params(x_grid, pars, alpha=0.05):
    """
    Build bootstrap CI band from bootstrap parameter draws.
    x_grid is on MODEL scale.
    Returns: lo, mid, hi
    """
    x_grid = np.asarray(x_grid, float).reshape(-1)
    pars = np.asarray(pars, float)

    if pars.ndim != 2 or pars.shape[0] == 0:
        nan = np.full_like(x_grid, np.nan, dtype=float)
        return nan, nan, nan

    curves = np.array([model_p(x_grid, bb) for bb in pars], float)
    q = np.quantile(curves, [alpha / 2, 0.5, 1.0 - alpha / 2], axis=0)
    return q[0], q[1], q[2]


# ============================================================
# 10) High-level wrapper for one panel
# ============================================================

def fit_ci_pack_rawgrid(
    x_raw,
    y,
    transform="raw",
    xmax_raw=None,
    grid_n=500,
    name="",
    l2=0.0,
    B=2000,
    seed=123,
    min_ae=2,
    z=1.959963984540054,
):
    """
    Fit one panel and return everything needed for plots/tables.
    """
    x_raw = np.asarray(x_raw, float).reshape(-1)
    y = np.asarray(y, int).reshape(-1)

    if transform not in ("raw", "log"):
        raise ValueError("transform must be 'raw' or 'log'")

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

    x_model = x_raw if transform == "raw" else np.log(x_raw)

    b = fit_newton(x_model, y, l2=l2)
    cov = covariance(x_model, b, l2=l2)
    gof = goodness_of_fit(x_model, y, b, l2=l2)

    xmin_raw = float(np.min(x_raw))
    xmax0 = float(np.max(x_raw))
    xmax_use = xmax0 if xmax_raw is None else max(float(xmax_raw), xmax0)
    x_grid_raw = np.linspace(xmin_raw, xmax_use, int(grid_n))

    x_grid_model = x_grid_raw if transform == "raw" else np.log(x_grid_raw)

    lo_n, md_n, hi_n = ci_band_normal(x_grid_model, b, cov, z=z)
    lo_d, md_d, hi_d = ci_band_delta(x_grid_model, b, cov, z=z)

    pars_np = bootstrap_params_stratified(
        x_model, y, B=B, seed=seed + 1, l2=l2, b_start=b
    )
    pars_pm = bootstrap_params_parametric(
        x_model, b, B=B, seed=seed + 2, l2=l2, min_ae=min_ae
    )

    lo_np, md_np, hi_np = bootstrap_band_from_params(x_grid_model, pars_np)
    lo_pm, md_pm, hi_pm = bootstrap_band_from_params(x_grid_model, pars_pm)

    return {
        "name": name,
        "transform": transform,
        "l2": float(l2),
        "x_raw": x_raw,
        "x_model": x_model,
        "y": y,
        "x_grid_raw": x_grid_raw,
        "x_grid_model": x_grid_model,
        "b": b,
        "cov": cov,
        "gof": gof,
        "bands": {
            "Normal":     (lo_n,  md_n,  hi_n),
            "Delta":      (lo_d,  md_d,  hi_d),
            "Nonparam":   (lo_np, md_np, hi_np),
            "Parametric": (lo_pm, md_pm, hi_pm),
        },
        "pars_nonparam": pars_np,
        "pars_parametric": pars_pm,
    }


# ============================================================
# 11) Model-band table with LL / UL
# ============================================================

def model_ci_table_4methods(
    P,
    keys=("FULL-RAW", "TRIM-RAW", "FULL-LOG", "TRIM-LOG"),
):
    """
    Long table of model CI bands on the grid.
    Includes:
      x_grid_model, x_grid_raw, fit, LL, UL
    """
    import pandas as pd

    rows = []

    for key in keys:
        pk = P[key]
        xg_raw = np.asarray(pk["x_grid_raw"], float)
        xg_mod = np.asarray(pk["x_grid_model"], float)
        trans = pk.get("transform", "")
        bands = pk["bands"]

        for method, (lo, md, hi) in bands.items():
            lo = np.asarray(lo, float)
            md = np.asarray(md, float)
            hi = np.asarray(hi, float)

            for i in range(len(xg_raw)):
                rows.append({
                    "Panel": key,
                    "Method": method,
                    "transform": trans,
                    "x_grid_raw": float(xg_raw[i]),
                    "x_grid_model": float(xg_mod[i]),
                    "fit": float(md[i]),
                    "LL": float(lo[i]),
                    "UL": float(hi[i]),
                })

    return pd.DataFrame(rows)


# ============================================================
# 12) Parameter/x50 CI summary table
# ============================================================

def param_ci_table_4methods(
    P,
    keys=("FULL-RAW", "TRIM-RAW", "FULL-LOG", "TRIM-LOG"),
    z=1.959963984540054,
    include_point_est=True,
):
    """
    Build tidy parameter/x50 CI table for:
      Normal, Delta, Nonparam, Parametric

    Notes
    -----
    - For parameters (b0, b1), Normal and Delta are the same analytic CI here.
    - x50 is on MODEL scale.
    - SUV50 is on RAW scale:
        raw panel -> same as x50
        log panel -> exp(x50)
    """
    import pandas as pd

    def _boot_ci_from_pars(pars, alpha=0.05):
        if pars is None or len(pars) == 0:
            nan2 = (np.nan, np.nan)
            return nan2, nan2, nan2, 0

        pars = np.asarray(pars, float)
        q = np.quantile(pars, [alpha / 2, 0.5, 1.0 - alpha / 2], axis=0)

        b0_ci = (float(q[0, 0]), float(q[2, 0]))
        b1_ci = (float(q[0, 1]), float(q[2, 1]))

        x50s = np.array([x50(bb) for bb in pars], float)
        x50s = x50s[np.isfinite(x50s)]
        if len(x50s) == 0:
            x50_ci = (np.nan, np.nan)
        else:
            xq = np.quantile(x50s, [alpha / 2, 1.0 - alpha / 2])
            x50_ci = (float(xq[0]), float(xq[1]))

        return b0_ci, b1_ci, x50_ci, int(len(pars))

    rows = []

    for key in keys:
        pk = P[key]
        b = np.asarray(pk["b"], float).reshape(2)
        cov = np.asarray(pk["cov"], float).reshape(2, 2)
        trans = pk.get("transform", "")
        l2 = float(pk.get("l2", 0.0))

        lcl, ucl = wald_ci(b, cov, z=z)
        x50_l, x50_u = x50_wald_ci(b, cov, z=z)
        x50_hat = x50(b)

        pars_np = pk.get("pars_nonparam", np.empty((0, 2)))
        pars_pm = pk.get("pars_parametric", np.empty((0, 2)))

        b0_np, b1_np, x50_np, n_np = _boot_ci_from_pars(pars_np)
        b0_pm, b1_pm, x50_pm, n_pm = _boot_ci_from_pars(pars_pm)

        def _to_suv50(x50_ci):
            lo, hi = x50_ci
            if trans == "log":
                return (float(np.exp(lo)), float(np.exp(hi)))
            return (float(lo), float(hi))

        suv50_w = _to_suv50((x50_l, x50_u))
        suv50_np = _to_suv50(x50_np)
        suv50_pm = _to_suv50(x50_pm)

        suv50_hat = float(np.exp(x50_hat)) if trans == "log" else float(x50_hat)

        def add_row(method, b0_ci, b1_ci, x50_ci, suv50_ci, B_used):
            row = {
                "Panel": key,
                "Method": method,
                "b0_LCL": float(b0_ci[0]),
                "b0_UCL": float(b0_ci[1]),
                "b1_LCL": float(b1_ci[0]),
                "b1_UCL": float(b1_ci[1]),
                "x50_LCL": float(x50_ci[0]),
                "x50_UCL": float(x50_ci[1]),
                "SUV50_LCL": float(suv50_ci[0]),
                "SUV50_UCL": float(suv50_ci[1]),
                "B_used": B_used,
            }
            if include_point_est:
                row.update({
                    "b0_hat": float(b[0]),
                    "b1_hat": float(b[1]),
                    "x50_hat": float(x50_hat),
                    "SUV50_hat": float(suv50_hat),
                    "transform": trans,
                    "l2": l2,
                })
            rows.append(row)

        add_row("Normal", (lcl[0], ucl[0]), (lcl[1], ucl[1]), (x50_l, x50_u), suv50_w, np.nan)
        add_row("Delta", (lcl[0], ucl[0]), (lcl[1], ucl[1]), (x50_l, x50_u), suv50_w, np.nan)
        add_row("Nonparam", b0_np, b1_np, x50_np, suv50_np, n_np)
        add_row("Parametric", b0_pm, b1_pm, x50_pm, suv50_pm, n_pm)

    return pd.DataFrame(rows)


# ============================================================
# 13) Elasticity analysis
# ============================================================

def elasticity_x50_slope(b, mode="raw"):
    """
    Elasticities for x50 and slope@50 with respect to b0 and b1.

    mode = 'raw'  : eta = b0 + b1*x
    mode = 'log'  : eta = b0 + b1*log(x)

    Returns dict with:
      x50, slope50,
      E_x50_b0, E_x50_b1,
      E_s50_b0, E_s50_b1
    """
    b0, b1 = map(float, np.asarray(b, float).reshape(2))

    out = {"b0": b0, "b1": b1, "mode": mode}

    if np.abs(b1) < 1e-12:
        out.update({
            "x50": np.nan,
            "slope50": np.nan,
            "E_x50_b0": np.nan,
            "E_x50_b1": np.nan,
            "E_s50_b0": np.nan,
            "E_s50_b1": np.nan,
        })
        return out

    if mode == "raw":
        # eta = b0 + b1*x
        x50 = -b0 / b1
        slope50 = 0.25 * b1

        # elasticities of x50
        if np.abs(x50) < 1e-12:
            E_x50_b0 = np.nan
            E_x50_b1 = np.nan
        else:
            E_x50_b0 = 1.0
            E_x50_b1 = -1.0

        # slope50 = 0.25*b1
        E_s50_b0 = 0.0
        E_s50_b1 = 1.0

    elif mode == "log":
        # eta = b0 + b1*log(x)
        x50 = float(np.exp(-b0 / b1))
        slope50 = 0.25 * b1 / x50

        # elasticities of x50
        E_x50_b0 = -b0 / b1
        E_x50_b1 =  b0 / b1

        # slope elasticity
        # slope50 = 0.25 * b1 * exp(b0/b1)
        if np.abs(slope50) < 1e-12:
            E_s50_b0 = np.nan
            E_s50_b1 = np.nan
        else:
            E_s50_b0 = b0 / b1
            E_s50_b1 = 1.0 - (b0 / b1)

    else:
        raise ValueError("mode must be 'raw' or 'log'")

    out.update({
        "x50": x50,
        "slope50": slope50,
        "E_x50_b0": E_x50_b0,
        "E_x50_b1": E_x50_b1,
        "E_s50_b0": E_s50_b0,
        "E_s50_b1": E_s50_b1,
    })
    return out


def elasticity_table_4panels(
    P,
    keys=("FULL-RAW", "FULL-LOG", "TRIM-RAW", "TRIM-LOG"),
):
    """
    Build a tidy elasticity table for the 4 fitted panels.

    Returns columns:
      Panel, transform, b0, b1, x50, slope50,
      E_x50_b0, E_x50_b1, E_s50_b0, E_s50_b1
    """
    import pandas as pd

    rows = []

    for key in keys:
        pk = P[key]
        b = np.asarray(pk["b"], float).reshape(2)
        transform = pk["transform"]

        res = elasticity_x50_slope(b, mode=transform)

        rows.append({
            "Panel": key,
            "transform": transform,
            "b0": res["b0"],
            "b1": res["b1"],
            "x50": res["x50"],
            "slope50": res["slope50"],
            "E_x50_b0": res["E_x50_b0"],
            "E_x50_b1": res["E_x50_b1"],
            "E_s50_b0": res["E_s50_b0"],
            "E_s50_b1": res["E_s50_b1"],
        })

    return pd.DataFrame(rows)

# ============================================================
# 14. NOISE ANALYSIS (logistic trained on log(X))
# ============================================================


from scipy.optimize import minimize
import matplotlib.pyplot as plt
from matplotlib.patches import Patch


# ------------------------------------------------------------
# logistic helpers
# ------------------------------------------------------------

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


def fit_logistic_logx(x_raw, y):
    """
    Fit logistic model

        p(y=1|x) = sigmoid(b0 + b1 log(x))
    """

    x = np.clip(np.asarray(x_raw).ravel(), 1e-12, None)
    y = np.asarray(y).astype(float)

    X = np.column_stack([np.ones_like(x), np.log(x)])

    def nll(b):
        z = X @ b
        p = logistic_sigmoid(z)
        eps = 1e-12
        ll = np.sum(y*np.log(p+eps) + (1-y)*np.log(1-p+eps))
        return -ll

    res = minimize(nll, np.zeros(2), method="L-BFGS-B")

    if not res.success:
        raise RuntimeError("Logistic optimisation failed")

    return res.x


def predict_curve_logx(b, x_grid):
    x = np.clip(np.asarray(x_grid), 1e-12, None)
    return logistic_sigmoid(b[0] + b[1]*np.log(x))


def x50_from_b(b):
    """
    p(x)=0.5 -> b0 + b1 log(x50)=0
    """
    b0, b1 = b
    return float(np.exp(-b0/b1))


# ------------------------------------------------------------
# noise models
# ------------------------------------------------------------

def add_noise_mult(x, sigma, rng):
    return np.clip(x*np.exp(rng.normal(0, sigma, size=x.shape)),1e-12,None)


def add_noise_add(x, sigma, rng):
    return np.clip(x+rng.normal(0, sigma, size=x.shape),1e-12,None)


# ------------------------------------------------------------
# band quantiles
# ------------------------------------------------------------

def band_quantiles(curves):
    C=np.vstack(curves)
    return np.quantile(C,[0.025,0.5,0.975],axis=0)


# ------------------------------------------------------------
# build noise bands
# ------------------------------------------------------------

def noise_logistic_bands(
        x_raw,
        y,
        sigma_mult=0.129,
        sigma_add=0.144,
        x_max=5,
        grid_n=1000,
        n_refit=200,
        n_tta=500,
        seed=1234
        ):

    rng=np.random.default_rng(seed)

    xc=np.linspace(0,x_max,grid_n)
    xc[0]=1e-12

    b_clean=fit_logistic_logx(x_raw,y)

    clean=predict_curve_logx(b_clean,xc)

    x50=x50_from_b(b_clean)

    # ---------- multiplicative ----------

    curves=[]
    for _ in range(n_refit):
        xn=add_noise_mult(x_raw,sigma_mult,rng)
        bn=fit_logistic_logx(xn,y)
        curves.append(predict_curve_logx(bn,xc))

    mult_refit=band_quantiles(curves)

    curves=[]
    for _ in range(n_tta):
        xn=add_noise_mult(xc,sigma_mult,rng)
        curves.append(predict_curve_logx(b_clean,xn))

    mult_tta=band_quantiles(curves)

    # ---------- additive ----------

    curves=[]
    for _ in range(n_refit):
        xn=add_noise_add(x_raw,sigma_add,rng)
        bn=fit_logistic_logx(xn,y)
        curves.append(predict_curve_logx(bn,xc))

    add_refit=band_quantiles(curves)

    curves=[]
    for _ in range(n_tta):
        xn=add_noise_add(xc,sigma_add,rng)
        curves.append(predict_curve_logx(b_clean,xn))

    add_tta=band_quantiles(curves)

    return dict(
        xc=xc,
        clean=clean,
        x50=x50,
        mult_refit=mult_refit,
        mult_tta=mult_tta,
        add_refit=add_refit,
        add_tta=add_tta
    )


# ------------------------------------------------------------
# plotting
# ------------------------------------------------------------
def plot_noise_panel(ax, pack, kind="mult", label="A", X=None, y=None):

    COL_MULT = "#1f78b4"
    COL_ADD  = "#e66101"

    xc = pack["xc"]
    clean = pack["clean"]
    x50 = pack["x50"]

    if kind == "mult":
        refit = pack["mult_refit"]
        tta = pack["mult_tta"]
        color = COL_MULT
    else:
        refit = pack["add_refit"]
        tta = pack["add_tta"]
        color = COL_ADD

    lo_r, _, hi_r = refit
    lo_t, _, hi_t = tta

    # TTA band (lighter)
    ax.fill_between(xc, lo_t, hi_t, color=color, alpha=0.10, zorder=1)

    # refit band (stronger)
    ax.fill_between(xc, lo_r, hi_r, color=color, alpha=0.24, zorder=2)

    # optional thin outlines for readability
    ax.plot(xc, lo_r, color=color, lw=1.0, alpha=0.65, zorder=3)
    ax.plot(xc, hi_r, color=color, lw=1.0, alpha=0.65, zorder=3)

    # clean fit
    ax.plot(xc, clean, color="black", lw=2.5, zorder=5)

    # x50 reference line
    # x50 reference line
    ax.axvline(x50, color="#666666", ls="--", lw=1.4, alpha=0.9, zorder=4)

    # widths at x50
    lo_r_x = np.interp(x50, xc, lo_r)
    hi_r_x = np.interp(x50, xc, hi_r)
    lo_t_x = np.interp(x50, xc, lo_t)
    hi_t_x = np.interp(x50, xc, hi_t)

    
    # data points
    if X is not None and y is not None:
        X = np.asarray(X).ravel()
        y = np.asarray(y).astype(int)

        ax.scatter(
            X[y == 0], np.zeros(np.sum(y == 0)),
            color="#2b8cbe", s=28, alpha=0.75, zorder=7
        )
        ax.scatter(
            X[y == 1], np.ones(np.sum(y == 1)),
            color="#d7301f", s=28, alpha=0.75, zorder=7
        )

    # panel label
    ax.text(
        0.50, 1.01, label,
        transform=ax.transAxes,
        ha="center", va="bottom",
        fontsize=18, fontweight="bold"
    )

    # only FULL/TRIM + x50
    variant_txt = "FULL (F)" if label in ["A", "B"] else "TRIM (T)"
    ax.text(0.02, 0.90, variant_txt, transform=ax.transAxes, fontsize=11, color="#111")
    ax.text(0.02, 0.82, f"x50={x50:.2f}", transform=ax.transAxes, fontsize=10, color="#111")

    # delta text box
    d_ref = hi_r_x - lo_r_x
    d_tta = hi_t_x - lo_t_x
    ax.text(
    0.95, 0.06,
    f"Δr={d_ref:.2f}  Δt={d_tta:.2f}",
    transform=ax.transAxes,
    fontsize=10,
    color="#222",
    ha="right",
    bbox=dict(facecolor="white", edgecolor=color, boxstyle="square,pad=0.2", alpha=0.9)
)
    ax.set_xlim(0, xc.max())
    ax.set_ylim(-0.05, 1.05)
    ax.grid(alpha=0.25)
    ax.tick_params(axis="x", labelbottom=True)
# ------------------------------------------------------------
# legend
# ------------------------------------------------------------
def noise_legend(fig):

    import matplotlib.lines as mlines
    from matplotlib.patches import Patch

    legend_handles = [

        mlines.Line2D(
            [0], [0],
            color="black",
            lw=2.5,
            label="clean logistic fit (trained on log(SUV))"
        ),

        Patch(
            facecolor="#999999",
            alpha=0.24,
            edgecolor="none",
            label="refit band (training perturbation; 95% CI)"
        ),

        Patch(
            facecolor="#999999",
            alpha=0.10,
            edgecolor="none",
            label="TTA band (inference-time noise; 95% CI)"
        ),

        Patch(
            facecolor="#1f78b4",
            alpha=0.24,
            edgecolor="none",
            label="multiplicative noise: σ=0.129 (blue)"
        ),

        Patch(
            facecolor="#e66101",
            alpha=0.24,
            edgecolor="none",
            label="additive noise: σ=0.144 (orange)"
        ),

        mlines.Line2D(
            [0], [0],
            color="#666666",
            lw=1.4,
            ls="--",
            label="x50 (P=0.5)"
        ),
    ]

    fig.legend(
        handles=legend_handles,
        labels=[h.get_label() for h in legend_handles],
        loc="lower center",
        ncol=3,
        frameon=False,
        bbox_to_anchor=(0.5, 0.02),
        fontsize=10
    )