# Constrained Bayesian fit (Gamma for NC, Beta-Prime for AE) — no regularization, no r-prior
import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize
from scipy.special import betaln, gammaln
import scipy.io as io

data_path = "../data/"
suv   = io.loadmat(data_path + "suv_percentilesSLOthenUWM.mat")['lung_SUVperc_COMBINED'][0:58, :, :]
flags = io.loadmat(data_path + "flags_combined.mat")['flags'][0:58, 3]   # 0=NC, 1=AE

# Feature X = max SUV_94 per subject; label y = flags
X = np.nanmax(suv[:, :, 94], axis=1).astype(float).ravel()
y = np.asarray(flags, int).ravel()

# Guard for logs
X = np.clip(X, 1e-12, None)
p_emp = float(y.mean())

# Small helpers

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

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

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

# dE(x) pieces for log-odds
def dE_ess(x, a, b, s, k, th):
    x = np.asarray(x, float)
    return (a - k) * np.log(x) - (a + b) * np.log1p(x / s) + x / th

def dE_const(a, b, s, k, th):
    return -(a * np.log(s)) - betaln(a, b) + k * np.log(th) + gammaln(k)

def dE_full(x, a, b, s, k, th):
    return dE_ess(x, a, b, s, k, th) + dE_const(a, b, s, k, th)

# Monotonicity cap for theta
def theta_max(a, b, k, s, eps=1e-12):
    A = a - k
    if A <= 0:
        return np.inf
    r = np.sqrt(a + b) - np.sqrt(max(A, eps))
    if r <= 1e-12:
        return np.inf
    return s / (r * r)

# φ = [p_raw, b_raw, s_raw, k_raw, d_raw, u_raw]  (unconstrained)
def unpack_phi_mono(phi):
    p_raw, b_raw, s_raw, k_raw, d_raw, u_raw = phi
    p = sigmoid(p_raw)                       # (0,1)
    b = softplus(b_raw) + 1e-6               # >0
    s = softplus(s_raw) + 1e-6               # >0
    k = softplus(k_raw) + 1e-6               # >0
    delta = softplus(d_raw) + 1e-6           # >0
    a = k + delta                            # enforce a > k
    th_cap = theta_max(a, b, k, s)           # theta cap from monotonicity
    th = th_cap * sigmoid(u_raw)             # 0 < theta <= th_cap
    return p, a, b, s, k, th



# Prior on r = theta / theta_max (softly avoid boundaries)

thcap = theta_max(a, b, k, s)
if np.isfinite(thcap) and thcap > 0:
    r = np.clip(th / thcap, 1e-9, 1 - 1e-9)
    # Negative log Beta prior:  -[(α-1)log r + (β-1)log(1-r)]  (const dropped)
    npr_r = -((0.5) * np.log(r) + (0.5) * np.log(1.0 - r))
else:
    npr_r = 0.0

#  Prior on p 
TAU = 25.0                                   # shrink toward empirical AE rate
alpha = max(TAU * p_emp, 1e-6)
beta  = max(TAU * (1.0 - p_emp), 1e-6)

# Objective: negative log-posterior (likelihood + Beta prior on p)
def neg_post_phi_mono(phi, X, y):
    p, a, b, s, k, th = unpack_phi_mono(phi)
    eps = 1e-12

    L  = (np.log(p) - np.log(1 - p)) + dE_full(X, a, b, s, k, th)
    px = logistic(L)
    nll = -np.sum(y * np.log(px + eps) + (1 - y) * np.log(1 - px + eps))

    # Beta(alpha, beta) prior on p → negative log-prior
    npr_p = -((alpha - 1) * np.log(p + eps) + (beta - 1) * np.log(1 - p + eps))

    return nll + npr_r

# Initialization (stable, simple)
def init_phi(X, y):
    # Gamma(k, theta) MoM for NC group (only need k0 as a safe size proxy)
    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**2)/(v0 + 1e-9), 1.5)

    # AE median to seed s0
    X1 = X[y == 1]
    m1 = np.median(X1) if X1.size else np.median(X)

    p0 = np.clip(float(y.mean()), 1e-3, 1 - 1e-3)
    b0, s0 = 1.5, 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  (delta)
        -0.2                               # u_raw  (keeps theta a bit below cap initially)
    ], float)

#  Fit wrapper (one retry)
def fit_bayes_mono(X, y, phi_start=None, rng=None):
    if rng is None:
        rng = np.random.default_rng(0)
    if phi_start is None:
        phi_start = init_phi(X, y)

    obj = lambda phi: neg_post_phi_mono(phi, X, y)

    res = optimize.minimize(
        obj, phi_start, method="L-BFGS-B",
        options={"maxiter": 6000, "ftol": 1e-9}
    )
    if not (res.success and np.isfinite(res.fun)):
        phi_try = phi_start + rng.normal(0, 0.2, size=phi_start.shape)
        res = optimize.minimize(
            obj, phi_try, method="L-BFGS-B",
            options={"maxiter": 6000, "ftol": 1e-9}
        )
    return unpack_phi_mono(res.x), res

#  prediction 
def P_with(theta, x):
    p, a, b, s, k, th = theta
    L = (np.log(p) - np.log(1 - p)) + dE_full(x, a, b, s, k, th)
    return logistic(L)

# Run fit + plot
theta_hat, res = fit_bayes_mono(X, y)
print("Optimization success:", res.success, " fval:", float(res.fun))
print("theta (p,a,b,s,k,theta):", tuple(float(t) for t in theta_hat))

# x-range (cap right end at 10 for readability)
x_lo = max(1e-6, float(X.min()) * 0.8)
x_hi = min(10.0, float(X.max()) * 1.2)
xg   = np.linspace(x_lo, x_hi, 600)
p_curve = P_with(theta_hat, xg)

fig, ax = plt.subplots(figsize=(7.0, 4.6), dpi=140)
ax.plot(xg, p_curve, color="#000000", lw=2.2, label="P(AE|x) (MAP)")

# overlay data with tiny vertical jitter so points don't overlap
rng_plot = np.random.default_rng(999)
jit = (rng_plot.random(len(y)) - 0.5) * 0.06
ax.scatter(X[y==0], (y + jit)[y==0], s=22, alpha=0.55, color="#2ca02c", edgecolors='none', label='NC')
ax.scatter(X[y==1], (y + jit)[y==1], s=26, alpha=0.75, color="#ff7f0e", edgecolors='none', label='AE')

ax.set_ylim(-0.05, 1.05)
ax.set_xlabel('x')
ax.set_ylabel('P(AE | x)')
ax.set_title('Constrained Bayesian fit (no regularization  prior on p and r:alpha and beta=1.5')
ax.grid(alpha=0.3)
ax.legend(loc='lower right', frameon=False)
plt.tight_layout()
plt.show()