"""
extreme_survival.py -- From-scratch EXTREME-VALUE, SURVIVAL & RELIABILITY distributions.

Backs the notebooks in  Z-Statistical Distributions-ExtremeValue-Survival  (Folder 4
of the Statistical Distributions catalog).

Two themes meet here:
  * EXTREME VALUE -- the laws of maxima (Gumbel, Frechet, Weibull) unified by the
    Generalized Extreme Value (GEV) distribution, and the peaks-over-threshold law,
    the Generalized Pareto (GPD).
  * SURVIVAL / RELIABILITY -- time-to-failure laws described by the HAZARD function
    h(x)=f(x)/S(x) with survival S(x)=1-F(x): Weibull (bathtub), log-logistic
    (non-monotone hazard), Gompertz-Makeham (mortality), and the actuarial loss /
    income family GB2 with its Burr and Dagum special cases.

Everything is implemented from the definitions using only elementary math and a few
special functions as primitives (the log-gamma / log-beta and the regularized
incomplete beta). No distribution objects. `scipy.stats` (Python) and R (R notebook)
are used as INDEPENDENT references.

Per distribution, where defined:
    *_pdf, *_cdf, *_rng(..., n, rng), *_moments   (NaN moments where they do not exist)
    survival S(x) and hazard h(x) via the generic helpers.

Parameterisations chosen to match a scipy reference where one exists:
    gumbel(mu, beta)                 scipy gumbel_r(mu, beta)
    weibull(k, lam)                  scipy weibull_min(k, scale=lam)
    frechet(alpha, s, m)             scipy invweibull(alpha, loc=m, scale=s)
    gev(mu, sigma, xi)               scipy genextreme(c=-xi, loc=mu, scale=sigma)   [note the sign]
    gpd(mu, sigma, xi)               scipy genpareto(xi, loc=mu, scale=sigma)
    loglogistic(alpha, beta)         scipy fisk(beta, scale=alpha)
    gompertz(eta, theta)             scipy gompertz(c=eta/theta, scale=1/theta)
    burr12(c, k, b)                  scipy burr12(c, k, scale=b)     (Singh-Maddala)
    dagum(a, b, p)                   scipy burr(a, p, scale=b)       (Burr III)
    gb2(a, b, p, q)                  no scipy; validated via Burr/Dagum/log-logistic special cases
    makeham(lam, eta, theta)         no scipy; reduces to Gompertz at lam=0
"""

import numpy as np
from scipy.special import gammaln, betaln, betainc, gamma as Gamma

EULER = 0.5772156649015329          # Euler-Mascheroni constant
_trapz = getattr(np, "trapezoid", getattr(np, "trapz", None))


# --------------------------------------------------------------------------- #
#  Generic helpers                                                              #
# --------------------------------------------------------------------------- #

def survival(cdf_vals):
    return 1.0 - np.asarray(cdf_vals, float)

def hazard(pdf_vals, cdf_vals):
    """h(x) = f(x) / S(x) -- the instantaneous failure rate."""
    S = 1.0 - np.asarray(cdf_vals, float)
    return np.asarray(pdf_vals, float) / np.where(S > 1e-300, S, 1e-300)

def emp_moments(x):
    x = np.asarray(x, float); m = x.mean(); c2 = x.var()
    return dict(mean=m, var=c2, skew=((x - m) ** 3).mean() / c2 ** 1.5,
                exkurt=((x - m) ** 4).mean() / c2 ** 2 - 3)

def _central_from_raw(m1, m2, m3, m4):
    var = m2 - m1 ** 2
    c3 = m3 - 3 * m1 * m2 + 2 * m1 ** 3
    c4 = m4 - 4 * m1 * m3 + 6 * m1 ** 2 * m2 - 3 * m1 ** 4
    sd = np.sqrt(var)
    return dict(mean=m1, var=var, skew=c3 / sd ** 3, exkurt=c4 / var ** 2 - 3)

def moments_from_pdf(pdf, lo, hi, n=800001):
    x = np.linspace(lo, hi, n); px = np.asarray(pdf(x), float)
    px = px / _trapz(px, x)
    m1 = _trapz(x * px, x); c2 = _trapz((x - m1) ** 2 * px, x)
    c3 = _trapz((x - m1) ** 3 * px, x); c4 = _trapz((x - m1) ** 4 * px, x)
    sd = np.sqrt(c2)
    return dict(mean=m1, var=c2, skew=c3 / sd ** 3, exkurt=c4 / c2 ** 2 - 3)

def inverse_transform_sample(cdf, lo, hi, n, rng, npts=200001):
    """Sample a distribution given only its CDF: invert F on a fine grid."""
    g = np.linspace(lo, hi, npts)
    F = np.asarray(cdf(g), float)
    F, idx = np.unique(F, return_index=True)          # enforce strict monotonicity
    return np.interp(rng.random(n), F, g[idx])


# ===========================================================================  #
#  PART A -- EXTREME VALUE (laws of maxima)                                     #
# ===========================================================================  #

# --- Gumbel(mu, beta) : Type I, light-tailed maxima -------------------------- #
def gumbel_pdf(x, mu, beta):
    z = (np.asarray(x, float) - mu) / beta
    return np.exp(-(z + np.exp(-z))) / beta

def gumbel_cdf(x, mu, beta):
    z = (np.asarray(x, float) - mu) / beta
    return np.exp(-np.exp(-z))

def gumbel_rng(mu, beta, n, rng):
    """Inverse transform: mu - beta*ln(-ln U)."""
    return mu - beta * np.log(-np.log(rng.random(n)))

def gumbel_moments(mu, beta):
    from scipy.special import zeta
    return dict(mean=mu + beta * EULER, var=np.pi ** 2 / 6 * beta ** 2,
                skew=12 * np.sqrt(6) * zeta(3) / np.pi ** 3, exkurt=12.0 / 5)


# --- Weibull(k, lam) : Type III / reliability workhorse ---------------------- #
def weibull_pdf(x, k, lam):
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1.0)
    return np.where(x > 0, (k / lam) * (xp / lam) ** (k - 1) * np.exp(-(xp / lam) ** k), 0.0)

def weibull_cdf(x, k, lam):
    x = np.asarray(x, float)
    return np.where(x > 0, 1 - np.exp(-(np.maximum(x, 0) / lam) ** k), 0.0)

def weibull_rng(k, lam, n, rng):
    """Inverse transform: lam * (-ln U)^{1/k}."""
    return lam * (-np.log(rng.random(n))) ** (1.0 / k)

def weibull_hazard(x, k, lam):
    x = np.asarray(x, float)
    return (k / lam) * (np.maximum(x, 0) / lam) ** (k - 1)     # increasing k>1, constant k=1, decreasing k<1

def weibull_moments(k, lam):
    g1 = Gamma(1 + 1 / k); g2 = Gamma(1 + 2 / k); g3 = Gamma(1 + 3 / k); g4 = Gamma(1 + 4 / k)
    return _central_from_raw(lam * g1, lam ** 2 * g2, lam ** 3 * g3, lam ** 4 * g4)


# --- Frechet(alpha, s, m) : Type II, heavy-tailed maxima --------------------- #
def frechet_pdf(x, alpha, s, m):
    z = (np.asarray(x, float) - m) / s; zp = np.where(z > 0, z, 1.0)
    return np.where(z > 0, (alpha / s) * zp ** (-1 - alpha) * np.exp(-zp ** (-alpha)), 0.0)

def frechet_cdf(x, alpha, s, m):
    z = (np.asarray(x, float) - m) / s
    return np.where(z > 0, np.exp(-np.where(z > 0, z, 1.0) ** (-alpha)), 0.0)

def frechet_rng(alpha, s, m, n, rng):
    """Inverse transform: m + s*(-ln U)^{-1/alpha}."""
    return m + s * (-np.log(rng.random(n))) ** (-1.0 / alpha)

def frechet_moments(alpha, s, m):
    g = lambda k: Gamma(1 - k / alpha) if alpha > k else np.nan
    if not (alpha > 1):
        return dict(mean=np.nan, var=np.nan, skew=np.nan, exkurt=np.nan)
    m1 = m + s * g(1)
    m2 = s ** 2 * g(2) + 2 * m * s * g(1) + m ** 2 if alpha > 2 else np.nan
    if alpha > 4:
        # central moments about the (shifted) Frechet; use raw moments of s*Z + m, Z standard Frechet
        z1, z2, z3, z4 = g(1), g(2), g(3), g(4)
        r1 = m + s * z1; r2 = m ** 2 + 2 * m * s * z1 + s ** 2 * z2
        r3 = m ** 3 + 3 * m ** 2 * s * z1 + 3 * m * s ** 2 * z2 + s ** 3 * z3
        r4 = m ** 4 + 4 * m ** 3 * s * z1 + 6 * m ** 2 * s ** 2 * z2 + 4 * m * s ** 3 * z3 + s ** 4 * z4
        return _central_from_raw(r1, r2, r3, r4)
    return dict(mean=m1, var=(m2 - m1 ** 2 if alpha > 2 else np.nan), skew=np.nan, exkurt=np.nan)


# --- GEV(mu, sigma, xi) : the unifying block-maxima law ---------------------- #
def gev_pdf(x, mu, sigma, xi):
    z = (np.asarray(x, float) - mu) / sigma
    if abs(xi) < 1e-12:
        t = np.exp(-z)
        return t * np.exp(-t) / sigma
    base = 1 + xi * z
    ok = base > 0
    tb = np.where(ok, base, 1.0) ** (-1.0 / xi)
    return np.where(ok, (1.0 / sigma) * np.where(ok, base, 1.0) ** (-1.0 / xi - 1) * np.exp(-tb), 0.0)

def gev_cdf(x, mu, sigma, xi):
    z = (np.asarray(x, float) - mu) / sigma
    if abs(xi) < 1e-12:
        return np.exp(-np.exp(-z))
    base = 1 + xi * z
    ok = base > 0
    out = np.exp(-np.where(ok, base, 1.0) ** (-1.0 / xi))
    # below/above the support endpoint the CDF is 0 or 1 depending on sign of xi
    return np.where(ok, out, np.where(xi > 0, 0.0, 1.0))

def gev_rng(mu, sigma, xi, n, rng):
    U = rng.random(n)
    if abs(xi) < 1e-12:
        return mu - sigma * np.log(-np.log(U))
    return mu + sigma * ((-np.log(U)) ** (-xi) - 1) / xi

def gev_moments(mu, sigma, xi):
    if abs(xi) < 1e-12:
        return gumbel_moments(mu, sigma)
    if xi >= 1:
        return dict(mean=np.nan, var=np.nan, skew=np.nan, exkurt=np.nan)
    g = lambda k: Gamma(1 - k * xi)
    g1, g2, g3, g4 = g(1), g(2), g(3), g(4)
    mean = mu + sigma * (g1 - 1) / xi
    if xi >= 0.5:
        return dict(mean=mean, var=np.nan, skew=np.nan, exkurt=np.nan)
    var = sigma ** 2 * (g2 - g1 ** 2) / xi ** 2
    if xi >= 1 / 3:
        return dict(mean=mean, var=var, skew=np.nan, exkurt=np.nan)
    sgn = np.sign(xi)
    skew = sgn * (g3 - 3 * g1 * g2 + 2 * g1 ** 3) / (g2 - g1 ** 2) ** 1.5
    if xi >= 0.25:
        return dict(mean=mean, var=var, skew=skew, exkurt=np.nan)
    exkurt = (g4 - 4 * g1 * g3 + 6 * g1 ** 2 * g2 - 3 * g1 ** 4) / (g2 - g1 ** 2) ** 2 - 3
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# --- GPD(mu, sigma, xi) : peaks over threshold ------------------------------- #
def gpd_pdf(x, mu, sigma, xi):
    z = (np.asarray(x, float) - mu) / sigma
    if abs(xi) < 1e-12:
        return np.where(z >= 0, np.exp(-z) / sigma, 0.0)
    base = 1 + xi * z
    ok = (z >= 0) & (base > 0)
    return np.where(ok, (1.0 / sigma) * np.where(ok, base, 1.0) ** (-1.0 / xi - 1), 0.0)

def gpd_cdf(x, mu, sigma, xi):
    z = (np.asarray(x, float) - mu) / sigma
    if abs(xi) < 1e-12:
        return np.where(z >= 0, 1 - np.exp(-z), 0.0)
    base = 1 + xi * z
    ok = (z >= 0) & (base > 0)
    out = 1 - np.where(ok, base, 1.0) ** (-1.0 / xi)
    return np.where(z < 0, 0.0, np.where(base > 0, out, 1.0))

def gpd_rng(mu, sigma, xi, n, rng):
    U = rng.random(n)
    if abs(xi) < 1e-12:
        return mu - sigma * np.log(U)
    return mu + sigma * (U ** (-xi) - 1) / xi

def gpd_moments(mu, sigma, xi):
    mean = mu + sigma / (1 - xi) if xi < 1 else np.nan
    var = sigma ** 2 / ((1 - xi) ** 2 * (1 - 2 * xi)) if xi < 0.5 else np.nan
    skew = (2 * (1 + xi) * np.sqrt(1 - 2 * xi) / (1 - 3 * xi)) if xi < 1 / 3 else np.nan
    exkurt = (3 * (1 - 2 * xi) * (2 * xi ** 2 + xi + 3) / ((1 - 3 * xi) * (1 - 4 * xi)) - 3) if xi < 0.25 else np.nan
    return dict(mean=mean, var=var, skew=skew, exkurt=exkurt)


# ===========================================================================  #
#  PART B -- SURVIVAL / RELIABILITY (hazard-driven)                             #
# ===========================================================================  #

# --- Log-logistic / Fisk(alpha scale, beta shape) : non-monotone hazard ------ #
def loglogistic_pdf(x, alpha, beta):
    x = np.asarray(x, float); z = np.where(x > 0, x, 1.0) / alpha
    return np.where(x > 0, (beta / alpha) * z ** (beta - 1) / (1 + z ** beta) ** 2, 0.0)

def loglogistic_cdf(x, alpha, beta):
    x = np.asarray(x, float); z = np.where(x > 0, x, 1.0) / alpha
    return np.where(x > 0, 1.0 / (1 + z ** (-beta)), 0.0)

def loglogistic_rng(alpha, beta, n, rng):
    """Inverse transform: alpha*(U/(1-U))^{1/beta}."""
    U = rng.random(n)
    return alpha * (U / (1 - U)) ** (1.0 / beta)

def loglogistic_hazard(x, alpha, beta):
    return hazard(loglogistic_pdf(x, alpha, beta), loglogistic_cdf(x, alpha, beta))

def loglogistic_moments(alpha, beta):
    def raw(k):
        if beta <= k: return np.nan
        b = k * np.pi / beta
        return alpha ** k * b / np.sin(b)
    r1, r2, r3, r4 = raw(1), raw(2), raw(3), raw(4)
    if np.isnan(r2): return dict(mean=r1, var=np.nan, skew=np.nan, exkurt=np.nan)
    if np.isnan(r4): return dict(mean=r1, var=r2 - r1 ** 2, skew=(np.nan if np.isnan(r3) else None), exkurt=np.nan)
    return _central_from_raw(r1, r2, r3, r4)


# --- Gompertz(eta, theta) and Gompertz-Makeham(lam, eta, theta) : mortality --- #
def gompertz_pdf(x, eta, theta):
    x = np.asarray(x, float)
    return np.where(x >= 0, eta * np.exp(theta * x) * np.exp(-(eta / theta) * (np.exp(theta * x) - 1)), 0.0)

def gompertz_cdf(x, eta, theta):
    x = np.asarray(x, float)
    return np.where(x >= 0, 1 - np.exp(-(eta / theta) * (np.exp(theta * np.maximum(x, 0)) - 1)), 0.0)

def gompertz_hazard(x, eta, theta):
    x = np.asarray(x, float)
    return eta * np.exp(theta * x)                       # exponentially increasing failure rate

def gompertz_rng(eta, theta, n, rng):
    """Closed-form inverse: x = (1/theta) * ln(1 - (theta/eta) ln U)."""
    return (1.0 / theta) * np.log(1 - (theta / eta) * np.log(rng.random(n)))

def gompertz_moments(eta, theta):
    return moments_from_pdf(lambda t: gompertz_pdf(t, eta, theta), 0.0, 50.0 / theta)

def makeham_pdf(x, lam, eta, theta):
    x = np.asarray(x, float)
    S = np.exp(-(lam * np.maximum(x, 0) + (eta / theta) * (np.exp(theta * np.maximum(x, 0)) - 1)))
    return np.where(x >= 0, (lam + eta * np.exp(theta * x)) * S, 0.0)

def makeham_cdf(x, lam, eta, theta):
    x = np.asarray(x, float)
    S = np.exp(-(lam * np.maximum(x, 0) + (eta / theta) * (np.exp(theta * np.maximum(x, 0)) - 1)))
    return np.where(x >= 0, 1 - S, 0.0)

def makeham_hazard(x, lam, eta, theta):
    return lam + eta * np.exp(theta * np.asarray(x, float))   # constant (accidents) + Gompertz (ageing)

def makeham_rng(lam, eta, theta, n, rng):
    """No closed-form quantile (transcendental) -> numerical inverse-transform."""
    return inverse_transform_sample(lambda t: makeham_cdf(t, lam, eta, theta), 0.0, 60.0 / theta, n, rng)

def makeham_moments(lam, eta, theta):
    return moments_from_pdf(lambda t: makeham_pdf(t, lam, eta, theta), 0.0, 60.0 / theta)


# ===========================================================================  #
#  PART C -- ACTUARIAL LOSS / INCOME (the GB2 family)                           #
# ===========================================================================  #

# --- GB2(a, b, p, q) : Generalized Beta of the 2nd kind ---------------------- #
def gb2_pdf(x, a, b, p, q):
    x = np.asarray(x, float); xp = np.where(x > 0, x, 1.0)
    logp = np.log(abs(a)) + (a * p - 1) * np.log(xp) - a * p * np.log(b) - betaln(p, q) \
           - (p + q) * np.log1p((xp / b) ** a)
    return np.where(x > 0, np.exp(logp), 0.0)

def gb2_cdf(x, a, b, p, q):
    x = np.asarray(x, float); z = (np.where(x > 0, x, 0.0) / b) ** a
    u = z / (1 + z)
    return np.where(x > 0, betainc(p, q, u), 0.0)

def gb2_rng(a, b, p, q, n, rng):
    """If W ~ Beta(p,q) then X = b*(W/(1-W))^{1/a} ~ GB2(a,b,p,q)."""
    w = _beta_rng(p, q, n, rng)
    return b * (w / (1 - w)) ** (1.0 / a)

def gb2_moments(a, b, p, q):
    def raw(k):
        if not (-a * p < k < a * q): return np.nan
        return b ** k * np.exp(betaln(p + k / a, q - k / a) - betaln(p, q))
    r1, r2, r3, r4 = raw(1), raw(2), raw(3), raw(4)
    if np.isnan(r2): return dict(mean=r1, var=np.nan, skew=np.nan, exkurt=np.nan)
    if np.isnan(r3): return dict(mean=r1, var=r2 - r1 ** 2, skew=np.nan, exkurt=np.nan)
    if np.isnan(r4):
        c = _central_from_raw(r1, r2, r3, r3); c["exkurt"] = np.nan; return c
    return _central_from_raw(r1, r2, r3, r4)


# --- Burr XII / Singh-Maddala(c, k, b) = GB2(c, b, 1, k) --------------------- #
def burr12_pdf(x, c, k, b):
    x = np.asarray(x, float); z = np.where(x > 0, x, 1.0) / b
    return np.where(x > 0, (c * k / b) * z ** (c - 1) * (1 + z ** c) ** (-k - 1), 0.0)

def burr12_cdf(x, c, k, b):
    x = np.asarray(x, float); z = np.where(x > 0, x, 0.0) / b
    return np.where(x > 0, 1 - (1 + z ** c) ** (-k), 0.0)

def burr12_rng(c, k, b, n, rng):
    """Inverse transform: b*((1-U)^{-1/k} - 1)^{1/c}."""
    return b * ((1 - rng.random(n)) ** (-1.0 / k) - 1) ** (1.0 / c)

def burr12_moments(c, k, b):
    return gb2_moments(c, b, 1.0, k)


# --- Dagum / Burr III(a, b, p) = GB2(a, b, p, 1) ----------------------------- #
def dagum_pdf(x, a, b, p):
    x = np.asarray(x, float); z = np.where(x > 0, x, 1.0) / b
    return np.where(x > 0, (a * p / b) * z ** (a * p - 1) * (1 + z ** a) ** (-p - 1), 0.0)

def dagum_cdf(x, a, b, p):
    x = np.asarray(x, float); z = np.where(x > 0, x, 1.0) / b
    return np.where(x > 0, (1 + z ** (-a)) ** (-p), 0.0)

def dagum_rng(a, b, p, n, rng):
    """Inverse transform: b*(U^{-1/p} - 1)^{-1/a}."""
    return b * (rng.random(n) ** (-1.0 / p) - 1) ** (-1.0 / a)

def dagum_moments(a, b, p):
    return gb2_moments(a, b, p, 1.0)


# --------------------------------------------------------------------------- #
#  Private RNG primitives (Box-Muller normal, Marsaglia-Tsang gamma, Beta)      #
# --------------------------------------------------------------------------- #

def _std_normal(n, rng):
    m = (n + 1) // 2
    u1 = rng.random(m); u2 = rng.random(m)
    r = np.sqrt(-2 * np.log(u1)); z = np.empty(2 * m)
    z[:m] = r * np.cos(2 * np.pi * u2); z[m:] = r * np.sin(2 * np.pi * u2)
    return z[:n]

def _gamma_shape(k, n, rng):
    n = int(n)
    if k < 1:
        g = _gamma_shape(k + 1.0, n, rng); u = rng.random(n)
        return g * u ** (1.0 / k)
    d = k - 1.0 / 3.0; c = 1.0 / np.sqrt(9.0 * d)
    res = np.empty(n); todo = np.ones(n, bool)
    while todo.any():
        m = int(todo.sum()); x = _std_normal(m, rng)
        v = (1.0 + c * x) ** 3; u = rng.random(m)
        ok = (v > 0) & (np.log(u) < 0.5 * x ** 2 + d - d * v + d * np.log(np.where(v > 0, v, 1.0)))
        cur = np.where(todo)[0]; acc = cur[ok]
        res[acc] = d * v[ok]; todo[acc] = False
    return res

def _beta_rng(p, q, n, rng):
    x = _gamma_shape(p, n, rng); y = _gamma_shape(q, n, rng)
    return x / (x + y)