"""Ridge estimator implementation for NeuroWeave."""

from __future__ import annotations

from typing import Any

import numpy as np

from neuroweave.core.base_estimator import BaseEstimator


class RidgeEstimator(BaseEstimator):
    """Minimal ridge regression estimator.

    Parameters
    ----------
    alpha : float, optional
        L2 regularization strength.
    """

    def __init__(self, alpha: float = 1.0) -> None:
        super().__init__(alpha=alpha)
        self.alpha = alpha
        self.coef_: np.ndarray | None = None

    def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Estimate ridge regression coefficients.

        Parameters
        ----------
        X : np.ndarray
            Design matrix with shape ``(samples, features)``.
        y : np.ndarray
            Target values with shape ``(samples,)``.

        Returns
        -------
        np.ndarray
            Estimated coefficient vector.
        """

        X_array = np.asarray(X, dtype=float)
        y_array = np.asarray(y, dtype=float).reshape(-1)

        xtx = X_array.T @ X_array
        regularizer = self.alpha * np.eye(xtx.shape[0], dtype=X_array.dtype)
        xty = X_array.T @ y_array

        self.coef_ = np.linalg.solve(xtx + regularizer, xty)
        return self.coef_

    def get_params(self) -> dict[str, Any]:
        """Return estimator parameters.

        Returns
        -------
        dict[str, Any]
            Estimator parameter mapping.
        """

        return {"alpha": self.alpha}

    def set_params(self, **params: Any) -> RidgeEstimator:
        """Set estimator parameters.

        Parameters
        ----------
        **params : Any
            Estimator parameter values.

        Returns
        -------
        RidgeEstimator
            The estimator instance.
        """

        if "alpha" in params:
            self.alpha = float(params["alpha"])
        return self
