"""Small reachable-set risk head for Gaussian linearized dynamics."""
import numpy as np
from scipy.special import ndtr


def propagate(mu, cov, jacobian, process_cov):
    """One Gaussian covariance recursion step."""
    mu = np.asarray(mu, dtype=float)
    cov = np.asarray(cov, dtype=float)
    J = np.asarray(jacobian, dtype=float)
    Q = np.asarray(process_cov, dtype=float)
    return mu, J @ cov @ J.T + Q


def halfspace_probability(mu, cov, c, b, eps=1e-12):
    """P(c^T X >= b), X ~ N(mu,cov)."""
    c = np.asarray(c, dtype=float)
    variance = max(float(c @ cov @ c), eps)
    margin = (float(b - c @ mu)) / np.sqrt(variance)
    return float(ndtr(-margin)), margin


def reachable_risk(mu0, cov0, dynamics, controls, unsafe_halfspaces, process_cov,
                   return_trace=False):
    """Propagate a locally linear Gaussian model and union risk over H steps.

    dynamics(mu, u) returns (next_mean, jacobian). Halfspaces are (c, b).
    """
    mu, cov = np.asarray(mu0, float), np.asarray(cov0, float)
    survival = 1.0
    trace = []
    for u in controls:
        next_mu, J = dynamics(mu, u)
        mu, cov = propagate(next_mu, cov, J, process_cov)
        qs, margins = [], []
        for c, b in unsafe_halfspaces:
            q, m = halfspace_probability(mu, cov, c, b)
            qs.append(q); margins.append(m)
            survival *= (1.0 - q)
        trace.append({"mu": mu.copy(), "cov": cov.copy(), "q": qs, "margin": margins})
    risk = 1.0 - survival
    return (risk, trace) if return_trace else risk


def fixed_false_warning_recall(scores, labels, false_positive_rate=0.05):
    """Recall at an empirical threshold allowing at most target false warnings."""
    scores, labels = np.asarray(scores), np.asarray(labels).astype(bool)
    negatives = scores[~labels]
    threshold = np.inf if len(negatives) == 0 else np.quantile(negatives, 1-false_positive_rate)
    return float(np.mean(scores[labels] >= threshold)), float(threshold)