# bayesianconstraints.py
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


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


def logit(p):
    p = np.clip(p, 1e-12, 1 - 1e-12)
    return np.log(p) - np.log(1 - p)


def dE_gamma(x, a, b, k, theta, s=1.0):
    x = np.asarray(x, float)
    return (a - k) * np.log(x) - (a + b) * np.log1p(x / s) + x / theta


def dE_lognorm(x, a, b, mu, sigma, s=1.0):
    x = np.asarray(x, float)
    return ((mu - np.log(x)) ** 2) / (2.0 * sigma ** 2) + a * np.log(x) - (a + b) * np.log1p(x / s)


def check_gamma_constraints(a, b, k, theta, s):
    positivity = (a > 0) and (b > 0) and (k > 0) and (theta > 0) and (s > 0)

    if not positivity:
        status = "FAIL"
        note = "all parameters must be > 0"
    elif a > k:
        status = "MEET"
        note = "a > k"
    elif a == k:
        status = "EDGE"
        note = "a = k"
    else:
        status = "FAIL"
        note = "a < k"

    return {
        "model": "BetaPrime vs Gamma",
        "a": a,
        "b": b,
        "k": k,
        "theta": theta,
        "s": s,
        "all_positive": positivity,
        "a_ge_k": positivity and (a >= k),
        "status": status,
        "note": note,
    }


def check_lognorm_constraints(a, b, mu, sigma, s):
    positivity = (a > 0) and (b > 0) and (sigma > 0) and (s > 0)

    if not positivity:
        status = "FAIL"
        note = "a,b,sigma,s must be > 0"
    else:
        status = "FAIL"
        note = "left tail goes to 1, not 0"

    return {
        "model": "BetaPrime vs LogNormal",
        "a": a,
        "b": b,
        "mu": mu,
        "sigma": sigma,
        "s": s,
        "all_positive": positivity,
        "status": status,
        "note": note,
    }


def make_gamma_constraint_table(param_list):
    return pd.DataFrame([check_gamma_constraints(**p) for p in param_list])


def make_lognorm_constraint_table(param_list):
    return pd.DataFrame([check_lognorm_constraints(**p) for p in param_list])


def plot_gamma_constraint_curves(param_meet, param_edge_fail, p_prior=5/58,
                                 x=None, figsize=(10, 6)):
    if x is None:
        x = np.logspace(-12, 3, 900)

    c0 = logit(p_prior)
    fig, ax = plt.subplots(figsize=figsize)

    for P in param_meet:
        p = logistic(c0 + dE_gamma(x, **P))
        label = f"MEET (a={P['a']}, b={P['b']}, k={P['k']}, θ={P['theta']}, s={P['s']})"
        ax.plot(x, p, label=label)

    for P in param_edge_fail:
        p = logistic(c0 + dE_gamma(x, **P))
        tag = "EDGE" if P["a"] == P["k"] else "FAIL"
        label = f"{tag} (a={P['a']}, b={P['b']}, k={P['k']}, θ={P['theta']}, s={P['s']})"
        ax.plot(x, p, linestyle="--", label=label)

    ax.set_xscale("log")
    ax.set_ylim(-0.05, 1.05)
    ax.set_xlabel("x")
    ax.set_ylabel("P(AE | x)")
    ax.set_title("Gamma NC — posterior-like curves")
    ax.grid(True, which="both")
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig, ax


def plot_lognorm_constraint_curves(param_list, p_prior=5/58,
                                   x=None, figsize=(10, 6)):
    if x is None:
        x = np.logspace(-12, 3, 900)

    c0 = logit(p_prior)
    fig, ax = plt.subplots(figsize=figsize)

    for P in param_list:
        p = logistic(c0 + dE_lognorm(x, **P))
        label = f"(a={P['a']}, b={P['b']}, μ={P['mu']}, σ={P['sigma']}, s={P['s']})"
        ax.plot(x, p, label=label)

    ax.set_xscale("log")
    ax.set_ylim(-0.05, 1.05)
    ax.set_xlabel("x")
    ax.set_ylabel("P(AE | x)")
    ax.set_title("LogNormal NC — posterior-like curves")
    ax.grid(True, which="both")
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig, ax