import inspect
import json
import logging
import warnings
from collections import OrderedDict

import gorilla
from slack_sdk import WebClient
from slack_sdk.errors import SlackApiError
import numpy as np
import tensorflow
from tensorflow.keras.callbacks import CSVLogger
from Accuinsight.modeler.core import func, path, get
from Accuinsight.modeler.core.func import get_time
from Accuinsight.modeler.core.LcConst import LcConst
from Accuinsight.modeler.core.LcConst.LcConst import RUN_NAME_TENSORFLOW
from Accuinsight.modeler.core.Run.RunInfo.RunInfo import set_current_runs, clear_runs, set_model_json_path, \
    set_visual_csv_path, set_visual_json_path, set_best_model_json_path, set_best_model_h5_path, \
    set_python_dependencies, set_run_name, set_model_file_path, set_prefix_path, set_shap_json_path
from Accuinsight.modeler.core.get_for_visual import roc_pr_curve, get_true_y, get_visual_info_regressor
from Accuinsight.modeler.clients.modeler_api import LifecycleRestApi
from Accuinsight.modeler.utils.dependency.dependencies import gather_sources_and_dependencies
from Accuinsight.modeler.utils.dl_utils import delete_files_except_best, get_best_model_path
from Accuinsight.modeler.utils.os_getenv import is_in_ipython, get_current_notebook, get_os_env
from Accuinsight.modeler.core.feature_contribution import shap_value
from Accuinsight.Lifecycle.common import Common

logging.basicConfig(level=logging.INFO,
                    format='%(message)s')

warnings.filterwarnings("ignore")


class accuinsight(Common):
    def __init__(self):
        super(Common, self).__init__()

    def autolog(self, tag=None, best_weights=False, model_monitor=False):
        global description, endpoint, var_model_file_path, message, thresholds, token, channel_id, best_weights_on, run_id, shap_on, feature_name
        description = tag
        endpoint = self.endpoint
        message = self.message
        thresholds = self.thresholds
        token = self.token
        channel_id = self.channel_id
    
        if best_weights:
            best_weights_on = True
        else:
            best_weights_on = False
        if model_monitor:
            shap_on = True
            if self.feature_name is None:
                try:
                    feature_name = get.feature_name(self.save_path, self.StorageInfo, self.target_name)

                except:
                    pass      # when user did not use <get_file> or <set_feature> function
            else:
                feature_name = self.feature_name
        else:
            shap_on = False
        run_id = None

        is_notebook = is_in_ipython()
        if is_notebook:
            var_model_file_path = get_current_notebook()
            _caller_globals = inspect.stack()[1][0].f_globals
            (
                mainfile,
                sources,
                dependencies
            ) = gather_sources_and_dependencies(
                globs=_caller_globals,
                save_git_info=False
            )
        else:
            _caller_globals = inspect.stack()[1][0].f_globals
            (
                mainfile,
                sources,
                dependencies
            ) = gather_sources_and_dependencies(
                globs=_caller_globals,
                save_git_info=True
            )
            var_model_file_path = mainfile['filename']

        class TrainHistoryCallbacks(tensorflow.keras.callbacks.Callback):
            def __init__(self, verbose=1, mode='auto'):
                super(TrainHistoryCallbacks, self).__init__()
                self.verbose = verbose
                self.best_epochs = 0
                self.epochs_since_last_save = 0
                self.mode = mode

            def on_train_begin(self, logs={}):
                if best_weights_on and shap_on:
                    logging.info('Using autolog(best_weights=True, model_monitor=True)')
                elif best_weights_on == True and shap_on == False:
                    logging.info('Using autolog(best_weights=True, model_monitor=False)')
                elif best_weights_on == False and shap_on == True:
                    logging.info('Using autolog(best_weights=False, model_monitor=True)')
                else:
                    logging.info('Using autolog(best_weights=False, model_monitor=False)')
                
                global start
                start = get_time.now()
                opt = self.model.optimizer.get_config()
                opt_key = list(opt.keys())[1:]
                opt_result = {k: np.float64(opt[k]) for k in opt_key}
                self.model_summary = OrderedDict()
                self.model_summary['data_version'] = endpoint
                self.model_summary['model_description'] = description
                self.model_summary['logging_time'] = get_time.logging_time()
                self.model_summary['run_id'] = func.get_run_id()
                self.model_summary['model_type'] = get.model_type(self.model)

                if hasattr(self.model.loss, 'get_config'):
                    self.model_summary['loss_function'] = self.model.loss.get_config()['name']
                else:
                    self.model_summary['loss_function'] = self.model.loss

                self.model_summary['optimizer_info'] = {opt['name']: opt_result}

                '''[get best model] on_train_begin '''
                self.best_weights = self.model.get_weights()

                self.dict_path = path.get_file_path(self.model, usedFramework='tensorflow')

                set_prefix_path(self.dict_path[LcConst.RUN_PREFIX_PATH])

                set_run_name(self.model_summary['model_type'], self.model_summary['run_id'])
                set_python_dependencies(py_depenpency=dependencies)

#                print('dict_path: ', self.dict_path)         ######################
        
            '''[get best model] on_epoch_end '''
            def on_epoch_end(self, epoch, logs=None):
                logs = logs or {}
                if epoch == 0:
                    if len(self.model.metrics_names) == 1 and 'loss' in self.model.metrics_names:
                        self.monitor = 'val_loss'
                    elif len(self.model.metrics_names) >= 2:
                        self.monitor = 'val_' + self.model.metrics_names[1]
                    # set monitoring option
                    if self.mode not in ['auto', 'min', 'max']:
                        warnings.warn('GetBest mode %s is unknown, '
                                      'fallback to auto mode.' % (self.mode), RuntimeWarning)
                        self.mode = 'auto'
                    if self.mode == 'min':
                        self.monitor_op = np.less
                        self.best = np.Inf
                    elif self.mode == 'max':
                        self.monitor_op = np.greater
                        self.best = -np.Inf
                    else:
                        if 'acc' in self.monitor or 'f1' in self.monitor:
                            self.monitor_op = np.greater
                            self.best = -np.Inf
                        else:
                            self.monitor_op = np.less
                            self.best = np.Inf
                else:
                    pass
                
                # Using best_weights
                if best_weights_on:
                    
                    # update best_weights
                    self.epochs_since_last_save += 1
                    if self.epochs_since_last_save >= 1:
                        self.epochs_since_last_save = 0
                        current = logs.get(self.monitor)
                        if current is None:
                            warnings.warn('Can pick best model only with %s available, '
                                      'skipping.' % (self.monitor), RuntimeWarning)
                        else:
                            if self.monitor_op(current, self.best):
                                self.best = current
                                self.best_epochs = epoch + 1
                                self.best_weights = self.model.get_weights()
                            else:
                                pass

                    self.current_value = current
    
                # Not using best_weights
                else:
                    self.last_epoch_metric = logs.get(self.monitor)
                    self.best_epochs = epoch + 1
                    current = logs.get(self.monitor)
                    self.current_value = current

                # model save path
                run_id = self.model_summary['model_type'] + '-' + self.model_summary['run_id']
                common_path = self.dict_path['save_model_path'] + run_id + '-epoch-' + str(epoch + 1).zfill(5) + '-' + self.monitor + '-' + str(current).zfill(5)
                save_model_path = common_path + '.json'
                save_weights_path = common_path + '.h5'
                # model to JSON
                model_json = self.model.to_json()
                with open(save_model_path, "w") as json_file:
                    json_file.write(model_json)
                # weights to H5
                self.model.save_weights(save_weights_path)

            def on_train_end(self, logs={}):
                '''[get best model] on_train_end '''
                if best_weights_on:
                    if self.verbose > 0:
                        print('\nUsing epoch %05d with %s: %0.5f' % (self.best_epochs, self.monitor, self.best))
                    self.model.set_weights(self.best_weights)  # set best model's weights
                    
                    self.model_summary['selected_metrics'] = {self.monitor: np.float64(self.best)}
                else:
                    self.model_summary['selected_metrics'] = {self.monitor: np.float64(self.last_epoch_metric)}
    
          #      print('model_summary: ', self.model_summary)   ##############################
                end = get_time.now()
                self.model_summary['time_delta'] = str(end - start)

                # path for model_info.json
                self.path_for_setting_model_json = self.dict_path['model_json']
                set_model_json_path(self.path_for_setting_model_json)

                model_json_full_path = self.dict_path['model_json_full']

                with open(model_json_full_path, 'w', encoding='utf-8') as save_file:
                    json.dump(self.model_summary, save_file, indent="\t")

                delete_files_except_best(run_id=self.model_summary['run_id'], epochs=str(self.best_epochs),
                                         path=self.dict_path)

                path_for_setting_model_json = self.dict_path['save_model_dir'] + \
                                              get_best_model_path(run_id=self.model_summary['run_id'],
                                                                  path=self.dict_path)['json']
                path_for_setting_model_h5 = self.dict_path['save_model_dir'] + \
                                            get_best_model_path(run_id=self.model_summary['run_id'],
                                                                path=self.dict_path)['h5']
                set_best_model_json_path(path_for_setting_model_json)
                set_best_model_h5_path(path_for_setting_model_h5)

                start_ts = int(start.timestamp())
                end_ts = int(end.timestamp())
                delta_ts = end_ts - start_ts

                clear_runs(start_ts, end_ts, delta_ts)
                accuinsight._send_message(metric=self.monitor,
                                          current_value=self.current_value,
                                          message=message,
                                          thresholds=thresholds,
                                          token=token,
                                          channel_id=channel_id)
                env_value = get_os_env('ENV')
                modeler_rest = LifecycleRestApi(env_value[LcConst.BACK_END_API_URL],
                                                env_value[LcConst.BACK_END_API_PORT],
                                                env_value[LcConst.BACK_END_API_URI])
                modeler_rest.lc_create_run()
                accuinsight._off_autolog()

        class visualCallbacks(tensorflow.keras.callbacks.Callback):
            def __init__(self, x_validation=None, y_validation=None):
                super(visualCallbacks, self).__init__()
                self.x_val = x_validation
                self.y_val = y_validation

            def on_train_end(self, logs={}):

                self.dict_path = path.get_file_path(self.model, usedFramework='tensorflow')
                # path for visual.json
                path_for_setting_visual_json = self.dict_path['visual_json']
                visual_json_full_path = self.dict_path['visual_json_full']
                set_visual_json_path(path_for_setting_visual_json)

                # classification
                if get.is_classification(self.model):
                    visual_classification_json = roc_pr_curve(self.model, self.x_val, self.y_val)

                    with open(visual_json_full_path, 'w', encoding='utf-8') as save_file:
                        json.dump(visual_classification_json, save_file, indent="\t")

                # regression
                else:
                    visual_regression_json = OrderedDict()
                    visual_regression_json['True_y'] = get_true_y(self.y_val)
                    visual_regression_json['Predicted_y'] = get_visual_info_regressor(self.model, self.x_val)

                    with open(visual_json_full_path, 'w', encoding='utf-8') as save_file:
                        json.dump(visual_regression_json, save_file, indent="\t")

        class shapCallbacks(tensorflow.keras.callbacks.Callback):
            def __init__(self, trainX, feature_name, run_id, trigger = shap_on):
                super(shapCallbacks, self).__init__()
                self.trainX = trainX
                self.trigger = shap_on
                self.run_id = run_id
                self.feature_name_in_shap = feature_name

            def on_train_end(self, logs={}):
                if self.trigger:
                    self.shap_value = shap_value(self.model, self.trainX, self.feature_name_in_shap)
                    
#                     func.insertDB(self.shap_value, 2)  # 수정: 2 -> self.run_id
                    
                    self.dict_path = path.get_file_path(self.model, usedFramework='tensorflow')

                    # path for shap.json
                    shap_json_full_path = self.dict_path['shap_json_full']
                    set_shap_json_path(self.dict_path['shap_json'])
                    
                    with open(shap_json_full_path, 'w', encoding='utf-8') as save_file:
                        json.dump(self.shap_value, save_file, indent='\t')

                else:
                    pass
                    
        def run_and_log_function(self, original, x, y, kwargs):
            dict_path = path.get_file_path(self, usedFramework='tensorflow')
            path_for_setting_visual_csv = dict_path['visual_csv']
            visual_csv_full_path = dict_path['visual_csv_full']

            # set current run
            set_current_runs(RUN_NAME_TENSORFLOW)
            set_model_file_path(var_model_file_path)
            set_visual_csv_path(path_for_setting_visual_csv)

            csv_logger = CSVLogger(visual_csv_full_path, append=True, separator=';')
            
            # get train data(x) for computing shap value
            if 'x':
                x_train = x
            if shap_on:
                get_shap = shapCallbacks(x_train, feature_name, run_id, trigger = shap_on)
            else:
                pass
            
            ''' save json for visualization '''
            # using validation_data argument
            if 'validation_data' in kwargs:
                validation_set = kwargs['validation_data']

                try:
                    x_val = validation_set[0]
                    y_val = validation_set[1]

                except:
                    iterator = iter(validation_set)
                    valid_set = next(iterator)
                    x_val = valid_set[0].numpy()
                    y_val = valid_set[1].numpy()

                get_visual = visualCallbacks(x_validation=x_val, y_validation=y_val)

            else:
                raise ValueError('"validation_data" does not exist.')

            if 'callbacks' in kwargs:
                kwargs['callbacks'] += [csv_logger]
                kwargs['callbacks'] += [get_visual]
                if shap_on:
                    kwargs['callbacks'] += [get_shap]
                else:
                    pass
                kwargs['callbacks'] += [TrainHistoryCallbacks()]
                

            else:
                kwargs['callbacks'] = [csv_logger]
                kwargs['callbacks'] += [get_visual]
                if shap_on:
                    kwargs['callbacks'] += [get_shap]
                else:
                    pass
                kwargs['callbacks'] += [TrainHistoryCallbacks()]

            return original(self, x, y, **kwargs)

        @gorilla.patch(tensorflow.keras.Model)
        def fit(self, x, y, **kwargs):
            original = gorilla.get_original_attribute(tensorflow.keras.Model, 'fit')
            unlogged_params = ['self', 'callbacks', 'validation_data', 'verbose']
            return run_and_log_function(self, original, x, y, kwargs)

        settings = gorilla.Settings(allow_hit=True, store_hit=True)
        gorilla.apply(gorilla.Patch(tensorflow.keras.Model, 'fit', fit, settings=settings))

    def _off_autolog():
        def stop_log(self, original, args, kwargs, unlogged_params):
            fit_model = original(self, *args, **kwargs)
            return (fit_model)

        @gorilla.patch(tensorflow.keras.Model)
        def fit(self, *args, **kwargs):
            original = gorilla.get_original_attribute(tensorflow.keras.Model, 'fit')
            unlogged_params = ['self', 'x', 'y', 'callbacks', 'validation_data', 'verbose']
            return stop_log(self, original, args, kwargs, unlogged_params)

        settings = gorilla.Settings(allow_hit=True, store_hit=True)
        gorilla.apply(gorilla.Patch(tensorflow.keras.Model, 'fit', fit, settings=settings))

