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

from numpy import ndarray

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

class BaseLabelPropagation(ClassifierMixin, BaseEstimator, metaclass=ABCMeta):
    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        kernel: Callable | Literal["rbf", "knn"] = "rbf",
        *,
        gamma: Float = 20,
        n_neighbors: Int = 7,
        alpha: Float = 1,
        max_iter: Int = 30,
        tol: Float = 1e-3,
        n_jobs=None,
    ) -> None: ...
    def predict(self, X: MatrixLike) -> ndarray: ...
    def predict_proba(self, X: MatrixLike) -> ndarray: ...
    def fit(self, X: MatrixLike, y: ArrayLike) -> Self | LabelSpreading: ...

class LabelPropagation(BaseLabelPropagation):
    n_iter_: int = ...
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    transduction_: ndarray = ...
    label_distributions_: ndarray = ...
    classes_: ndarray = ...
    X_: ndarray = ...

    _variant: ClassVar[str] = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        kernel: Callable | Literal["knn", "rbf"] = "rbf",
        *,
        gamma: Float = 20,
        n_neighbors: Int = 7,
        max_iter: Int = 1000,
        tol: float = 1e-3,
        n_jobs: None | Int = None,
    ) -> None: ...
    def fit(self, X: MatrixLike, y: ArrayLike) -> Self: ...

class LabelSpreading(BaseLabelPropagation):
    n_iter_: int = ...
    feature_names_in_: ndarray = ...
    n_features_in_: int = ...
    transduction_: ndarray = ...
    label_distributions_: ndarray = ...
    classes_: ndarray = ...
    X_: ndarray = ...

    _variant: ClassVar[str] = ...

    _parameter_constraints: ClassVar[dict] = ...

    def __init__(
        self,
        kernel: Callable | Literal["rbf", "knn"] = "rbf",
        *,
        gamma: Float = 20,
        n_neighbors: Int = 7,
        alpha: Float = 0.2,
        max_iter: Int = 30,
        tol: Float = 1e-3,
        n_jobs: None | Int = None,
    ) -> None: ...