import importlib
import inspect
import io
import os
import posixpath
import sys
from typing import Any, Dict, List, Optional, Tuple, Union

import cloudpickle as cp
import numpy as np
from scipy.stats import rankdata
from sklearn import model_selection

from snowflake.ml._internal import telemetry
from snowflake.ml._internal.utils import identifier, snowpark_dataframe_utils
from snowflake.ml._internal.utils.temp_file_utils import (
    cleanup_temp_files,
    get_temp_file_path,
)
from snowflake.ml.modeling._internal.model_specifications import (
    ModelSpecificationsBuilder,
)
from snowflake.ml.modeling._internal.snowpark_trainer import SnowparkModelTrainer
from snowflake.snowpark import DataFrame, Session, functions as F
from snowflake.snowpark._internal.utils import (
    TempObjectType,
    random_name_for_temp_object,
)
from snowflake.snowpark.functions import col, sproc, udtf
from snowflake.snowpark.types import IntegerType, StringType, StructField, StructType

cp.register_pickle_by_value(inspect.getmodule(get_temp_file_path))
cp.register_pickle_by_value(inspect.getmodule(identifier.get_inferred_name))

_PROJECT = "ModelDevelopment"
DEFAULT_UDTF_NJOBS = 3


class DistributedHPOTrainer(SnowparkModelTrainer):
    """
    A class for performing distributed hyperparameter optimization (HPO) using Snowpark.

    This class inherits from SnowparkModelTrainer and extends its functionality
    to support distributed HPO for machine learning models. It enables optimization
    of hyperparameters by distributing the tasks across the warehouse using Snowpark.
    """

    def __init__(
        self,
        estimator: object,
        dataset: DataFrame,
        session: Session,
        input_cols: List[str],
        label_cols: Optional[List[str]],
        sample_weight_col: Optional[str],
        autogenerated: bool = False,
        subproject: str = "",
    ) -> None:
        """
        Initializes the DistributedHPOTrainer with a model, a Snowpark DataFrame, feature, and label column names, etc.

        Args:
            estimator: SKLearn compatible estimator or transformer object.
            dataset: The dataset used for training the model.
            session: Snowflake session object to be used for training.
            input_cols: The name(s) of one or more columns in a DataFrame containing a feature to be used for training.
            label_cols: The name(s) of one or more columns in a DataFrame representing the target variable(s) to learn.
            sample_weight_col: The column name representing the weight of training examples.
            autogenerated: A boolean denoting if the trainer is being used by autogenerated code or not.
            subproject: subproject name to be used in telemetry.
        """
        super().__init__(
            estimator=estimator,
            dataset=dataset,
            session=session,
            input_cols=input_cols,
            label_cols=label_cols,
            sample_weight_col=sample_weight_col,
            autogenerated=autogenerated,
            subproject=subproject,
        )

    # TODO(snandamuri): Copied this code as it is from the snowpark_handler.
    #   Update it to improve the readability.
    def fit_search_snowpark(
        self,
        param_grid: Union[model_selection.ParameterGrid, model_selection.ParameterSampler],
        dataset: DataFrame,
        session: Session,
        estimator: Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV],
        dependencies: List[str],
        udf_imports: List[str],
        input_cols: List[str],
        label_cols: Optional[List[str]],
        sample_weight_col: Optional[str],
    ) -> Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV]:
        from itertools import product

        import cachetools
        from sklearn.base import clone, is_classifier
        from sklearn.calibration import check_cv

        # Create one stage for data and for estimators.
        temp_stage_name = random_name_for_temp_object(TempObjectType.STAGE)
        temp_stage_creation_query = f"CREATE OR REPLACE TEMP STAGE {temp_stage_name};"
        session.sql(temp_stage_creation_query).collect()

        # Stage data.
        dataset = snowpark_dataframe_utils.cast_snowpark_dataframe(dataset)
        remote_file_path = f"{temp_stage_name}/{temp_stage_name}.parquet"
        dataset.write.copy_into_location(  # type:ignore[call-overload]
            remote_file_path, file_format_type="parquet", header=True, overwrite=True
        )
        imports = [f"@{row.name}" for row in session.sql(f"LIST @{temp_stage_name}").collect()]

        # Store GridSearchCV's refit variable. If user set it as False, we don't need to refit it again
        original_refit = estimator.refit

        # Create a temp file and dump the estimator to that file.
        estimator_file_name = get_temp_file_path()
        params_to_evaluate = []
        for param_to_eval in list(param_grid):
            for k, v in param_to_eval.items():
                param_to_eval[k] = [v]
            params_to_evaluate.append([param_to_eval])

        with open(estimator_file_name, mode="w+b") as local_estimator_file_obj:
            # Set GridSearchCV refit as False and fit it again after retrieving the best param
            estimator.refit = False
            cp.dump(dict(estimator=estimator, param_grid=params_to_evaluate), local_estimator_file_obj)
        stage_estimator_file_name = posixpath.join(temp_stage_name, os.path.basename(estimator_file_name))
        sproc_statement_params = telemetry.get_function_usage_statement_params(
            project=_PROJECT,
            subproject=self._subproject,
            function_name=telemetry.get_statement_params_full_func_name(
                inspect.currentframe(), self.__class__.__name__
            ),
            api_calls=[sproc],
            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
        )
        udtf_statement_params = telemetry.get_function_usage_statement_params(
            project=_PROJECT,
            subproject=self._subproject,
            function_name=telemetry.get_statement_params_full_func_name(
                inspect.currentframe(), self.__class__.__name__
            ),
            api_calls=[udtf],
            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
        )

        # Put locally serialized estimator on stage.
        put_result = session.file.put(
            estimator_file_name,
            temp_stage_name,
            auto_compress=False,
            overwrite=True,
        )
        estimator_location = put_result[0].target
        imports.append(f"@{temp_stage_name}/{estimator_location}")

        search_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
        random_udtf_name = random_name_for_temp_object(TempObjectType.FUNCTION)

        required_deps = dependencies + [
            "snowflake-snowpark-python<2",
            "fastparquet<2023.11",
            "pyarrow<14",
            "cachetools<5",
        ]

        @sproc(  # type: ignore[misc]
            is_permanent=False,
            name=search_sproc_name,
            packages=required_deps,  # type: ignore[arg-type]
            replace=True,
            session=session,
            anonymous=True,
            imports=imports,  # type: ignore[arg-type]
            statement_params=sproc_statement_params,
        )
        def _distributed_search(
            session: Session,
            imports: List[str],
            stage_estimator_file_name: str,
            input_cols: List[str],
            label_cols: Optional[List[str]],
        ) -> str:
            import os
            import time
            from typing import Iterator

            import cloudpickle as cp
            import pandas as pd
            import pyarrow.parquet as pq
            from sklearn.metrics import check_scoring
            from sklearn.metrics._scorer import _check_multimetric_scoring

            for import_name in udf_imports:
                importlib.import_module(import_name)

            data_files = [
                filename
                for filename in os.listdir(sys._xoptions["snowflake_import_directory"])
                if filename.startswith(temp_stage_name)
            ]
            partial_df = [
                pq.read_table(os.path.join(sys._xoptions["snowflake_import_directory"], file_name)).to_pandas()
                for file_name in data_files
            ]
            df = pd.concat(partial_df, ignore_index=True)
            df.columns = [identifier.get_inferred_name(col) for col in df.columns]

            X = df[input_cols]
            y = df[label_cols].squeeze() if label_cols else None

            local_estimator_file_name = get_temp_file_path()
            session.file.get(stage_estimator_file_name, local_estimator_file_name)

            local_estimator_file_path = os.path.join(
                local_estimator_file_name, os.listdir(local_estimator_file_name)[0]
            )
            with open(local_estimator_file_path, mode="r+b") as local_estimator_file_obj:
                estimator = cp.load(local_estimator_file_obj)["estimator"]

            cv_orig = check_cv(estimator.cv, y, classifier=is_classifier(estimator.estimator))
            indices = [test for _, test in cv_orig.split(X, y)]
            local_indices_file_name = get_temp_file_path()
            with open(local_indices_file_name, mode="w+b") as local_indices_file_obj:
                cp.dump(indices, local_indices_file_obj)

            # Put locally serialized indices on stage.
            put_result = session.file.put(
                local_indices_file_name,
                temp_stage_name,
                auto_compress=False,
                overwrite=True,
            )
            indices_location = put_result[0].target
            imports.append(f"@{temp_stage_name}/{indices_location}")
            indices_len = len(indices)

            assert estimator is not None

            @cachetools.cached(cache={})
            def _load_data_into_udf() -> Tuple[
                Dict[str, pd.DataFrame],
                Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV],
                pd.DataFrame,
                int,
                List[Dict[str, Any]],
            ]:
                import pyarrow.parquet as pq

                data_files = [
                    filename
                    for filename in os.listdir(sys._xoptions["snowflake_import_directory"])
                    if filename.startswith(temp_stage_name)
                ]
                partial_df = [
                    pq.read_table(os.path.join(sys._xoptions["snowflake_import_directory"], file_name)).to_pandas()
                    for file_name in data_files
                ]
                df = pd.concat(partial_df, ignore_index=True)
                df.columns = [identifier.get_inferred_name(col) for col in df.columns]

                # load estimator
                local_estimator_file_path = os.path.join(
                    sys._xoptions["snowflake_import_directory"], f"{estimator_location}"
                )
                with open(local_estimator_file_path, mode="rb") as local_estimator_file_obj:
                    estimator_objects = cp.load(local_estimator_file_obj)
                    estimator = estimator_objects["estimator"]
                    params_to_evaluate = estimator_objects["param_grid"]

                # load indices
                local_indices_file_path = os.path.join(
                    sys._xoptions["snowflake_import_directory"], f"{indices_location}"
                )
                with open(local_indices_file_path, mode="rb") as local_indices_file_obj:
                    indices = cp.load(local_indices_file_obj)

                argspec = inspect.getfullargspec(estimator.fit)
                args = {"X": df[input_cols]}

                if label_cols:
                    label_arg_name = "Y" if "Y" in argspec.args else "y"
                    args[label_arg_name] = df[label_cols].squeeze()

                if sample_weight_col is not None and "sample_weight" in argspec.args:
                    args["sample_weight"] = df[sample_weight_col].squeeze()
                return args, estimator, indices, len(df), params_to_evaluate

            class SearchCV:
                def __init__(self) -> None:
                    args, estimator, indices, data_length, params_to_evaluate = _load_data_into_udf()
                    self.args = args
                    self.estimator = estimator
                    self.indices = indices
                    self.data_length = data_length
                    self.params_to_evaluate = params_to_evaluate

                def process(self, params_idx: int, idx: int) -> Iterator[Tuple[str]]:
                    if hasattr(estimator, "param_grid"):
                        self.estimator.param_grid = self.params_to_evaluate[params_idx]
                    else:
                        self.estimator.param_distributions = self.params_to_evaluate[params_idx]
                    full_indices = np.array([i for i in range(self.data_length)])
                    test_indice = self.indices[idx]
                    train_indice = np.setdiff1d(full_indices, test_indice)
                    self.estimator.cv = [(train_indice, test_indice)]
                    self.estimator.fit(**self.args)
                    binary_cv_results = None
                    with io.BytesIO() as f:
                        cp.dump(self.estimator.cv_results_, f)
                        f.seek(0)
                        binary_cv_results = f.getvalue().hex()
                    yield (binary_cv_results,)

                def end_partition(self) -> None:
                    ...

            session.udtf.register(
                SearchCV,
                output_schema=StructType([StructField("CV_RESULTS", StringType())]),
                input_types=[IntegerType(), IntegerType()],
                name=random_udtf_name,
                packages=required_deps,  # type: ignore[arg-type]
                replace=True,
                is_permanent=False,
                imports=imports,  # type: ignore[arg-type]
                statement_params=udtf_statement_params,
            )

            HP_TUNING = F.table_function(random_udtf_name)

            idx_length = int(indices_len)
            params_length = len(param_grid)
            idxs = [i for i in range(idx_length)]
            param_indices, training_indices = [], []
            for param_idx, cv_idx in product([param_index for param_index in range(params_length)], idxs):
                param_indices.append(param_idx)
                training_indices.append(cv_idx)

            pd_df = pd.DataFrame(
                {
                    "PARAMS": param_indices,
                    "TRAIN_IND": training_indices,
                    "PARAM_INDEX": [i for i in range(idx_length * params_length)],
                }
            )
            df = session.create_dataframe(pd_df)
            results = df.select(
                F.cast(df["PARAM_INDEX"], IntegerType()).as_("PARAM_INDEX"),
                (HP_TUNING(df["PARAMS"], df["TRAIN_IND"]).over(partition_by=df["PARAM_INDEX"])),
            )

            # cv_result maintains the original order
            multimetric = False
            cv_results_ = dict()
            scorers = set()
            for i, val in enumerate(results.select("CV_RESULTS").sort(col("PARAM_INDEX")).collect()):
                # retrieved string had one more double quote in the front and end of the string.
                # use [1:-1] to remove the extra double quotes
                hex_str = bytes.fromhex(val[0])
                with io.BytesIO(hex_str) as f_reload:
                    each_cv_result = cp.load(f_reload)
                    for k, v in each_cv_result.items():
                        cur_cv = i % idx_length
                        key = k
                        if "split0_test_" in k:
                            # For multi-metric evaluation, the scores for all the scorers are available in the
                            # cv_results_ dict at the keys ending with that scorer’s name ('_<scorer_name>')
                            # instead of '_score'.
                            scorers.add(k[len("split0_test_") :])
                            key = k.replace("split0_test", f"split{cur_cv}_test")
                        elif k.startswith("param"):
                            if cur_cv != 0:
                                key = False
                        if key:
                            if key not in cv_results_:
                                cv_results_[key] = v
                            else:
                                cv_results_[key] = np.concatenate([cv_results_[key], v])

            multimetric = len(scorers) > 1
            # Use numpy to re-calculate all the information in cv_results_ again
            # Generally speaking, reshape all the results into the (scorers+2, idx_length, params_length) shape,
            # and average them by the idx_length;
            # idx_length is the number of cv folds; params_length is the number of parameter combinations
            scores = [
                np.reshape(
                    np.concatenate([cv_results_[f"split{cur_cv}_test_{score}"] for cur_cv in range(idx_length)]),
                    (idx_length, -1),
                )
                for score in scorers
            ]

            fit_score_test_matrix = np.stack(
                [
                    np.reshape(cv_results_["mean_fit_time"], (idx_length, -1)),
                    np.reshape(cv_results_["mean_score_time"], (idx_length, -1)),
                ]
                + scores
            )

            mean_fit_score_test_matrix = np.mean(fit_score_test_matrix, axis=1)
            std_fit_score_test_matrix = np.std(fit_score_test_matrix, axis=1)
            cv_results_["std_fit_time"] = std_fit_score_test_matrix[0]
            cv_results_["mean_fit_time"] = mean_fit_score_test_matrix[0]
            cv_results_["std_score_time"] = std_fit_score_test_matrix[1]
            cv_results_["mean_score_time"] = mean_fit_score_test_matrix[1]
            for idx, score in enumerate(scorers):
                cv_results_[f"std_test_{score}"] = std_fit_score_test_matrix[idx + 2]
                cv_results_[f"mean_test_{score}"] = mean_fit_score_test_matrix[idx + 2]
                # re-compute the ranking again with mean_test_<score>.
                cv_results_[f"rank_test_{score}"] = rankdata(-cv_results_[f"mean_test_{score}"], method="min")
                # The best param is the highest ranking (which is 1) and we choose the first time ranking 1 appeared.
                # If all scores are `nan`, `rankdata` will also produce an array of `nan` values.
                # In that case, default to first index.
                best_param_index = (
                    np.where(cv_results_[f"rank_test_{score}"] == 1)[0][0]
                    if not np.isnan(cv_results_[f"rank_test_{score}"]).all()
                    else 0
                )

            estimator.cv_results_ = cv_results_
            estimator.multimetric_ = multimetric

            # Reconstruct the sklearn estimator.
            refit_metric = "score"
            if callable(estimator.scoring):
                scorers = estimator.scoring
            elif estimator.scoring is None or isinstance(estimator.scoring, str):
                scorers = check_scoring(estimator.estimator, estimator.scoring)
            else:
                scorers = _check_multimetric_scoring(estimator.estimator, estimator.scoring)
                estimator._check_refit_for_multimetric(scorers)
                refit_metric = original_refit

            estimator.scorer_ = scorers

            # check refit_metric now for a callabe scorer that is multimetric
            if callable(estimator.scoring) and estimator.multimetric_:
                refit_metric = original_refit

            # For multi-metric evaluation, store the best_index_, best_params_ and
            # best_score_ iff refit is one of the scorer names
            # In single metric evaluation, refit_metric is "score"
            if original_refit or not estimator.multimetric_:
                estimator.best_index_ = estimator._select_best_index(original_refit, refit_metric, cv_results_)
                if not callable(original_refit):
                    # With a non-custom callable, we can select the best score
                    # based on the best index
                    estimator.best_score_ = cv_results_[f"mean_test_{refit_metric}"][estimator.best_index_]
                estimator.best_params_ = cv_results_["params"][best_param_index]

            if original_refit:
                estimator.best_estimator_ = clone(estimator.estimator).set_params(
                    **clone(estimator.best_params_, safe=False)
                )

                # Let the sproc use all cores to refit.
                estimator.n_jobs = -1 if not estimator.n_jobs else estimator.n_jobs

                # process the input as args
                argspec = inspect.getfullargspec(estimator.fit)
                args = {"X": X}
                if label_cols:
                    label_arg_name = "Y" if "Y" in argspec.args else "y"
                    args[label_arg_name] = y
                if sample_weight_col is not None and "sample_weight" in argspec.args:
                    args["sample_weight"] = df[sample_weight_col].squeeze()
                estimator.refit = original_refit
                refit_start_time = time.time()
                estimator.best_estimator_.fit(**args)
                refit_end_time = time.time()
                estimator.refit_time_ = refit_end_time - refit_start_time

                if hasattr(estimator.best_estimator_, "feature_names_in_"):
                    estimator.feature_names_in_ = estimator.best_estimator_.feature_names_in_

            local_result_file_name = get_temp_file_path()

            with open(local_result_file_name, mode="w+b") as local_result_file_obj:
                cp.dump(estimator, local_result_file_obj)

            session.file.put(
                local_result_file_name,
                temp_stage_name,
                auto_compress=False,
                overwrite=True,
            )

            # Note: you can add something like  + "|" + str(df) to the return string
            # to pass debug information to the caller.
            return str(os.path.basename(local_result_file_name))

        sproc_export_file_name = _distributed_search(
            session,
            imports,
            stage_estimator_file_name,
            input_cols,
            label_cols,
        )

        local_estimator_path = get_temp_file_path()
        session.file.get(
            posixpath.join(temp_stage_name, sproc_export_file_name),
            local_estimator_path,
        )

        with open(os.path.join(local_estimator_path, sproc_export_file_name), mode="r+b") as result_file_obj:
            fit_estimator = cp.load(result_file_obj)

        cleanup_temp_files([local_estimator_path])

        return fit_estimator

    def train(self) -> object:
        """
        Runs hyper parameter optimization by distributing the tasks across warehouse.

        Returns:
            Trained model
        """
        model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
        assert isinstance(self.estimator, model_selection.GridSearchCV) or isinstance(
            self.estimator, model_selection.RandomizedSearchCV
        )
        if hasattr(self.estimator.estimator, "n_jobs") and self.estimator.estimator.n_jobs in [
            None,
            -1,
        ]:
            self.estimator.estimator.n_jobs = DEFAULT_UDTF_NJOBS

        if isinstance(self.estimator, model_selection.GridSearchCV):
            param_grid = model_selection.ParameterGrid(self.estimator.param_grid)
        elif isinstance(self.estimator, model_selection.RandomizedSearchCV):
            param_grid = model_selection.ParameterSampler(
                self.estimator.param_distributions,
                n_iter=self.estimator.n_iter,
                random_state=self.estimator.random_state,
            )
        return self.fit_search_snowpark(
            param_grid=param_grid,
            dataset=self.dataset,
            session=self.session,
            estimator=self.estimator,
            dependencies=model_spec.pkgDependencies,
            udf_imports=["sklearn"],
            input_cols=self.input_cols,
            label_cols=self.label_cols,
            sample_weight_col=self.sample_weight_col,
        )
