"""Certified L2 radius of a linear model is exact."""

import numpy as np
import pytest
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression

from clfkit.robustness import certified_radius, linear_weights


def _model(n_classes, seed=0):
    X, y = make_classification(n_samples=400, n_features=8, n_informative=5, n_classes=n_classes,
                               n_clusters_per_class=1, random_state=seed)
    return LogisticRegression(max_iter=2000).fit(X, y), X


def _predict(W, b, Z):
    return (Z @ W.T + b).argmax(1)


@pytest.mark.parametrize("n_classes", [2, 4])
def test_radius_is_exact(n_classes):
    clf, X = _model(n_classes)
    W, b = linear_weights(clf)
    Z = X[:60]
    pred, rad, D = certified_radius(Z, W, b, return_direction=True)
    assert np.array_equal(pred, clf.predict(Z))  # our scores reproduce sklearn's argmax
    rng = np.random.default_rng(0)
    for i in range(len(Z)):
        r = rad[i]
        # 1) Random perturbations strictly inside the radius never flip the label.
        U = rng.normal(size=(200, Z.shape[1]))
        U /= np.linalg.norm(U, axis=1, keepdims=True)
        scales = r * rng.uniform(0, 0.999, size=(200, 1))
        assert np.all(_predict(W, b, Z[i] + scales * U) == pred[i])
        # 2) The worst-case direction, just inside: still the same label.
        assert _predict(W, b, (Z[i] + 0.999 * r * D[i])[None])[0] == pred[i]
        # 3) Just beyond the radius along that direction: label flips (tight).
        assert _predict(W, b, (Z[i] + 1.001 * r * D[i])[None])[0] != pred[i]


def test_binary_matches_closed_form():
    clf, X = _model(2, seed=1)
    _, rad = certified_radius(X, *linear_weights(clf))
    w, b0 = clf.coef_[0], clf.intercept_[0]
    assert np.allclose(rad, np.abs(X @ w + b0) / np.linalg.norm(w))