"""
This type stub file was generated by pyright.
"""

import numpy as np
import scipy.sparse as ss
from enum import Enum
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, Union
from .compat import Client, dask_Array, pd_DataFrame, pd_Series
from .sklearn import LGBMClassifier, LGBMRanker, LGBMRegressor

"""Distributed training with LightGBM and dask.distributed.

This module enables you to perform distributed training with LightGBM on
dask.Array and dask.DataFrame collections.

It is based on dask-lightgbm, which was based on dask-xgboost.
"""
_DaskCollection = ...
_DaskMatrixLike = ...
_DaskVectorLike = ...
_DaskPart = Union[np.ndarray, pd_DataFrame, pd_Series, ss.spmatrix]
_PredictionDtype = Union[Type[np.float32], Type[np.float64], Type[np.int32], Type[np.int64]]
_HostWorkers = ...
class _DatasetNames(Enum):
    """Placeholder names used by lightgbm.dask internals to say 'also evaluate the training data'.

    Avoid duplicating the training data when the validation set refers to elements of training data.
    """
    TRAINSET = ...
    SAMPLE_WEIGHT = ...
    INIT_SCORE = ...
    GROUP = ...


class _DaskLGBMModel:
    @property
    def client_(self) -> Client:
        """:obj:`dask.distributed.Client`: Dask client.

        This property can be passed in the constructor or updated
        with ``model.set_params(client=client)``.
        """
        ...
    


class DaskLGBMClassifier(LGBMClassifier, _DaskLGBMModel):
    """Distributed version of lightgbm.LGBMClassifier."""
    def __init__(self, boosting_type: str = ..., num_leaves: int = ..., max_depth: int = ..., learning_rate: float = ..., n_estimators: int = ..., subsample_for_bin: int = ..., objective: Optional[Union[Callable, str]] = ..., class_weight: Optional[Union[dict, str]] = ..., min_split_gain: float = ..., min_child_weight: float = ..., min_child_samples: int = ..., subsample: float = ..., subsample_freq: int = ..., colsample_bytree: float = ..., reg_alpha: float = ..., reg_lambda: float = ..., random_state: Optional[Union[int, np.random.RandomState]] = ..., n_jobs: int = ..., silent: bool = ..., importance_type: str = ..., client: Optional[Client] = ..., **kwargs: Any) -> None:
        """Docstring is inherited from the lightgbm.LGBMClassifier.__init__."""
        ...
    
    _base_doc = ...

    def __getstate__(self) -> Dict[Any, Any]:
        ...
    
    def fit(self, X: _DaskMatrixLike, y: _DaskCollection, sample_weight: Optional[_DaskVectorLike] = ..., init_score: Optional[_DaskCollection] = ..., eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = ..., eval_names: Optional[List[str]] = ..., eval_sample_weight: Optional[List[_DaskVectorLike]] = ..., eval_class_weight: Optional[List[Union[dict, str]]] = ..., eval_init_score: Optional[List[_DaskCollection]] = ..., eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = ..., early_stopping_rounds: Optional[int] = ..., **kwargs: Any) -> DaskLGBMClassifier:
        """Docstring is inherited from the lightgbm.LGBMClassifier.fit."""
        ...
    
    _base_doc = ...

    def predict(self, X: _DaskMatrixLike, **kwargs: Any) -> dask_Array:
        """Docstring is inherited from the lightgbm.LGBMClassifier.predict."""
        ...
    
    def predict_proba(self, X: _DaskMatrixLike, **kwargs: Any) -> dask_Array:
        """Docstring is inherited from the lightgbm.LGBMClassifier.predict_proba."""
        ...
    
    def to_local(self) -> LGBMClassifier:
        """Create regular version of lightgbm.LGBMClassifier from the distributed version.

        Returns
        -------
        model : lightgbm.LGBMClassifier
            Local underlying model.
        """
        ...
    


class DaskLGBMRegressor(LGBMRegressor, _DaskLGBMModel):
    """Distributed version of lightgbm.LGBMRegressor."""
    def __init__(self, boosting_type: str = ..., num_leaves: int = ..., max_depth: int = ..., learning_rate: float = ..., n_estimators: int = ..., subsample_for_bin: int = ..., objective: Optional[Union[Callable, str]] = ..., class_weight: Optional[Union[dict, str]] = ..., min_split_gain: float = ..., min_child_weight: float = ..., min_child_samples: int = ..., subsample: float = ..., subsample_freq: int = ..., colsample_bytree: float = ..., reg_alpha: float = ..., reg_lambda: float = ..., random_state: Optional[Union[int, np.random.RandomState]] = ..., n_jobs: int = ..., silent: bool = ..., importance_type: str = ..., client: Optional[Client] = ..., **kwargs: Any) -> None:
        """Docstring is inherited from the lightgbm.LGBMRegressor.__init__."""
        ...
    
    _base_doc = ...
    
    def __getstate__(self) -> Dict[Any, Any]:
        ...
    
    def fit(self, X: _DaskMatrixLike, y: _DaskCollection, sample_weight: Optional[_DaskVectorLike] = ..., init_score: Optional[_DaskVectorLike] = ..., eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = ..., eval_names: Optional[List[str]] = ..., eval_sample_weight: Optional[List[_DaskVectorLike]] = ..., eval_init_score: Optional[List[_DaskVectorLike]] = ..., eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = ..., early_stopping_rounds: Optional[int] = ..., **kwargs: Any) -> DaskLGBMRegressor:
        """Docstring is inherited from the lightgbm.LGBMRegressor.fit."""
        ...

    def predict(self, X: _DaskMatrixLike, **kwargs) -> dask_Array:
        """Docstring is inherited from the lightgbm.LGBMRegressor.predict."""
        ...
    
    def to_local(self) -> LGBMRegressor:
        """Create regular version of lightgbm.LGBMRegressor from the distributed version.

        Returns
        -------
        model : lightgbm.LGBMRegressor
            Local underlying model.
        """
        ...
    


class DaskLGBMRanker(LGBMRanker, _DaskLGBMModel):
    """Distributed version of lightgbm.LGBMRanker."""
    def __init__(self, boosting_type: str = ..., num_leaves: int = ..., max_depth: int = ..., learning_rate: float = ..., n_estimators: int = ..., subsample_for_bin: int = ..., objective: Optional[Union[Callable, str]] = ..., class_weight: Optional[Union[dict, str]] = ..., min_split_gain: float = ..., min_child_weight: float = ..., min_child_samples: int = ..., subsample: float = ..., subsample_freq: int = ..., colsample_bytree: float = ..., reg_alpha: float = ..., reg_lambda: float = ..., random_state: Optional[Union[int, np.random.RandomState]] = ..., n_jobs: int = ..., silent: bool = ..., importance_type: str = ..., client: Optional[Client] = ..., **kwargs: Any) -> None:
        """Docstring is inherited from the lightgbm.LGBMRanker.__init__."""
        ...
    
    _base_doc = ...
    
    def __getstate__(self) -> Dict[Any, Any]:
        ...
    
    def fit(self, X: _DaskMatrixLike, y: _DaskCollection, sample_weight: Optional[_DaskVectorLike] = ..., init_score: Optional[_DaskVectorLike] = ..., group: Optional[_DaskVectorLike] = ..., eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = ..., eval_names: Optional[List[str]] = ..., eval_sample_weight: Optional[List[_DaskVectorLike]] = ..., eval_init_score: Optional[List[_DaskVectorLike]] = ..., eval_group: Optional[List[_DaskVectorLike]] = ..., eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = ..., eval_at: Iterable[int] = ..., early_stopping_rounds: Optional[int] = ..., **kwargs: Any) -> DaskLGBMRanker:
        """Docstring is inherited from the lightgbm.LGBMRanker.fit."""
        ...
    
    def predict(self, X: _DaskMatrixLike, **kwargs: Any) -> dask_Array:
        """Docstring is inherited from the lightgbm.LGBMRanker.predict."""
        ...
    
    def to_local(self) -> LGBMRanker:
        """Create regular version of lightgbm.LGBMRanker from the distributed version.

        Returns
        -------
        model : lightgbm.LGBMRanker
            Local underlying model.
        """
        ...
    


