"""From-scratch support-vector and kernel machines.

Two load-bearing algorithms, built by hand and validated against scikit-learn:

* PegasosSVM -- the soft-margin linear support-vector classifier trained by
  stochastic sub-gradient descent on the hinge-loss objective
      (lambda/2)||w||^2 + (1/n) sum_i max(0, 1 - y_i (w.x_i)),
  Shalev-Shwartz et al., "Pegasos" (ICML 2007).  The max-margin idea and the
  hinge loss are the whole story of the SVM; this is the optimiser.

* KernelRidge -- kernel ridge regression in closed form,
      alpha = (K + n*lambda*I)^{-1} y,   f(x) = sum_i alpha_i k(x, x_i),
  which solves the SAME linear system as scikit-learn's KernelRidge (so the
  fits match to numerical tolerance) and, with an RBF kernel and noise
  variance = n*lambda, equals the POSTERIOR MEAN of a Gaussian process --
  the exact bridge to the Bayesian-nonparametric arc demonstrated in the
  notebook.

The kernel trick: replace every inner product x.x' by k(x,x') = <phi(x),phi(x')>
and a linear method becomes nonlinear without ever forming phi.  rbf_kernel and
poly_kernel below are the two used in the notebook.
"""
import numpy as np


def rbf_kernel(A, B, gamma):
    """Gaussian / squared-exponential kernel exp(-gamma ||a-b||^2)."""
    a2 = np.sum(A**2, axis=1)[:, None]
    b2 = np.sum(B**2, axis=1)[None, :]
    sq = np.maximum(a2 + b2 - 2 * A @ B.T, 0.0)
    return np.exp(-gamma * sq)


def poly_kernel(A, B, degree=3, coef0=1.0, gamma=1.0):
    """Polynomial kernel (gamma <a,b> + coef0)^degree."""
    return (gamma * (A @ B.T) + coef0) ** degree


class PegasosSVM:
    """Linear soft-margin SVM via the Pegasos stochastic sub-gradient method.

    Labels are taken in {0,1} and mapped internally to {-1,+1}.  `lam` is the
    L2 regularisation strength (larger = wider margin, more regularised).
    decision_function returns the signed distance w.x + b; predict thresholds
    it at 0.
    """

    def __init__(self, lam=1e-4, n_epochs=20, seed=0):
        self.lam = lam
        self.n_epochs = n_epochs
        self.seed = seed

    def fit(self, X, y):
        X = np.asarray(X, float); n, d = X.shape
        yy = np.where(np.asarray(y) > 0, 1.0, -1.0)
        Xb = np.hstack([X, np.ones((n, 1))])           # absorb bias as an extra feature
        w = np.zeros(d + 1)
        rng = np.random.default_rng(self.seed)
        t = 0
        for _ in range(self.n_epochs):
            for i in rng.permutation(n):
                t += 1
                eta = 1.0 / (self.lam * t)             # decaying step size
                if yy[i] * (Xb[i] @ w) < 1:            # inside the margin -> hinge active
                    w = (1 - eta * self.lam) * w + eta * yy[i] * Xb[i]
                else:
                    w = (1 - eta * self.lam) * w
        self.coef_ = w[:d]; self.intercept_ = w[d]
        return self

    def decision_function(self, X):
        return np.asarray(X, float) @ self.coef_ + self.intercept_

    def predict(self, X):
        return (self.decision_function(X) > 0).astype(int)


class KernelRidge:
    """Kernel ridge regression in closed form (= GP posterior mean).

    kernel: 'rbf' or 'poly'.  alpha is the ridge penalty (scikit-learn's alpha).
    Fits dual_ = (K + alpha*I)^{-1} y and predicts k(x, X) @ dual_ -- no centering
    of y, matching scikit-learn's KernelRidge, so standardize the target outside
    the class if you want a data-mean prior rather than a zero-mean one.
    """

    def __init__(self, kernel="rbf", gamma=1.0, alpha=1.0, degree=3, coef0=1.0):
        self.kernel = kernel; self.gamma = gamma; self.alpha = alpha
        self.degree = degree; self.coef0 = coef0

    def _K(self, A, B):
        if self.kernel == "rbf":
            return rbf_kernel(A, B, self.gamma)
        return poly_kernel(A, B, self.degree, self.coef0, self.gamma)

    def fit(self, X, y):
        # matches scikit-learn KernelRidge exactly (no centering): dual = (K+alpha*I)^{-1} y.
        # With an RBF kernel and alpha = noise variance this is the posterior mean of a
        # zero-mean Gaussian process -- standardize y outside for accuracy / a data-mean prior.
        self.X_ = np.asarray(X, float); y = np.asarray(y, float)
        K = self._K(self.X_, self.X_)
        n = K.shape[0]
        self.dual_ = np.linalg.solve(K + self.alpha * np.eye(n), y)
        return self

    def predict(self, X):
        return self._K(np.asarray(X, float), self.X_) @ self.dual_