"""Hilbert-Schmidt-regularized operator estimators."""

from __future__ import annotations

from dataclasses import dataclass
import numpy as np


def _one_hot(y: np.ndarray, n_classes: int) -> np.ndarray:
    y = np.asarray(y, dtype=np.int64)
    out = np.zeros((len(y), n_classes), dtype=np.float64)
    out[np.arange(len(y)), y] = 1.0
    return out


@dataclass
class HilbertSchmidtOperatorClassifier:
    """Closed-form task operator T: H_X -> R^C with HS/Frobenius regularization.

    Given feature vectors Phi(x_i), the classifier solves
        min_W ||X W - Y||_F^2 + lambda ||W||_F^2.
    This is the finite-dimensional realization of the Hilbert-space operator
    estimation objective in the paper.
    """

    ridge_lambda: float = 1e-3
    standardize: bool = True
    add_bias: bool = True

    def fit(self, X: np.ndarray, y: np.ndarray) -> "HilbertSchmidtOperatorClassifier":
        X = np.asarray(X, dtype=np.float64)
        y = np.asarray(y, dtype=np.int64)
        self.classes_ = np.unique(y)
        if not np.array_equal(self.classes_, np.arange(len(self.classes_))):
            raise ValueError("Labels must be contiguous from 0 to C-1")
        self.n_classes_ = len(self.classes_)

        if self.standardize:
            self.mean_ = X.mean(axis=0, keepdims=True)
            self.std_ = X.std(axis=0, keepdims=True) + 1e-8
            Xs = (X - self.mean_) / self.std_
        else:
            self.mean_ = np.zeros((1, X.shape[1]))
            self.std_ = np.ones((1, X.shape[1]))
            Xs = X
        if self.add_bias:
            Xs = np.concatenate([Xs, np.ones((Xs.shape[0], 1))], axis=1)

        Y = _one_hot(y, self.n_classes_)
        # Solve (X^T X + lambda I) W = X^T Y
        A = Xs.T @ Xs
        reg = self.ridge_lambda * np.eye(A.shape[0], dtype=np.float64)
        if self.add_bias:
            reg[-1, -1] = 0.0  # do not regularize bias
        B = Xs.T @ Y
        self.W_ = np.linalg.solve(A + reg, B)
        self.param_count_ = int(self.W_.size)
        return self

    def decision_function(self, X: np.ndarray) -> np.ndarray:
        X = np.asarray(X, dtype=np.float64)
        Xs = (X - self.mean_) / self.std_
        if self.add_bias:
            Xs = np.concatenate([Xs, np.ones((Xs.shape[0], 1))], axis=1)
        return Xs @ self.W_

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        logits = self.decision_function(X)
        logits = logits - logits.max(axis=1, keepdims=True)
        exp = np.exp(logits)
        return exp / exp.sum(axis=1, keepdims=True)

    def predict(self, X: np.ndarray) -> np.ndarray:
        return self.decision_function(X).argmax(axis=1)


def fit_rbf_kernel_ridge(X_train: np.ndarray, y_train: np.ndarray, gamma: float = 0.03, lam: float = 1e-2, n_landmarks: int = 256, seed: int = 7):
    """Fast RKHS baseline using Nystrom RBF features and operator ridge.

    A full kernel solve is intentionally avoided so the repository remains fast
    on CPU-only environments. The learned operator still lives in an explicit
    finite-dimensional RKHS approximation.
    """
    rng = np.random.default_rng(seed)
    X_train = np.asarray(X_train, dtype=np.float64)
    y_train = np.asarray(y_train, dtype=np.int64)
    n = len(X_train)
    m = min(n_landmarks, n)
    idx = rng.choice(n, size=m, replace=False)
    landmarks = X_train[idx]

    def features(X):
        X = np.asarray(X, dtype=np.float64)
        sq = ((X[:, None, :] - landmarks[None, :, :]) ** 2).sum(axis=-1)
        return np.exp(-gamma * sq)

    Z = features(X_train)
    clf = HilbertSchmidtOperatorClassifier(ridge_lambda=lam, standardize=True).fit(Z, y_train)

    class RBFKernelRidge:
        param_count_ = int(clf.param_count_ + landmarks.size)

        def decision_function(self, X):
            return clf.decision_function(features(X))

        def predict_proba(self, X):
            return clf.predict_proba(features(X))

        def predict(self, X):
            return clf.predict(features(X))

    return RBFKernelRidge()