# Stubs for pyspark.ml.pipeline (Python 3)

from typing import Any, Dict, List, Optional, Tuple, Type, Union

from pyspark.ml._typing import PipelineStage
from pyspark.context import SparkContext
from pyspark.ml.base import Estimator, Model, Transformer
from pyspark.ml.param import Param
from pyspark.ml.util import *

class Pipeline(Estimator[PipelineModel], MLReadable[Pipeline], MLWritable):
    stages: List[PipelineStage]
    def __init__(self, *, stages: Optional[List[PipelineStage]] = ...) -> None: ...
    def setStages(self, stages: List[PipelineStage]) -> Pipeline: ...
    def getStages(self) -> List[PipelineStage]: ...
    def setParams(self, *, stages: Optional[List[PipelineStage]] = ...) -> Pipeline: ...
    def copy(self, extra: Optional[Dict[Param, str]] = ...) -> Pipeline: ...
    def write(self) -> JavaMLWriter: ...
    def save(self, path: str) -> None: ...
    @classmethod
    def read(cls) -> PipelineReader: ...

class PipelineWriter(MLWriter):
    instance: Pipeline
    def __init__(self, instance: Pipeline) -> None: ...
    def saveImpl(self, path: str) -> None: ...

class PipelineReader(MLReader):
    cls: Type[Pipeline]
    def __init__(self, cls: Type[Pipeline]) -> None: ...
    def load(self, path: str) -> Pipeline: ...

class PipelineModelWriter(MLWriter):
    instance: PipelineModel
    def __init__(self, instance: PipelineModel) -> None: ...
    def saveImpl(self, path: str) -> None: ...

class PipelineModelReader(MLReader):
    cls: Type[PipelineModel]
    def __init__(self, cls: Type[PipelineModel]) -> None: ...
    def load(self, path: str) -> PipelineModel: ...

class PipelineModel(Model, MLReadable[PipelineModel], MLWritable):
    stages: List[PipelineStage]
    def __init__(self, stages: List[Transformer]) -> None: ...
    def copy(self, extra: Optional[Dict[Param, Any]] = ...) -> PipelineModel: ...
    def write(self) -> JavaMLWriter: ...
    def save(self, path: str) -> None: ...
    @classmethod
    def read(cls) -> PipelineModelReader: ...

class PipelineSharedReadWrite:
    @staticmethod
    def checkStagesForJava(stages: List[PipelineStage]) -> bool: ...
    @staticmethod
    def validateStages(stages: List[PipelineStage]) -> None: ...
    @staticmethod
    def saveImpl(
        instance: Union[Pipeline, PipelineModel],
        stages: List[PipelineStage],
        sc: SparkContext,
        path: str,
    ) -> None: ...
    @staticmethod
    def load(
        metadata: Dict[str, Any], sc: SparkContext, path: str
    ) -> Tuple[str, List[PipelineStage]]: ...
    @staticmethod
    def getStagePath(
        stageUid: str, stageIdx: int, numStages: int, stagesDir: str
    ) -> str: ...
