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

from typing import Any, Dict, List, Optional, Tuple, Type, Union
from pyspark.context import SparkContext
from pyspark.ml.base import Estimator, Model, Transformer
from pyspark.ml.param import Param
from pyspark.ml.util import *

PipelineStage = Union[Estimator, Transformer]

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: ...
