"""Monte Carlo summaries, confidence intervals and coverage."""
from __future__ import annotations
import numpy as np
from scipy import stats


def summarize(errors: np.ndarray):
    """errors: (B,) or (B,p) array -> dict with bias, sd, rmse, MC standard errors."""
    e = np.asarray(errors, float)
    B = e.shape[0]
    mean = e.mean(axis=0); sd = e.std(axis=0, ddof=1)
    rmse = np.sqrt(np.mean(e ** 2, axis=0))
    return {"n": B, "bias": mean, "sd": sd, "rmse": rmse, "bias_se": sd / np.sqrt(B),
            "rmse_se": np.sqrt(np.var(e ** 2, axis=0, ddof=1) / B) / (2 * np.maximum(rmse, 1e-300)),
            "p95_abs": np.quantile(np.abs(e), 0.95, axis=0)}


def wald_ci(theta, cov, N, level=0.95):
    z = stats.norm.ppf(0.5 + level / 2)
    se = np.sqrt(np.diag(cov) / N)
    return theta - z * se, theta + z * se


def coverage(theta_hat, cov_hat, N, theta0, levels=(0.90, 0.95, 0.99)):
    """theta_hat (B,p), cov_hat (B,p,p) [asymptotic covariance of sqrt(N)(theta-theta0)] -> coverage per level and param."""
    out = {}
    for lv in levels:
        z = stats.norm.ppf(0.5 + lv / 2)
        se = np.sqrt(np.einsum('bpp->bp', cov_hat) / N)
        hit = np.abs(theta_hat - theta0) <= z * se
        out[lv] = {"coverage": hit.mean(axis=0), "mean_length": (2 * z * se).mean(axis=0)}
    return out