from abc import ABCMeta, abstractmethod
from collections.abc import Iterator
from typing import Any, ClassVar, Literal
from typing_extensions import Self

from numpy import ndarray
from numpy.random import RandomState

from .._typing import ArrayLike, Float, Int, MatrixLike
from ..base import BaseEstimator, ClassifierMixin, RegressorMixin
from ._base import BaseEnsemble

__all__ = [
    "AdaBoostClassifier",
    "AdaBoostRegressor",
]

class BaseWeightBoosting(BaseEnsemble, metaclass=ABCMeta):
    _parameter_constraints: ClassVar[dict] = ...

    @abstractmethod
    def __init__(
        self,
        estimator=None,
        *,
        n_estimators: int = 50,
        estimator_params=...,
        learning_rate: float = 1.0,
        random_state=None,
        base_estimator: str = "deprecated",
    ) -> None: ...
    def fit(
        self,
        X: MatrixLike | ArrayLike,
        y: ArrayLike,
        sample_weight: None | ArrayLike = None,
    ) -> Self: ...
    def staged_score(
        self,
        X: MatrixLike | ArrayLike,
        y: ArrayLike,
        sample_weight: None | ArrayLike = None,
    ): ...
    @property
    def feature_importances_(self) -> ndarray: ...

class AdaBoostClassifier(ClassifierMixin, BaseWeightBoosting):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    feature_importances_: ndarray = ...
    estimator_errors_: ndarray = ...
    estimator_weights_: ndarray = ...
    n_classes_: int = ...
    classes_: ndarray = ...
    estimators_: list[ClassifierMixin] = ...
    base_estimator_: BaseEstimator = ...
    estimator_: BaseEstimator = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        estimator: Any = None,
        *,
        n_estimators: Int = 50,
        learning_rate: Float = 1.0,
        algorithm: Literal["SAMME", "SAMME.R"] = "SAMME.R",
        random_state: RandomState | None | Int = None,
        base_estimator: Any = "deprecated",
    ) -> None: ...
    def predict(self, X: MatrixLike | ArrayLike) -> ndarray: ...
    def staged_predict(self, X: MatrixLike) -> Iterator[ndarray]: ...
    def decision_function(self, X: MatrixLike | ArrayLike) -> ndarray: ...
    def staged_decision_function(self, X: MatrixLike | ArrayLike) -> Iterator[ndarray]: ...
    def predict_proba(self, X: MatrixLike | ArrayLike) -> ndarray: ...
    def staged_predict_proba(self, X: MatrixLike | ArrayLike): ...
    def predict_log_proba(self, X: MatrixLike | ArrayLike) -> ndarray: ...

class AdaBoostRegressor(RegressorMixin, BaseWeightBoosting):
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    feature_importances_: ndarray = ...
    estimator_errors_: ndarray = ...
    estimator_weights_: ndarray = ...
    estimators_: list[RegressorMixin] = ...
    base_estimator_: BaseEstimator = ...
    estimator_: BaseEstimator = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        estimator: Any = None,
        *,
        n_estimators: Int = 50,
        learning_rate: Float = 1.0,
        loss: Literal["linear", "square", "exponential"] = "linear",
        random_state: RandomState | None | Int = None,
        base_estimator: Any = "deprecated",
    ) -> None: ...
    def predict(self, X: MatrixLike | ArrayLike) -> ndarray: ...
    def staged_predict(self, X: MatrixLike | ArrayLike): ...