# (c) 2026 William J. Dwyer. MIT License. Cite Zenodo DOI 10.5281/zenodo.20709963.
#
# verify_distribution_atlas.py
# Reproduces every closed form in "The Log-Domain Distributional Atlas" (companion to M1):
# the Var(ln X) bridges, the constant log-skewness signatures, the shape-from-CV solvers,
# and an end-to-end test of model_atlas() (generate from each family -> recover it).
#
# Run:  python3 verify_distribution_atlas.py   (exits 0 iff all checks pass)
import sys
import numpy as np
from scipy import stats, optimize
from scipy.special import polygamma, gamma as Gamma, gammaln

PASS=[]
def check(name, cond, detail=""):
    PASS.append(bool(cond)); print(f"[{'PASS' if cond else 'FAIL'}] {name}"+(f"  — {detail}" if detail else ""))

GUMBEL_SKEW = -12*np.sqrt(6)*1.2020569031595942/np.pi**3   # -12 sqrt6 zeta(3)/pi^3 = min-Gumbel skew
# ============================ shape-from-CV solvers ============================
def k_gamma(cv):        return 1.0/cv**2
def alpha_invgamma(cv): return 2.0 + 1.0/cv**2
def alpha_pareto(cv):   return 1.0 + np.sqrt(1.0 + 1.0/cv**2)        # closed: a(a-2)=1/cv^2
def weibull_cv2(k):  # overflow-safe via log-gamma
    return np.exp(gammaln(1+2/k) - 2*gammaln(1+1/k)) - 1
def k_weibull(cv):
    return optimize.brentq(lambda k: weibull_cv2(k) - cv**2, 1e-2, 1e4)
def beta_loglogistic(cv):
    # mean=b1/sin(b1) (beta>1); E[X^2]=b2/sin(b2) (beta>2); cv^2=var/mean^2
    def cv2(beta):
        b1=np.pi/beta; b2=2*np.pi/beta
        mean=b1/np.sin(b1); m2=b2/np.sin(b2); return m2/mean**2 - 1
    # beta>2 required; cv2 decreases from +inf (beta->2) to 0 (beta->inf)
    return optimize.brentq(lambda b: cv2(b)-cv**2, 2.0001, 1e4)

# ============================ predicted log-cumulants ==========================
def pred_lognormal(cv):  s2=np.log1p(cv**2); return dict(lv=s2, ls=0.0, lk=0.0)
def pred_gamma(cv):
    k=k_gamma(cv); v=polygamma(1,k)
    return dict(lv=v, ls=polygamma(2,k)/v**1.5, lk=polygamma(3,k)/v**2)
def pred_invgamma(cv):
    a=alpha_invgamma(cv); v=polygamma(1,a)
    return dict(lv=v, ls=-polygamma(2,a)/v**1.5, lk=polygamma(3,a)/v**2)   # lnX=-lnGamma: odd cumulant flips
def pred_weibull(cv):
    k=k_weibull(cv); return dict(lv=np.pi**2/(6*k**2), ls=GUMBEL_SKEW, lk=2.4)
def pred_loglogistic(cv):
    b=beta_loglogistic(cv); return dict(lv=np.pi**2/(3*b**2), ls=0.0, lk=1.2)
def pred_pareto(cv):
    a=alpha_pareto(cv); return dict(lv=1.0/a**2, ls=2.0, lk=6.0)

FAMILIES={"lognormal":pred_lognormal,"gamma":pred_gamma,"inverse-gamma":pred_invgamma,
          "weibull":pred_weibull,"log-logistic":pred_loglogistic,"pareto":pred_pareto}

# ============================ the atlas classifier ============================
def log_cumulants(x):
    y=np.log(np.asarray(x,float)); m=y.mean(); v=y.var()
    return v, ((y-m)**3).mean()/v**1.5, ((y-m)**4).mean()/v**2 - 3.0

def model_atlas(cv, lv, ls, lk=None, wv=1.0, ws=1.0, wk=0.25):
    """Rank positive families by how well their CV-implied log-cumulants match the observed
    (log-variance lv, log-skewness ls, optional log-excess-kurtosis lk)."""
    out=[]
    for name,fn in FAMILIES.items():
        try: p=fn(cv)
        except Exception: continue
        d = wv*((lv-p["lv"])/max(abs(lv),1e-6))**2 + ws*(ls-p["ls"])**2
        if lk is not None: d += wk*(lk-p["lk"])**2
        out.append((name, np.sqrt(d), p))
    out.sort(key=lambda t:t[1])
    return out

# ============================================================================
print("="*72); print("LOG-DOMAIN DISTRIBUTIONAL ATLAS — verification"); print("="*72)

# ---- Closed-form Var(lnX) and log-skewness vs simulation ----
print("\n-- closed forms vs simulation (n = 2,000,000) --")
rng=np.random.default_rng(0); N=2_000_000
def emp(dist):
    x=dist.rvs(N,random_state=rng); return log_cumulants(x)
cases=[
 ("lognormal", stats.lognorm(s=0.6), np.log1p((np.exp(0.36)-1)), 0.0),  # cv^2=e^{0.36}-1
 ("gamma k=4", stats.gamma(a=4), polygamma(1,4), polygamma(2,4)/polygamma(1,4)**1.5),
 ("weibull k=1.5", stats.weibull_min(c=1.5), np.pi**2/(6*1.5**2), GUMBEL_SKEW),
 ("log-logistic b=3", stats.fisk(c=3), np.pi**2/(3*3**2), 0.0),
 ("pareto a=4", stats.pareto(b=4), 1/16.0, 2.0),
 ("inverse-gamma a=5", stats.invgamma(a=5), polygamma(1,5), -polygamma(2,5)/polygamma(1,5)**1.5),
]
okv=oks=True
for name,dist,lv_cf,ls_cf in cases:
    v,s,k=emp(dist)
    dv=abs(v-lv_cf); ds=abs(s-ls_cf)
    okv &= dv<0.01*max(lv_cf,0.05)+0.01; oks &= ds<0.03
    print(f"  {name:<18} Var(lnX) {v:.4f} (cf {lv_cf:.4f})  logskew {s:+.4f} (cf {ls_cf:+.4f})")
check("Var(ln X) closed forms match simulation", okv)
check("log-skewness closed forms match simulation", oks)

# ---- The constant signatures (independent of shape) ----
print("\n-- constant log-skewness signatures (shape-independent) --")
wb=[log_cumulants(stats.weibull_min(c=k).rvs(1_000_000,random_state=rng))[1] for k in (0.7,2.0,5.0)]
pa=[log_cumulants(stats.pareto(b=a).rvs(1_000_000,random_state=rng))[1]     for a in (3.0,6.0,12.0)]
check("Weibull log-skewness is constant ≈ -1.1395 for all shapes", max(abs(np.array(wb)-GUMBEL_SKEW))<0.02, f"{np.round(wb,3)}")
check("Pareto log-skewness is constant ≈ +2 for all tail indices",  max(abs(np.array(pa)-2.0))<0.03, f"{np.round(pa,3)}")
check("Lognormal & log-logistic have log-skewness 0 (symmetric logs)",
      abs(log_cumulants(stats.lognorm(s=0.5).rvs(1_000_000,random_state=rng))[1])<0.02 and
      abs(log_cumulants(stats.fisk(c=4).rvs(1_000_000,random_state=rng))[1])<0.02)

# ---- shape-from-CV solvers invert correctly ----
print("\n-- shape-from-CV solvers --")
oks2=True
for cv in [0.3,0.6,0.9]:
    k=k_weibull(cv); cv_back=np.sqrt(weibull_cv2(k)); oks2 &= abs(cv_back-cv)<1e-6
    b=beta_loglogistic(cv);
    a=alpha_pareto(cv); cvp=np.sqrt(1/(a*(a-2))); oks2 &= abs(cvp-cv)<1e-9
print(f"  Weibull/Pareto CV inversion exact at CV=0.3,0.6,0.9")
check("shape-from-CV solvers invert to the target CV", oks2)

# ---- End-to-end: model_atlas recovers the generating family ----
print("\n-- model_atlas: generate from each family, recover it (n = 400,000) --")
def cv_of(x): x=np.asarray(x,float); return x.std(ddof=1)/x.mean()
gens={
 "lognormal":   stats.lognorm(s=0.55),
 "gamma":       stats.gamma(a=3.0),
 "weibull":     stats.weibull_min(c=1.6),
 "log-logistic":stats.fisk(c=3.5),
 "pareto":      stats.pareto(b=4.0),
 "inverse-gamma":stats.invgamma(a=6.0),
}
hits=0
for truth,dist in gens.items():
    x=dist.rvs(400_000,random_state=rng); x=x[x>0]
    cv=cv_of(x); lv,ls,lk=log_cumulants(x)
    rank=model_atlas(cv,lv,ls,lk)
    top=rank[0][0]; ok=(top==truth)
    hits+=ok
    print(f"  truth {truth:<13} CV={cv:.3f} -> nearest {top:<13} {'OK' if ok else '(2nd: '+rank[1][0]+')'}")
check("model_atlas recovers the generating family for all six", hits==len(gens), f"{hits}/{len(gens)}")

print("\n"+"="*72)
ok=all(PASS); print(f"RESULT: {sum(PASS)}/{len(PASS)} checks passed — {'ALL PASS' if ok else 'FAILURES'}")
print("="*72); sys.exit(0 if ok else 1)
