"""mcmc_checks: the fit-checking pass from the AthlyticZ session "MCMC in Practice" (Chris Fonnesbeck),
as one reusable function you can call on anything pm.sample() returns.

    import mcmc_checks as mc
    report = mc.check_fit(idata)          # prints the verdict, returns a dict
    mc.ppc_overlay(idata, "y")            # one posterior predictive check, as a plot

What it checks, in the order the session teaches it:
  1. divergences            any at all means the sampler hit geometry it could not follow (think non-centered)
  2. R-hat                  > 1.01 on any parameter: the chains disagree, do not trust the posterior yet
  3. effective sample size  bulk or tail ESS under 400: the estimate is noisier than the draw count suggests
  4. BFMI                   under 0.3 on any chain: the energy transitions are poor, often a funnel
  5. tree depth             hits at the cap mean NUTS was cut short; raise the cap or reparameterize

Requires: arviz >= 0.17, matplotlib. Works on InferenceData from PyMC 5 and 6 (nutpie included).
License: MIT. Written for athlyticz.com/skills from the method the session teaches; the session's own
notebook and datasets come with the recording.
"""
from __future__ import annotations

import numpy as np

try:
    import arviz as az
except ImportError as e:  # pragma: no cover
    raise SystemExit("mcmc_checks needs arviz: pip install arviz") from e


def check_fit(idata, var_names=None, rhat_max=1.01, ess_min=400, bfmi_min=0.3, verbose=True):
    """Run the five sampler checks on an InferenceData and return {check: (ok, detail)}.

    var_names narrows the R-hat and ESS tables to the parameters you care about (defaults to all).
    The thresholds are the session's: R-hat 1.01, ESS 400, BFMI 0.3.
    """
    out = {}
    ss = idata.get("sample_stats")

    # 1. divergences
    if ss is not None and "diverging" in ss:
        n_div = int(ss["diverging"].sum())
        out["divergences"] = (n_div == 0, f"{n_div} divergent transition(s)")
    else:
        out["divergences"] = (True, "no sample_stats.diverging recorded")

    # 2 and 3. R-hat and ESS, over the posterior
    summ = az.summary(idata, var_names=var_names, kind="diagnostics")
    worst_rhat = float(summ["r_hat"].max())
    out["r_hat"] = (worst_rhat <= rhat_max, f"max R-hat {worst_rhat:.3f} ({summ['r_hat'].idxmax()})")
    ess_bulk, ess_tail = float(summ["ess_bulk"].min()), float(summ["ess_tail"].min())
    out["ess"] = (min(ess_bulk, ess_tail) >= ess_min,
                  f"min bulk ESS {ess_bulk:.0f} ({summ['ess_bulk'].idxmin()}), min tail ESS {ess_tail:.0f} ({summ['ess_tail'].idxmin()})")

    # 4. BFMI, per chain
    try:
        bfmi = az.bfmi(idata)
        out["bfmi"] = (bool(np.all(bfmi >= bfmi_min)), "per-chain BFMI " + ", ".join(f"{b:.2f}" for b in bfmi))
    except Exception as e:  # no energy recorded (e.g. a non-HMC sampler)
        out["bfmi"] = (True, f"BFMI not available ({type(e).__name__})")

    # 5. tree depth saturation
    if ss is not None and "tree_depth" in ss:
        depth = ss["tree_depth"].values
        cap = int(depth.max())
        frac = float((depth == cap).mean())
        out["tree_depth"] = (frac < 0.05 or cap < 10, f"max depth {cap}, {frac:.1%} of transitions at it")
    else:
        out["tree_depth"] = (True, "no tree_depth recorded")

    if verbose:
        width = max(len(k) for k in out)
        for k, (ok, detail) in out.items():
            print(f"{'ok  ' if ok else 'FAIL'} {k:<{width}}  {detail}")
        bad = [k for k, (ok, _) in out.items() if not ok]
        print("verdict:", "the computation is sound; now check the predictions" if not bad
              else "do not interpret this posterior yet: " + ", ".join(bad))
    return out


def ppc_overlay(idata, var, model=None, draws=100, ax=None):
    """One posterior predictive check: simulated datasets over the observed one.

    If idata has no posterior_predictive group, pass the PyMC model and it is sampled here.
    Returns the matplotlib Axes. The question to ask of the picture is the session's: do the
    simulated datasets look like the real one, or is the model missing something the
    sampler diagnostics could never show?
    """
    import matplotlib.pyplot as plt

    if "posterior_predictive" not in idata:
        if model is None:
            raise ValueError("no posterior_predictive group: pass model= so it can be sampled")
        import pymc as pm
        with model:
            pm.sample_posterior_predictive(idata, extend_inferencedata=True, random_seed=1)
    ax = az.plot_ppc(idata, var_names=[var], num_pp_samples=draws, ax=ax)
    ax.set_title(f"posterior predictive check: {var}")
    return ax


if __name__ == "__main__":
    # a two-minute self-test: a well-specified model, then a funnel that should fail
    import pymc as pm

    rng = np.random.default_rng(1)
    y = rng.normal(2.0, 1.0, size=200)
    with pm.Model() as good:
        mu = pm.Normal("mu", 0, 5)
        sd = pm.HalfNormal("sd", 2)
        pm.Normal("y", mu, sd, observed=y)
        idata = pm.sample(1000, tune=1000, chains=4, progressbar=False, random_seed=1)
    print("well-specified model:")
    check_fit(idata)

    with pm.Model() as funnel:
        tau = pm.Normal("tau", 0, 3)
        x = pm.Normal("x", 0, pm.math.exp(tau), shape=9)
        idata_f = pm.sample(1000, tune=1000, chains=4, progressbar=False, random_seed=1)
    print("\nNeal's funnel (centered):")
    check_fit(idata_f)