"""Temporal receptive field model implementation."""

from __future__ import annotations

from typing import Any

import numpy as np

from neuroweave.core.base_estimator import BaseEstimator
from neuroweave.core.base_model import BaseEncodingModel
from neuroweave.core.result import EncodingResult
from neuroweave.design.lag import make_lagged_matrix
from neuroweave.estimators.linear.ridge import RidgeEstimator


class TRFModel(BaseEncodingModel):
    """Minimal temporal receptive field model.

    Parameters
    ----------
    lags : tuple[int, int], optional
        Inclusive lag window.
    alpha : float, optional
        Ridge regularization strength.
    estimator : BaseEstimator | None, optional
        Estimator used to fit model parameters.
    """

    def __init__(
        self,
        lags: tuple[int, int] = (-10, 10),
        alpha: float = 1.0,
        estimator: BaseEstimator | None = None,
    ) -> None:
        super().__init__(lags=lags, alpha=alpha)
        self.lags = lags
        self.alpha = alpha
        self.estimator = estimator if estimator is not None else RidgeEstimator(alpha=alpha)
        self.kernel_: np.ndarray | None = None
        self.result_: EncodingResult | None = None

    def fit(
        self,
        X: np.ndarray,
        y: np.ndarray,
        estimator: BaseEstimator | None = None,
    ) -> EncodingResult:
        """Fit the TRF model.

        Parameters
        ----------
        X : np.ndarray
            Feature matrix with shape ``(time, features)``.
        y : np.ndarray
            Target vector with shape ``(time,)``.
        estimator : BaseEstimator | None, optional
            Optional estimator override.

        Returns
        -------
        EncodingResult
            Result object containing fitted kernel and predictions.
        """

        X_array = np.asarray(X, dtype=float)
        y_array = np.asarray(y, dtype=float).reshape(-1)
        active_estimator = estimator if estimator is not None else self.estimator

        X_lagged = make_lagged_matrix(X_array, self.lags)
        beta = active_estimator.fit(X_lagged, y_array)
        predictions = X_lagged @ beta

        self.estimator = active_estimator
        self.kernel_ = beta
        self.result_ = EncodingResult(
            kernel=beta,
            predictions=predictions,
            score=None,
            metadata={"y_true": y_array, "lags": self.lags},
        )
        return self.result_

    def predict(self, X: np.ndarray) -> np.ndarray:
        """Predict responses from the fitted TRF model.

        Parameters
        ----------
        X : np.ndarray
            Feature matrix with shape ``(time, features)``.

        Returns
        -------
        np.ndarray
            Predicted responses.
        """

        if self.kernel_ is None:
            raise ValueError("The model must be fitted before calling predict().")

        X_lagged = make_lagged_matrix(np.asarray(X, dtype=float), self.lags)
        return X_lagged @ self.kernel_

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        """Compute the coefficient of determination.

        Parameters
        ----------
        X : np.ndarray
            Feature matrix with shape ``(time, features)``.
        y : np.ndarray
            Target vector with shape ``(time,)``.

        Returns
        -------
        float
            Coefficient of determination.
        """

        predictions = self.predict(X)
        result = EncodingResult(predictions=predictions, metadata={"y_true": y})
        return result.r2()

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

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

        return {
            "lags": self.lags,
            "alpha": self.alpha,
            "estimator": self.estimator,
        }

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

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

        Returns
        -------
        TRFModel
            The model instance.
        """

        if "lags" in params:
            self.lags = tuple(params["lags"])
        if "alpha" in params:
            self.alpha = float(params["alpha"])
            if isinstance(self.estimator, RidgeEstimator):
                self.estimator.set_params(alpha=self.alpha)
        if "estimator" in params:
            self.estimator = params["estimator"]
        return self
