"""ERT metrics for conditional coverage (arXiv:2512.11779v1).

Table 1 of the paper, transcribed:

  name      proper score l(p,y)                        l-ERT formula
  L1-ERT    sgn(p-(1-a)) (1-a-y)                       E_X |1-a - p(X)|
  L2-ERT    Brier, (y-p)^2                             E_X (1-a - p(X))^2
  KL-ERT    log-loss, -log p_y                         E_X D_KL(p(X) || 1-a)

with l-ERT(h) := R_l(1-a) - R_l(h), i.e. the risk of the constant 1-a
predictor minus the risk of a fitted classifier h of the coverage indicator
Z = 1{y in C(x)}.  Conditional coverage holds iff no classifier beats the
constant, so ERT >= 0 measures the violation.

Design choice for the synthetic experiments
-------------------------------------------
The data-generating process here is chosen so that the TRUE conditional
coverage p(x) is available in closed form.  With y|x ~ N(f(x), s(x)^2) and a
symmetric interval [yhat(x) - q, yhat(x) + q],

    p(x) = Phi((q + yhat(x) - f(x))/s(x)) - Phi((-q + yhat(x) - f(x))/s(x))

so the true L1, L2 and KL deviations can be computed to Monte-Carlo accuracy
and every estimate can be scored against ground truth rather than against
another estimate.  The oracle interval [f(x) -+ z_{1-a/2} s(x)] gives
p(x) = 1-a exactly, which is the negative control every metric must return
(near) zero on.
"""
import numpy as np
from scipy.stats import norm
from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor
from sklearn.model_selection import KFold
from sklearn.tree import DecisionTreeClassifier

EPS = 1e-6


# ------------------------------------------------------------------- losses
def loss_L1(p, y, alpha):
    return np.sign(p - (1 - alpha)) * ((1 - alpha) - y)


def loss_L2(p, y, alpha=None):
    return (y - p) ** 2


def loss_KL(p, y, alpha=None):
    p = np.clip(p, EPS, 1 - EPS)
    return -(y * np.log(p) + (1 - y) * np.log(1 - p))


LOSSES = {"L1": loss_L1, "L2": loss_L2, "KL": loss_KL}


def ert(h, z, alpha, kind):
    """l-ERT(h) = R_l(1-a) - R_l(h)."""
    const = np.full_like(np.asarray(h, float), 1 - alpha)
    f = LOSSES[kind]
    if kind == "L1":
        return float(np.mean(f(const, z, alpha) - f(h, z, alpha)))
    return float(np.mean(f(const, z) - f(h, z)))


def ert_signed(h, z, alpha, kind):
    """Asymmetric decomposition of Section 3.3: restrict the fitted
    probability to over-coverage (h > 1-a) and under-coverage (h < 1-a),
    replacing it by the constant elsewhere so each part is itself an ERT."""
    h = np.asarray(h, float)
    h_over = np.where(h > 1 - alpha, h, 1 - alpha)     # conservatism
    h_under = np.where(h < 1 - alpha, h, 1 - alpha)    # aggressiveness
    return ert(h_over, z, alpha, kind), ert(h_under, z, alpha, kind)


# --------------------------------------------------------- true (oracle) ERT
def true_ert(p, alpha, kind):
    p = np.clip(np.asarray(p, float), EPS, 1 - EPS)
    a = 1 - alpha
    if kind == "L1":
        return float(np.mean(np.abs(a - p)))
    if kind == "L2":
        return float(np.mean((a - p) ** 2))
    return float(np.mean(p * np.log(p / a) + (1 - p) * np.log((1 - p) / (1 - a))))


def true_ert_signed(p, alpha, kind):
    p = np.asarray(p, float)
    a = 1 - alpha
    over = np.where(p > a, p, a)
    under = np.where(p < a, p, a)
    return true_ert(over, alpha, kind), true_ert(under, alpha, kind)


# ------------------------------------------------------------------- the DGP
def f_mean(X):
    return X[:, 0] + X[:, 1] ** 2


def s_sd(X, hetero=True):
    if not hetero:
        return np.full(len(X), 0.8)
    return 0.4 + np.abs(X[:, 0]) + 0.5 * np.abs(X[:, 1])


def sample(n, d, rng, hetero=True, skew=0.0):
    X = rng.uniform(-1, 1, size=(n, d))
    e = rng.standard_normal(n)
    if skew:                       # asymmetric law for the claim-4 experiment
        e = e + skew * (X[:, 2] > 0) * np.abs(rng.standard_normal(n))
    y = f_mean(X) + s_sd(X, hetero) * e
    return X, y


def true_coverage(X, yhat, q, hetero=True):
    """Closed-form P(|y - yhat(x)| <= q | x) for the Gaussian DGP."""
    s = s_sd(X, hetero)
    c = yhat - f_mean(X)
    return norm.cdf((q + c) / s) - norm.cdf((-q + c) / s)


# ------------------------------------------------------- conformal predictors
def split_conformal(Xtr, ytr, Xcal, ycal, alpha, seed=0):
    """Standard split conformal with a homoskedastic base model: the interval
    width is constant, so conditional coverage is violated wherever s(x)
    departs from its average."""
    m = HistGradientBoostingRegressor(random_state=seed, max_iter=200).fit(Xtr, ytr)
    res = np.abs(ycal - m.predict(Xcal))
    n = len(res)
    k = int(np.ceil((n + 1) * (1 - alpha)))
    q = float(np.sort(res)[min(k, n) - 1])
    return m, q


def coverage_indicator(y, yhat, q):
    return (np.abs(y - yhat) <= q).astype(float)


# ------------------------------------------------------------ estimating ERT
def make_classifier(seed=0, model="hgb"):
    """The classifier used inside Algorithm 1.

    The configuration matters more than it looks: an *untuned* boosted-tree
    default (max_iter=200, no regularisation) is flexible enough to overfit the
    Bernoulli coverage indicator out of fold, and then loses to the constant
    1-alpha predictor on the strictly proper losses -- L2- and KL-ERT come out
    NEGATIVE.  A shallow, strongly regularised model with early stopping
    recovers ~88% of the true L1 deviation instead.  This is the paper's own
    thesis (the classifier is what gives the metric its power) showing up as a
    reproduction hazard.
    """
    if model == "tree":
        return DecisionTreeClassifier(random_state=seed)
    return HistGradientBoostingClassifier(
        random_state=seed, max_iter=400, learning_rate=0.03, max_leaf_nodes=7,
        l2_regularization=5.0, min_samples_leaf=60, early_stopping=True,
        validation_fraction=0.15)


def cv_predict(X, z, seed=0, folds=5, model="hgb"):
    """Algorithm 1: out-of-fold predictions so the classifier is never
    evaluated on the points it was fitted on."""
    out = np.empty(len(z), float)
    kf = KFold(folds, shuffle=True, random_state=seed)
    for tr, te in kf.split(X):
        g = make_classifier(seed, model)
        if len(np.unique(z[tr])) < 2:
            out[te] = z[tr].mean()
            continue
        g.fit(X[tr], z[tr])
        out[te] = g.predict_proba(X[te])[:, 1]
    return np.clip(out, EPS, 1 - EPS)


def in_sample_predict(X, z, seed=0, model="tree"):
    if len(np.unique(z)) < 2:
        return np.full(len(z), z.mean())
    g = make_classifier(seed, model)
    g.fit(X, z)
    return np.clip(g.predict_proba(X)[:, 1], EPS, 1 - EPS)


# --------------------------------------------------------------- CovGap
def covgap(X, z, alpha, n_groups=10, seed=0):
    """Group-conditional coverage gap: partition X and average the absolute
    deviation of each group's empirical coverage from 1-alpha."""
    rng = np.random.default_rng(seed)
    from sklearn.cluster import KMeans
    km = KMeans(n_clusters=n_groups, n_init=4, random_state=seed).fit(X)
    g = km.labels_
    gaps = []
    for k in range(n_groups):
        m = g == k
        if m.sum() == 0:
            continue
        gaps.append(abs(z[m].mean() - (1 - alpha)))
    return float(np.mean(gaps))