from __future__ import annotations

import numpy as np
from scipy import stats


def pit_histogram(u: np.ndarray, bins: int = 20) -> dict:
    h, edges = np.histogram(np.clip(u, 0, 1), bins=bins, range=(0, 1))
    n = len(u)
    expected = n / bins
    chi2 = float(((h - expected) ** 2 / expected).sum())
    ks = stats.kstest(u, "uniform")
    return {"hist": (h / n).tolist(), "edges": edges.tolist(), "chi2": chi2, "chi2_p": float(1 - stats.chi2.cdf(chi2, bins - 1)),
            "ks_stat": float(ks.statistic), "ks_p": float(ks.pvalue), "mean": float(u.mean()), "var": float(u.var()),
            "n": int(n)}


def reliability(y: np.ndarray, s: np.ndarray, levels=(0.5, 0.8, 0.9, 0.95, 0.99)) -> list[dict]:
    out = []
    for lv in levels:
        a = 1 - lv
        lo, hi = np.quantile(s, [a / 2, 1 - a / 2], axis=0)
        cov = float(((y >= lo) & (y <= hi)).mean())
        out.append({"level": lv, "coverage": cov, "width": float((hi - lo).mean()), "gap": cov - lv})
    return out