"""Measured endpoint metrics and grouped uncertainty summaries."""
import numpy as np
import pandas as pd

def intervention_metrics(probabilities, observed_losses, tolerance, seed=0, draws=10, repeats=200):
    p, loss = np.asarray(probabilities, float), np.asarray(observed_losses, float)
    if p.shape != loss.shape or (p < 0).any() or not np.isclose(p.sum(), 1):
        raise ValueError("A normalized action distribution and aligned losses are required")
    if not np.isfinite(loss).all(): raise ValueError("Observed losses must be finite")
    selected = np.random.default_rng(seed).choice(len(p), size=(repeats, draws), p=p)
    expected = float(p@loss)
    return {"endpoint_loss": expected, "regret": expected-float(loss.min()),
            "success": float(p@(loss <= tolerance)), "best_of_10": float(loss[selected].min(1).mean()),
            "entropy": float(-np.sum(p*np.log(np.maximum(p, 1e-300)))),
            "effective_actions": float(np.exp(-np.sum(p*np.log(np.maximum(p, 1e-300)))))}

def grouped_summary(frame, metric="endpoint_loss", group="target_molecule", bootstrap=2000, seed=0):
    rng = np.random.default_rng(seed)
    values = frame.groupby([group, "method"])[metric].mean().unstack()
    if values.isna().any().any():
        raise ValueError("Methods must share the same evaluated molecule groups")
    n = len(values)
    if n == 0: raise ValueError("No evaluation groups")
    draws = values.to_numpy()[rng.integers(n, size=(bootstrap, n))].mean(1)
    rows = []
    for j, name in enumerate(values.columns):
        lo, hi = np.quantile(draws[:, j], [.025, .975])
        rows.append({"method": name, "metric": metric, "mean": float(values.iloc[:, j].mean()),
                     "ci_low": float(lo), "ci_high": float(hi), "molecules": n})
    return pd.DataFrame(rows)

def paired_difference(frame, first, second, metric="endpoint_loss", seed=0, bootstrap=2000):
    pivot = frame.groupby(["target_molecule", "method"])[metric].mean().unstack()
    diff = (pivot[first]-pivot[second]).dropna().to_numpy()
    if not len(diff): raise ValueError("No paired molecule groups")
    rng = np.random.default_rng(seed)
    values = diff[rng.integers(len(diff), size=(bootstrap, len(diff)))].mean(1)
    return {"first": first, "second": second, "metric": metric, "difference": float(diff.mean()),
            "ci_low": float(np.quantile(values, .025)), "ci_high": float(np.quantile(values, .975))}

def rbf_mmd(x, y, bandwidth):
    from scipy.spatial.distance import cdist
    if bandwidth <= 0 or len(x) < 2 or len(y) < 2: raise ValueError("MMD needs positive bandwidth and two cells per group")
    xx, yy, xy = [np.exp(-d/(2*bandwidth**2)) for d in [cdist(x, x, "sqeuclidean"), cdist(y, y, "sqeuclidean"), cdist(x, y, "sqeuclidean")]]
    return float((xx.sum()-np.trace(xx))/(len(x)*(len(x)-1))+(yy.sum()-np.trace(yy))/(len(y)*(len(y)-1))-2*xy.mean())