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

from numpy import ndarray
from scipy.sparse import spmatrix

from .._typing import ArrayLike, Int, MatrixLike
from ..base import BaseEstimator, TransformerMixin

# Authors: Nicolas Tresegnie <nicolas.tresegnie@gmail.com>
#          Sergey Feldman <sergeyfeldman@gmail.com>
# License: BSD 3 clause

class _BaseImputer(TransformerMixin, BaseEstimator):
    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        *,
        missing_values=...,
        add_indicator: bool = False,
        keep_empty_features: bool = False,
    ) -> None: ...

class SimpleImputer(_BaseImputer):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    indicator_: MissingIndicator = ...
    statistics_: ndarray = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        *,
        missing_values: float | None | str = ...,
        strategy: str = "mean",
        fill_value: float | None | str = None,
        verbose: str | Int = "deprecated",
        copy: bool = True,
        add_indicator: bool = False,
        keep_empty_features: bool = False,
    ) -> None: ...
    def fit(self, X: MatrixLike, y: Any = None) -> Self: ...
    def transform(self, X: MatrixLike) -> ndarray | spmatrix: ...
    def inverse_transform(self, X: MatrixLike) -> ndarray: ...
    def get_feature_names_out(self, input_features: None | ArrayLike = None) -> ndarray: ...

class MissingIndicator(TransformerMixin, BaseEstimator):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    features_: ndarray = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        *,
        missing_values: float | None | str = ...,
        features: Literal["missing-only", "all"] = "missing-only",
        sparse: Literal["auto"] | bool = "auto",
        error_on_new: bool = True,
    ) -> None: ...
    def fit(self, X: MatrixLike | ArrayLike, y: Any = None) -> Self: ...
    def transform(self, X: MatrixLike | ArrayLike) -> ndarray | spmatrix: ...
    def fit_transform(self, X: MatrixLike | ArrayLike, y: Any = None) -> ndarray | spmatrix: ...
    def get_feature_names_out(self, input_features: None | ArrayLike = None) -> ndarray: ...
