"""Robustness: exact certificates for linear models, empirical checks for any model.

Exact L2 certificate for a linear classifier
--------------------------------------------
A linear multiclass model scores f_k(z) = w_k . z + b_k and predicts
c = argmax_k f_k(z). The set where c wins is the polyhedron
{z : (w_c - w_k) . z + (b_c - b_k) >= 0 for all k != c}. The L2 distance from
z to leaving an intersection of half-spaces is the minimum distance to any of
their boundaries, so the largest radius r such that NO perturbation with
||delta||_2 < r can change the prediction is exactly

    r(z) = min_{k != c} (f_c(z) - f_k(z)) / ||w_c - w_k||_2 .

Binary (one weight vector w): r(z) = |w . z + b| / ||w||_2.
This is tight: moving by r + eps along the minimising direction flips (or, at
exactly r, ties) the prediction.

Scope - be precise about what this certifies:
* It is a statement in the model's INPUT space, i.e. after preprocessing
  (standardised numerics, one-hot categoricals, TF-IDF vectors). A radius of
  0.5 means "0.5 standard deviations of Euclidean change across standardised
  features", not raw units.
* It certifies the argmax label of THIS linear model, which is invariant to
  temperature scaling, but not conformal sets or other models.
* It does not cover perturbations that are not small in L2 (e.g. swapping a
  word changes a TF-IDF vector by a lot).

For non-linear models we only offer an *empirical* stability check: the
fraction of predictions unchanged under random Gaussian noise. Random noise is
far weaker than adversarial noise, so a high stability rate is NOT a
certificate - it is a smoke test for brittleness.
"""

from __future__ import annotations

import numpy as np
import pandas as pd


def linear_weights(clf) -> tuple[np.ndarray, np.ndarray]:
    """Return (W, b) with one row per class, for sklearn linear classifiers.

    sklearn stores a single row for binary problems (score of class 1); we
    expand to two rows [0, w] so the multiclass formula applies uniformly.
    """
    W = np.asarray(clf.coef_, dtype=float)
    b = np.asarray(clf.intercept_, dtype=float)
    if W.shape[0] == 1:
        W = np.vstack([np.zeros_like(W[0]), W[0]])
        b = np.array([0.0, b[0]])
    return W, b


def certified_radius(Z, W: np.ndarray, b: np.ndarray, return_direction: bool = False):
    """Exact L2 certified radius of the argmax prediction for each row of Z.

    Z may be dense or scipy-sparse. Returns ``(pred, radius)`` and, if asked,
    the unit direction that reaches the decision boundary fastest.
    """
    F = np.asarray(Z @ W.T) + b  # (n, K)
    pred = F.argmax(1)
    n, K = F.shape
    # Pairwise ||w_i - w_j|| computed once (K x K) - avoids an (n x d) array.
    G = W @ W.T
    sq = np.diag(G)
    pair_norm = np.sqrt(np.clip(sq[:, None] + sq[None, :] - 2 * G, 0.0, None))
    radius = np.full(n, np.inf)
    runner = np.zeros(n, dtype=int)
    for k in range(K):
        norm = pair_norm[pred, k]
        gap = F[np.arange(n), pred] - F[:, k]
        with np.errstate(divide="ignore", invalid="ignore"):
            r = np.where((k != pred) & (norm > 0), gap / norm, np.inf)
        better = r < radius
        radius = np.where(better, r, radius)
        runner = np.where(better, k, runner)
    if not return_direction:
        return pred, radius
    D = W[runner] - W[pred]
    D = D / np.linalg.norm(D, axis=1, keepdims=True)
    return pred, radius, D


def certified_accuracy_curve(pred, radius, y, radii) -> np.ndarray:
    """Fraction of points that are correct AND certified at radius >= r, for each r."""
    correct = pred == y
    return np.array([np.mean(correct & (radius >= r)) for r in radii])


def perturb_dataframe(
    X: pd.DataFrame, numeric_cols, text_cols, sigma: float, rng: np.random.Generator,
    col_std: dict | None = None,
) -> pd.DataFrame:
    """Random perturbation of raw inputs.

    Numeric columns: add N(0, (sigma * column std)^2).
    Text columns: drop each token independently with probability ``sigma``
    (clipped to [0, 0.9]). Categoricals are left alone.
    """
    Xp = X.copy()
    for c in numeric_cols:
        sd = (col_std or {}).get(c, float(X[c].std()) or 1.0)
        Xp[c] = X[c] + rng.normal(0.0, sigma * sd, size=len(X))
    p_drop = float(np.clip(sigma, 0.0, 0.9))
    for c in text_cols:
        Xp[c] = [
            " ".join(t for t in str(doc).split() if rng.uniform() >= p_drop) for doc in X[c]
        ]
    return Xp


def empirical_stability(
    model, X: pd.DataFrame, numeric_cols, text_cols, sigmas=(0.05, 0.1, 0.2, 0.5),
    n_repeats: int = 5, seed: int = 0, col_std: dict | None = None,
) -> list[dict]:
    """Fraction of predictions unchanged under random perturbations, per noise level.

    ``stable_all`` requires the prediction to survive all ``n_repeats`` draws.
    """
    rng = np.random.default_rng(seed)
    base = model.predict(X)
    rows = []
    for s in sigmas:
        same = np.ones(len(X), dtype=bool)
        agree = []
        for _ in range(n_repeats):
            p = model.predict(perturb_dataframe(X, numeric_cols, text_cols, s, rng, col_std))
            same &= p == base
            agree.append(np.mean(p == base))
        rows.append({"sigma": float(s), "agreement": float(np.mean(agree)), "stable_all": float(same.mean())})
    return rows