from typing import Callable, ClassVar, Literal
from typing_extensions import Self

from joblib.memory import MemorizedFunc
from numpy import ndarray

from .._typing import ArrayLike, Float, Int, MatrixLike
from ..base import BaseEstimator
from ._base import SelectorMixin

# Authors: V. Michel, B. Thirion, G. Varoquaux, A. Gramfort, E. Duchesnay.
#          L. Buitinck, A. Joly
# License: BSD 3 clause

######################################################################
# Scoring functions

# The following function is a rewriting of scipy.stats.f_oneway
# Contrary to the scipy.stats.f_oneway implementation it does not
# copy the data while keeping the inputs unchanged.
def f_oneway(*args) -> tuple[ndarray, ndarray] | tuple[float, float]: ...
def f_classif(X: MatrixLike | ArrayLike, y: ArrayLike) -> tuple[ndarray, ndarray]: ...
def chi2(X: MatrixLike | ArrayLike, y: ArrayLike) -> tuple[ndarray, ndarray]: ...
def r_regression(X: MatrixLike | ArrayLike, y: ArrayLike, *, center: bool = True, force_finite: bool = True) -> ndarray: ...
def f_regression(
    X: MatrixLike | ArrayLike, y: ArrayLike, *, center: bool = True, force_finite: bool = True
) -> tuple[ndarray, ndarray]: ...

######################################################################
# Base classes

class _BaseFilter(SelectorMixin, BaseEstimator):
    _parameter_constraints: ClassVar[dict] = ...

    def __init__(self, score_func: Callable | MemorizedFunc) -> None: ...
    def fit(self, X: MatrixLike, y: ArrayLike) -> Self: ...

######################################################################
# Specific filters
######################################################################
class SelectPercentile(_BaseFilter):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    pvalues_: ArrayLike = ...
    scores_: ArrayLike = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(self, score_func: Callable | MemorizedFunc = ..., *, percentile: Int = 10) -> None: ...

class SelectKBest(_BaseFilter):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    pvalues_: ArrayLike = ...
    scores_: ArrayLike = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(self, score_func: Callable = ..., *, k: str | int = 10) -> None: ...

class SelectFpr(_BaseFilter):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    pvalues_: ArrayLike = ...
    scores_: ArrayLike = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(self, score_func: Callable = ..., *, alpha: Float = 5e-2) -> None: ...

class SelectFdr(_BaseFilter):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    pvalues_: ArrayLike = ...
    scores_: ArrayLike = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(self, score_func: Callable = ..., *, alpha: Float = 5e-2) -> None: ...

class SelectFwe(_BaseFilter):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    pvalues_: ArrayLike = ...
    scores_: ArrayLike = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(self, score_func: Callable = ..., *, alpha: Float = 5e-2) -> None: ...

######################################################################
# Generic filter
######################################################################

class GenericUnivariateSelect(_BaseFilter):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    pvalues_: ArrayLike = ...
    scores_: ArrayLike = ...

    _selection_modes: ClassVar[dict] = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        score_func: Callable = ...,
        *,
        mode: Literal["percentile", "k_best", "fpr", "fdr", "fwe"] = "percentile",
        param: float | str = 1e-5,
    ) -> None: ...
