"""Public model entry points for NeuroWeave."""

from __future__ import annotations

from typing import Any

import numpy as np

from neuroweave.core.result import EncodingResult
from neuroweave.models.trf.trf import TRFModel


class BaseEncodingModel:
    """Base class for public API model wrappers."""

    def fit(self, X: np.ndarray, y: np.ndarray) -> EncodingResult:
        """Fit the encoding model.

        Parameters
        ----------
        X : np.ndarray
            Feature matrix.
        y : np.ndarray
            Target vector.

        Returns
        -------
        EncodingResult
            Fitted model result.
        """

        raise NotImplementedError()

    def predict(self, X: np.ndarray) -> np.ndarray:
        """Predict targets from features.

        Parameters
        ----------
        X : np.ndarray
            Feature matrix.

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

        raise NotImplementedError()

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        """Score the encoding model.

        Parameters
        ----------
        X : np.ndarray
            Feature matrix.
        y : np.ndarray
            Target vector.

        Returns
        -------
        float
            Model score.
        """

        raise NotImplementedError()


class TRF(BaseEncodingModel):
    """User-facing temporal receptive field model."""

    def __init__(
        self,
        lags: tuple[int, int] = (-10, 10),
        alpha: float = 1.0,
    ) -> None:
        self.model = TRFModel(lags=lags, alpha=alpha)

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

        Parameters
        ----------
        X : np.ndarray
            Feature matrix.
        y : np.ndarray
            Target vector.

        Returns
        -------
        EncodingResult
            Fitted model result.
        """

        return self.model.fit(X, y)

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

        Parameters
        ----------
        X : np.ndarray
            Feature matrix.

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

        return self.model.predict(X)

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        """Score the fitted TRF model.

        Parameters
        ----------
        X : np.ndarray
            Feature matrix.
        y : np.ndarray
            Target vector.

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

        return self.model.score(X, y)


class ComplexTRF(BaseEncodingModel):
    """Placeholder wrapper for a complex-valued TRF model."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.args = args
        self.kwargs = kwargs


class PhaseTRF(BaseEncodingModel):
    """Placeholder wrapper for a phase-aware TRF model."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.args = args
        self.kwargs = kwargs


class DynamicTRF(BaseEncodingModel):
    """Placeholder wrapper for a dynamic TRF model."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.args = args
        self.kwargs = kwargs


class MediatedTRF(BaseEncodingModel):
    """Placeholder wrapper for a mediated TRF model."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.args = args
        self.kwargs = kwargs


class SwitchingTRF(BaseEncodingModel):
    """Placeholder wrapper for a switching TRF model."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.args = args
        self.kwargs = kwargs


__all__ = [
    "BaseEncodingModel",
    "TRF",
    "ComplexTRF",
    "PhaseTRF",
    "DynamicTRF",
    "MediatedTRF",
    "SwitchingTRF",
]
