#  (C) Copyright IBM Corp. 2021.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import abc
from os import environ

import unittest
import time

from sklearn.pipeline import Pipeline

from ibm_watson_machine_learning import APIClient
from ibm_watson_machine_learning.experiment import AutoAI
from ibm_watson_machine_learning.deployment import WebService, Batch
from ibm_watson_machine_learning.workspace import WorkSpace
from ibm_watson_machine_learning.helpers.connections import DataConnection
from ibm_watson_machine_learning.experiment.autoai.optimizers import RemoteAutoPipelines
from ibm_watson_machine_learning.tests.utils import (get_wml_credentials, get_cos_credentials, get_space_id,
                                                     is_cp4d)
from ibm_watson_machine_learning.tests.utils.cleanup import space_cleanup, delete_model_deployment
from ibm_watson_machine_learning.utils.autoai.enums import PredictionType, RunStateTypes

from ibm_watson_machine_learning.tests.utils.assertions import get_and_predict_all_pipelines_as_lale, \
    validate_autoai_experiment

from ibm_watson_machine_learning.utils.autoai.utils import chose_model_output


class AbstractTestAutoAISyncNotebookFlow(abc.ABC):
    """
    The abstract tests which covers:
    - training AutoAI model on a dataset
    - autogenerated notebook scenario

    In order to execute test connection definitions must be provided
    in inheriting classes.
    """

    bucket_name = environ.get('BUCKET_NAME', "wml-autoaitests-qa")
    pod_version = environ.get('KB_VERSION', None)
    space_name = environ.get('SPACE_NAME', 'regression_tests_sdk_space')

    cos_endpoint = "https://s3.us.cloud-object-storage.appdomain.cloud"
    results_cos_path = 'results_wml_autoai'

    required_experiment_metadata_fields = ['prediction_type', 'prediction_column', 'holdout_size', 'scoring',
                                           'csv_separator', 'random_state', 'max_number_of_estimators',
                                           'training_data_reference', 'training_result_reference']

    experiment_metadata_fields_to_be_updated_after_training = ['training_data_reference', 'training_result_reference',
                                                               'holdout_size', 'scoring', 'csv_separator',
                                                               'random_state', 'max_number_of_estimators']

    # to be set in every child class:
    OPTIMIZER_NAME = "AutoAI regression test"

    SPACE_ONLY = True
    HISTORICAL_RUNS_CHECK = True

    experiment_info = dict(name=OPTIMIZER_NAME,
                           desc='test description',
                           prediction_type=PredictionType.MULTICLASS,
                           prediction_column='species',
                           autoai_pod_version=pod_version
                           )

    experiment_metadata: dict = None

    wml_client: 'APIClient' = None
    experiment: 'AutoAI' = None
    remote_auto_pipelines: 'RemoteAutoPipelines' = None
    wml_credentials = None
    cos_credentials = None
    pipeline_opt: 'RemoteAutoPipelines' = None
    service: 'WebService' = None
    service_batch: 'Batch' = None

    cos_resource_instance_id = None
    experiment_info: dict = None

    best_pipeline_name = False
    trained_pipeline_details = None
    run_id = None
    prev_run_id = None
    data_connection = None
    results_connection = None
    train_data = None

    pipeline: 'Pipeline' = None
    lale_pipeline = None
    deployed_pipeline = None
    hyperopt_pipelines = None
    new_pipeline = None
    new_sklearn_pipeline = None
    X_train = None
    y_train = None
    X_holdout = None
    y_holdout = None

    project_id = None
    space_id = None

    asset_id = None
    connection_id = None

    @classmethod
    def setUpClass(cls) -> None:
        """
        Load WML credentials from config.ini file based on ENV variable.
        """
        cls.wml_credentials = get_wml_credentials()
        cls.wml_client = APIClient(wml_credentials=cls.wml_credentials.copy())

        cls.cos_credentials = get_cos_credentials()
        cls.cos_endpoint = cls.cos_credentials.get('endpoint_url')
        cls.cos_resource_instance_id = cls.cos_credentials.get('resource_instance_id')

        cls.project_id = cls.wml_credentials.get('project_id')

    def test_00a_space_cleanup(self):
        space_checked = False
        while not space_checked:
            space_cleanup(self.wml_client,
                          get_space_id(self.wml_client, self.space_name,
                                       cos_resource_instance_id=self.cos_resource_instance_id),
                          days_old=7)
            space_id = get_space_id(self.wml_client, self.space_name,
                                    cos_resource_instance_id=self.cos_resource_instance_id)
            try:
                self.assertIsNotNone(space_id, msg="space_id is None")
                space_checked = True
            except AssertionError:
                space_checked = False

        AbstractTestAutoAISync.space_id = space_id

        if self.SPACE_ONLY:
            self.wml_client.set.default_space(self.space_id)
        else:
            self.wml_client.set.default_project(self.project_id)

    def test_01_initialize_AutoAI_experiment__pass_credentials__object_initialized(self):
        if self.SPACE_ONLY:
            AbstractTestAutoAISync.experiment_run_notebook = AutoAI(wml_credentials=self.wml_credentials.copy(),
                                                                    space_id=self.space_id)
        else:
            AbstractTestAutoAISync.experiment_run_notebook = AutoAI(wml_credentials=self.wml_credentials.copy(),
                                                                    project_id=self.project_id)

        self.assertIsInstance(self.experiment_run_notebook, AutoAI, msg="Experiment is not of type AutoAI.")

    @abc.abstractmethod
    def test_02_data_reference_setup(self):
        pass

    @abc.abstractmethod
    def test_02b_experiment_metadata_setup(self):
        pass

    def test_03_initialize_optimizer(self):
        AbstractTestAutoAISync.experiment_info = validate_autoai_experiment(self.experiment_info, self.pod_version)

        AbstractTestAutoAISync.remote_auto_pipelines = self.experiment_run_notebook.optimizer(
            **AbstractTestAutoAISync.experiment_info)

        self.assertIsInstance(self.remote_auto_pipelines, RemoteAutoPipelines,
                              msg="experiment.optimizer did not return RemoteAutoPipelines object")

    def test_04_get_configuration_parameters_of_remote_auto_pipeline(self):
        parameters = self.remote_auto_pipelines.get_params()
        print(parameters)

        self.assertIsInstance(parameters, dict, msg='Config parameters are not a dictionary instance.')

    def test_05_fit_run_training_of_auto_ai_in_wml(self):
        AbstractTestAutoAISync.trained_pipeline_details = self.remote_auto_pipelines.fit(
            training_data_reference=[self.data_connection],
            training_results_reference=self.results_connection,
            background_mode=True)

        AbstractTestAutoAISync.run_id = self.trained_pipeline_details['metadata']['id']
        self.assertIsNotNone(self.data_connection.auto_pipeline_params,
                             msg='DataConnection auto_pipeline_params was not updated.')

    def test_06a_get_run_status(self):
        status = self.remote_auto_pipelines.get_run_status()
        self.assertNotEqual(status, RunStateTypes.COMPLETED)

    def test_06b_get_run_details(self):
        parameters = self.remote_auto_pipelines.get_run_details()
        # print(parameters)
        self.assertIsNotNone(parameters)
        self.assertIn(self.run_id, str(parameters))

    def test_07_wait_for_first_pipeline(self):

        print(f"Run status = {self.remote_auto_pipelines.get_run_status()}")
        # note: check if first pipeline was generated

        metrics = self.wml_client.training.get_details(self.run_id)['entity']['status'].get('metrics', [])
        while chose_model_output("1") not in str(metrics) and self.remote_auto_pipelines.get_run_status() not in [
            'failed', 'canceled']:
            time.sleep(5)
            print(".", end=" ")
            metrics = self.wml_client.training.get_details(self.run_id)['entity']['status'].get('metrics', [])
        # end note

        status = self.remote_auto_pipelines.get_run_status()
        run_details = self.remote_auto_pipelines.get_run_details().get('entity')
        self.assertNotIn(status, ['failed', 'canceled'], msg=f"Training finished with status {status}. \n"
                                                             f"Details: {run_details.get('status')}")
        print("\n 1st pipeline completed")

    def test_08_update_experiment_metadata_data_referances(self):
        pipeline_parameters = self.remote_auto_pipelines.get_params()
        for param_to_add in self.experiment_metadata_fields_to_be_updated_after_training:
            if param_to_add == 'training_data_reference':
                self.experiment_metadata[param_to_add] = self.remote_auto_pipelines.get_data_connections()
            elif param_to_add == 'test_data_references':
                self.experiment_metadata[param_to_add] = self.remote_auto_pipelines.get_data_connections()
            elif param_to_add == 'training_result_reference':
                results_reference = self.remote_auto_pipelines.get_run_details()['entity']['results_reference']
                self.experiment_metadata[param_to_add] = DataConnection._from_dict(results_reference)
            if param_to_add not in self.experiment_metadata:
                self.experiment_metadata[param_to_add] = pipeline_parameters.get(param_to_add)

    def test_09_validate_experiment_metadata(self):
        missing_params = []
        for r_parameter in self.required_experiment_metadata_fields:
            if r_parameter not in self.experiment_metadata:
                missing_params.append(r_parameter)

        self.assertEqual(len(missing_params), 0, f"Missing in experiment metadata: {missing_params}")

    ###########################################
    #             NOTEBOOK CELLS              #
    ###########################################

    def test_10_notebook_wml_credential_and_experiment(self):
        wml_credentials = {
            "apikey": self.wml_credentials.get('apikey'),
            "url": self.experiment_metadata['deployment_url']
        }

        self.pipeline_optimizer = AutoAI(wml_credentials,
                                         project_id=self.experiment_metadata['project_id']).runs.get_optimizer(
            metadata=self.experiment_metadata)

    def test_10b_download_data_with_holdout_not_in_experiment_notebook(self):
        AbstractTestAutoAISync.X_train, AbstractTestAutoAISync.X_holdout, AbstractTestAutoAISync.y_train, AbstractTestAutoAISync.y_holdout = \
            self.pipeline_optimizer.get_data_connections()[0].read(
                with_holdout_split=True)

        print("train data sample:")
        print(AbstractTestAutoAISync.X_train)
        print(AbstractTestAutoAISync.y_train)
        print("holdout data sample:")
        print(AbstractTestAutoAISync.X_holdout)
        print(AbstractTestAutoAISync.y_holdout)

        self.assertGreater(len(AbstractTestAutoAISync.X_train), 0)
        self.assertGreater(len(AbstractTestAutoAISync.X_holdout), 0)
        self.assertEqual(len(AbstractTestAutoAISync.X_train), len(AbstractTestAutoAISync.y_train))
        self.assertEqual(len(AbstractTestAutoAISync.X_holdout), len(AbstractTestAutoAISync.y_holdout))

    def test_11_get_params(self):
        parameters = self.pipeline_optimizer.get_params()
        print(parameters)

        self.assertIsInstance(parameters, dict, msg='Config parameters are not a dictionary instance.')

    def test_12_get_summary(self):
        summary = self.pipeline_optimizer.summary()
        self.best_pipeline_name = list(summary.index)[0]

        print(summary)

        self.assertGreater(len(summary), 0,
                           msg=f"Summary DataFrame is empty. ")

        # check if pipelines are not duplicated
        self.assertEqual(len(summary.index.unique()), len(summary),
                         msg="Some pipeline names are duplicated in the summary")

    def test_13_pipeline_model(self):
        self.pipeline_model = self.pipeline_optimizer.get_pipeline()
        from lale.operators import TrainablePipeline
        self.assertIsInstance(self.lale_pipeline, TrainablePipeline,
                              msg="Fetched pipeline is not of TrainablePipeline instance.")

        feature_importance = self.pipeline_optimizer.get_pipeline_details()['features_importance']

        from pandas import DataFrame

        self.assertIsInstance(feature_importance, DataFrame,
                              msg="Feature importance is not returned as pandas.DataFrame.")

        self.assertGreater(len(feature_importance), 0,
                           msg=f"feature_importance DataFrame is empty.")

    def test_14a_pipeline_pretty_print(self):
        pipeline_code = self.pipeline_model.pretty_print(combinators=False)

        global pipeline
        exception = None
        try:
            exec(pipeline_code, globals())
        except Exception as exception:
            self.assertIsNone(exception, msg="Pretty print from lale pipeline was not successful")
        pipeline_str = str(pipeline.export_to_sklearn_pipeline()).replace("pipelinePipeline", "Pipeline")
        scikit_pipeline_str = str(self.pipeline_model.export_to_sklearn_pipeline())
        print(f"Pretty_print pipeline {pipeline_str}\n")
        print(f"Client pipeline {scikit_pipeline_str}")

        self.assertEqual(pipeline_str, scikit_pipeline_str, msg=f"Code is not equal, pretty_print_score:\
        {scikit_pipeline_str}, pipeline score {pipeline_str}")

    @unittest.skip
    def test_14b_pipeline_pretty_print_check_if_generated_python_pipeline_code_is_correct(self):
        variables = {"X_train": self.X_train.values,
                     "X_test": self.X_test.values,
                     "y_train": self.y_train.values,
                     "y_test": self.y_test.values
                     }
        global pretty_print_score
        global pipeline

        base_string: str = "\
        import pandas as pd\r\n\
        from sklearn.metrics import r2_score"
        end_string = "\
        pipeline.fit(X_train,y_train)\r\n\
        pretty_print_score = r2_score(y_test,pipeline.predict(X_test))"

        pipeline_code = self.pipeline_model.pretty_print(combinators=False)

        from sklearn.metrics import roc_auc_score

        scikit_pipeline = self.pipeline_model.export_to_sklearn_pipeline()
        scikit_pipeline.fit(self.X_train.values, self.y_train.values)
        print(f"scikit_pipeline_params{scikit_pipeline.get_params(deep=True)}")
        pipeline_score = roc_auc_score(self.y_test.values, scikit_pipeline.predict(self.X_test.values))

        exception = None
        try:
            # print(base_string + "\r\n" + pipeline_code + "\r\n" + end_string)
            exec(base_string + "\r\n" + pipeline_code + "\r\n" + end_string, variables, globals())
            pipeline = pipeline.export_to_sklearn_pipeline()
            print(f"pretty_pipeline_params{pipeline.get_params(deep=True)}")

        except Exception as exception:
            self.assertIsNone(exception, msg="Pretty print from lale pipeline was not successful")
        print(pipeline_score)
        print(pretty_print_score)
        self.assertAlmostEqual(first=pretty_print_score, second=pipeline_score, places=1, msg=f"Scores are different at first decimal place:\
                {pretty_print_score}, pipeline score {pipeline_score}")

    ###############################
    #         DEPLOYMENT          #
    ###############################

    def test_20_deployment_creation(self):
        target_space_id = self.space_id

        from ibm_watson_machine_learning.deployment import WebService

        if self.SPACE_ONLY:
            self.assertIsNotNone(self.experiment_metadata.get('space_id'), msg="Missing space id in metadata")
            AbstractTestAutoAISync.service = WebService(
                source_wml_credentials=self.wml_credentials,
                target_wml_credentials=self.wml_credentials,
                source_project_id=self.experiment_metadata['space_id'],
                target_space_id=target_space_id
            )
        else:
            AbstractTestAutoAISync.self.assertIsNotNone(self.experiment_metadata.get('project_id'),
                                                        msg="Missing project id in metadata")
            service = WebService(
                source_wml_credentials=self.wml_credentials,
                target_wml_credentials=self.wml_credentials,
                source_project_id=self.experiment_metadata['project_id'],
                target_space_id=target_space_id
            )

        AbstractTestAutoAISync.service.create(
            model=self.best_pipeline_name,
            metadata=self.experiment_metadata,
            deployment_name='Best_pipeline_webservice'
        )
        print(AbstractTestAutoAISync.service)

    def test_21_webservice_get_params(self):
        params = AbstractTestAutoAISync.service.get_params()
        print(params)
        self.assertIsNotNone(params)

    def test_22_webservice_score_markdown_in_experiment_notebook(self):
        """
        MARKDOWN:
        If you want to work with the web service in an external Python application, follow these steps to retrieve the service object:

        Initialize the service by service = WebService(target_wml_credentials=wml_credentials,target_space_id=experiment_metadata['space_id'])
        Get deployment_id by service.list() method
        Get webservice object by service.get('deployment_id') method

        After that you can call service.score(score_records_df) method. The score() method accepts pandas.DataFrame object.
        """
        score_records_df = self.X_holdout

        predictions = self.service.score(score_records_df)

        print(predictions)
        self.assertIsNotNone(predictions)
        self.assertEqual(len(predictions['predictions'][0]['values']), len(score_records_df))

    ########
    # LALE #
    ########

    def test_31__get_pipeline__load_lale_pipeline__pipeline_loaded(self):
        AbstractTestAutoAISync.lale_pipeline = self.remote_auto_pipelines.get_pipeline()
        print(f"Fetched pipeline type: {type(self.lale_pipeline)}")

        from lale.operators import TrainablePipeline
        self.assertIsInstance(self.lale_pipeline, TrainablePipeline,
                              msg="Fetched pipeline is not of TrainablePipeline instance.")
        predictions = self.lale_pipeline.predict(
            X=self.X_values[:5])
        print(predictions)

    def test_32_get_all_pipelines_as_lale(self):
        get_and_predict_all_pipelines_as_lale(self.remote_auto_pipelines, self.X_values)
