"""
    Working on Bayesian formula model + penalty  == MAP approach with prior = penalty
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize, stats

#here is data
assert np.all(X > 0), "All X must be > 0"
n1, n0 = int(y.sum()), int((1-y).sum())

m_emp = n1 / (n1 + n0)  # approx Prob(AE)

print(f"AE={n1}, NC={n0}, empirical p_AE={m_emp:.6f}")

# Model: AE ~ BetaPrime(a,b,scale), NC ~ LogNormal(mu, sigma)
# p_ae = Prob(AE)
# θ = [p_ae, a, b, sc, mu_nc, sig_nc]
def logpdf_betaprime(x, a, b, scale):
    return stats.betaprime.logpdf(x, a=a, b=b, scale=scale)

def logpdf_lognorm(x, mu, sigma):
    return stats.lognorm.logpdf(x, s=sigma, scale=np.exp(mu))

def neg_conditional_ll(theta, X, y, eps=1e-12):
    
    p_ae, a, b, sc, mu_nc, sig_nc = theta

    logf1 = logpdf_betaprime(X, a, b, sc)     # AE
    logf0 = logpdf_lognorm(X, mu_nc, sig_nc)  # NC

    # computing p(AE|x) = 1/(1 + exp(-logit))
    logit = np.log(p_ae) - np.log(1.0 - p_ae) + (logf1 - logf0)
    p = 1.0 / (1.0 + np.exp(-np.clip(logit, -50, 50)))

    # computing general nllf = -llf
    #  llf = sum_i log(p(y_i|x_i)) 
    #      = sum_i y_i log( p(AE|x) + (1- y_i) log(1 - p(AE|x));  p(NC|x) = 1 -p(AE|x)

    nllf = -np.sum(y*np.log(p+eps) + (1-y)*np.log(1-p+eps))
    return nllf

"""
    MAP regularization
        mean = m_emp, concentration = TAU
    using beta distribution B(alpha, beta). The parameters are set as
        alpha = m_emp * TAU, beta = (1-m_emp) * TAU  
    this yields E[X] = alpha/(alpha+beta) = m_emp
        Var[X] = m_emp(1- m_emp)/(tau +1)  
    this could be variance between institutions collected by Katja, 
    my estimate is 
        var = 5%^2 
    This yields tau about 30. 
"""
TAU = 100.0            # ↑ increase to pull p_AE closer to empirical prior
alpha = max(m_emp * TAU, 1e-6)
beta  = max((1 - m_emp) * TAU, 1e-6)


def neg_log_prior_p(p, eps=1e-12):
    # -log Beta(p | alpha, beta) up to a constant
    return -( (alpha - 1)*np.log(p + eps) + (beta - 1)*np.log(1 - p + eps) )


def neg_posterior(theta, X, y):
    nll = neg_conditional_ll(theta, X, y)
    return nll + neg_log_prior_p(theta[0])   # add prior penalty on p_AE only


# Bounds (optionally enforce heavy AE tail with b ≤ 1)
# =========================
HEAVY_TAIL = True     # set False if you don't want to force the right asymptote
b_upper = 1.0 if HEAVY_TAIL else 50.0
bounds = [
    (1e-3, 1-1e-3),                           # p_ae (free, but regularized by Beta prior)
    (0.20, 50.0),                             # a (AE BetaPrime)
    (0.20, b_upper),                          # b (AE BetaPrime)  <-- heavy tail if ≤ 1
    (0.01, 10.0),                             # scale (AE BetaPrime)
    (np.log(X).min()-2.0, np.log(X).max()+2.0),  # mu_nc (LogNormal)
    (0.05, 2.0),                              # sigma_nc (LogNormal)
]
# Initialization
p0 = np.clip(m_emp, bounds[0][0], bounds[0][1])
X0 = X[y==0]
mu0 = float(np.mean(np.log(X0)))
sig0 = float(np.std(np.log(X0), ddof=0))
mu0 = np.clip(mu0, bounds[4][0], bounds[4][1])
sig0 = np.clip(sig0, bounds[5][0], bounds[5][1])
X1 = X[y==1]
m1 = float(np.mean(X1))
a0, b0 = 2.5, min(0.8, b_upper)  # start with heavy-tail-ish b if allowed
sc0 = np.clip(m1 * (b0 - 1 + 1e-6) / max(a0, 1e-6), bounds[3][0], bounds[3][1])
theta0 = np.array([p0, a0, b0, sc0, mu0, sig0], dtype=float)
# FREE prior (for comparison)
# =========================
res_free = optimize.minimize(
    fun=neg_conditional_ll,
    x0=theta0,
    args=(X, y),
    method="L-BFGS-B",
    bounds=bounds,
    options=dict(maxiter=4000, ftol=1e-12)
)
theta_free = res_free.x
print("\n[FREE prior] p_AE =", float(theta_free[0]), "   CLL =", -res_free.fun)
# MAP prior (regularized toward empirical)
# =========================
res_map = optimize.minimize(
    fun=neg_posterior,
    x0=theta0,
    args=(X, y),
    method="L-BFGS-B",
    bounds=bounds,
    options=dict(maxiter=4000, ftol=1e-12)
)
theta_map = res_map.x
print("[MAP prior]  p_AE =", float(theta_map[0]), "   CLL(post) =", -res_map.fun)
# MAP parameters
p_hat, a_hat, b_hat, sc_hat, mu_hat, sig_hat = theta_map
print("\nFitted (MAP) parameters:")
print(f"  p_AE = {p_hat:.6f}  (empirical {m_emp:.6f}, TAU={TAU})")
print(f"  AE BetaPrime: a={a_hat:.4f}, b={b_hat:.4f}, scale={sc_hat:.4f}")
print(f"  NC LogNormal: mu={mu_hat:.4f}, sigma={sig_hat:.4f}")
# Posterior & plot
def predict_proba(x):
    x = np.asarray(x, dtype=float)
    logf1 = logpdf_betaprime(x, a_hat, b_hat, sc_hat)
    logf0 = logpdf_lognorm(x, mu_hat, sig_hat)
    logit = np.log(p_hat) - np.log(1.0 - p_hat) + (logf1 - logf0)
    return 1.0 / (1.0 + np.exp(-np.clip(logit, -50, 50)))
x_grid = np.linspace(max(1e-6, X.min()*0.6), max(X.max()*2.0, 8.0), 600)
p_grid = predict_proba(x_grid)
plt.figure(figsize=(8,5))
plt.plot(x_grid, p_grid, 'k-', linewidth=2, label="P(AE | X) [MAP]")
plt.scatter(X[y==1], np.ones(n1), marker='x', label="AE samples")
plt.scatter(X[y==0], np.zeros(n0), marker='o', label="NC samples")
plt.xlabel("SUV feature X"); plt.ylabel("Predicted P(AE | X)")
plt.title("BetaPrime–LogNormal with MAP prior on p_AE")
plt.ylim(-0.05, 1.05); plt.legend(); plt.grid(True)
plt.show() (edited)