from typing import List, Optional, Union

import pandas as pd
from sklearn import model_selection

from snowflake.ml.modeling._internal.distributed_hpo_trainer import (
    DistributedHPOTrainer,
)
from snowflake.ml.modeling._internal.estimator_utils import is_single_node
from snowflake.ml.modeling._internal.model_trainer import ModelTrainer
from snowflake.ml.modeling._internal.pandas_trainer import PandasModelTrainer
from snowflake.ml.modeling._internal.snowpark_trainer import SnowparkModelTrainer
from snowflake.snowpark import DataFrame, Session

_PROJECT = "ModelDevelopment"


class ModelTrainerBuilder:
    """
    A builder class to create instances of ModelTrainer for different models and training conditions.

    This class provides methods to build instances of ModelTrainer tailored to specific machine learning
    models and training configurations like dataset's location etc. It abstracts the creation process,
    allowing the user to obtain a configured ModelTrainer for a particular model architecture or configuration.
    """

    _ENABLE_DISTRIBUTED = True

    @classmethod
    def _check_if_distributed_hpo_enabled(cls, session: Session) -> bool:
        return not is_single_node(session) and ModelTrainerBuilder._ENABLE_DISTRIBUTED is True

    @classmethod
    def build(
        cls,
        estimator: object,
        dataset: Union[DataFrame, pd.DataFrame],
        input_cols: Optional[List[str]] = None,
        label_cols: Optional[List[str]] = None,
        sample_weight_col: Optional[str] = None,
        autogenerated: bool = False,
        subproject: str = "",
    ) -> ModelTrainer:
        """
        Builder method that creates an approproiate ModelTrainer instance based on the given params.
        """
        assert input_cols is not None  # Make MyPy happpy
        if isinstance(dataset, pd.DataFrame):
            return PandasModelTrainer(
                estimator=estimator,
                dataset=dataset,
                input_cols=input_cols,
                label_cols=label_cols,
                sample_weight_col=sample_weight_col,
            )
        elif isinstance(dataset, DataFrame):
            trainer_klass = SnowparkModelTrainer
            assert dataset._session is not None  # Make MyPy happpy
            if isinstance(estimator, model_selection.GridSearchCV) or isinstance(
                estimator, model_selection.RandomizedSearchCV
            ):
                if ModelTrainerBuilder._check_if_distributed_hpo_enabled(session=dataset._session):
                    trainer_klass = DistributedHPOTrainer
            return trainer_klass(
                estimator=estimator,
                dataset=dataset,
                session=dataset._session,
                input_cols=input_cols,
                label_cols=label_cols,
                sample_weight_col=sample_weight_col,
                autogenerated=autogenerated,
                subproject=subproject,
            )
        else:
            raise TypeError(
                f"Unexpected dataset type: {type(dataset)}."
                "Supported dataset types: snowpark.DataFrame, pandas.DataFrame."
            )
