import inspect
import json
import logging
import os
import sys
import warnings
from collections import OrderedDict
from pathlib import Path

import boto3
import gorilla
import numpy as np
import tensorflow
from slack import WebClient
from slack.errors import SlackApiError
from tensorflow.keras.callbacks import CSVLogger

from Accuinsight.modeler.core import func, path, get
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, print_run_info, \
    set_python_dependencies, set_run_name, set_model_file_path, set_prefix_path
from Accuinsight.modeler.core.func import get_time
from Accuinsight.modeler.core.get_for_visual import roc_pr_curve, 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
from Accuinsight.modeler.utils.runs_utils import get_aws_info, ProgressPercentage

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

warnings.filterwarnings("ignore")


class accuinsight(object):
    def __init__(self):
        self.BucketInfo = None
        self.endpoint = None
        self.thresholds = None
        self.token = None
        self.channel_id = None
        self.message = None

    def get_file(self, sub_path):
        self.BucketInfo = get_aws_info(sub_path)
        BUCKET_TYPE = self.BucketInfo['bucket_type']
        BUCKET_NAME = self.BucketInfo['bucket_name']
        FILE_PATH = self.BucketInfo['file_path']
        FILE_NAME = self.BucketInfo['file_name']
        FILE_TYPE = self.BucketInfo['file_type']
        FILE_DELIM = self.BucketInfo['file_delim']
        ACCESS_KEY = self.BucketInfo['my_access_key']
        SECRET_KEY = self.BucketInfo['my_secret_key']
        REGION = self.BucketInfo['region']
        URL = self.BucketInfo['endpoint']

        ## path for saving data
        save_dir = os.path.join(Path.home(), 'data_from_catalog')
        if os.path.exists(save_dir) == False:
            os.mkdir(save_dir)
        else:
            pass

        save_path = os.path.join(str(Path.home()), save_dir, FILE_NAME)

        # endpoint
        pre_url = 'https://' + BUCKET_NAME + '.' + URL
        self.endpoint = os.path.join(pre_url, FILE_PATH)

        client = boto3.client(BUCKET_TYPE,
                              aws_access_key_id=ACCESS_KEY,
                              aws_secret_access_key=SECRET_KEY,
                              region_name=REGION)

        transfer = boto3.s3.transfer.S3Transfer(client)

        progress = ProgressPercentage(client, BUCKET_NAME, FILE_PATH)

        sys.stdout.write('%s %s %s' % ('Downloading file...', FILE_NAME, '\n'))
        transfer.download_file(BUCKET_NAME, FILE_PATH, save_path, callback=progress)
        logging.info(save_path)

    def set_slack(self, token=None, channel_id=None):
        self.token = token
        self.channel_id = channel_id

    def send_message(self, message=None, thresholds=None):
        if message is not None and thresholds is not None:
            raise ValueError("'message'와 'thresholds' 두 개의 arguments를 동시에 입력할 수 없습니다.")
        else:
            self.message = message
            self.thresholds = thresholds

    def autolog(self, tag=None):
        global description, endpoint, var_model_file_path, message, thresholds, token, channel_id
        description = tag
        endpoint = self.endpoint
        message = self.message
        thresholds = self.thresholds
        token = self.token
        channel_id = self.channel_id

        is_notebook = is_in_ipython()
        if is_notebook == True:
            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={}):
                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()

                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]

                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

                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_git_meta(fileinfo=mainfile)
                set_python_dependencies(py_depenpency=dependencies)

            '''[get best model] on_epoch_end '''

            def on_epoch_end(self, epoch, logs=None):
                logs = logs or {}
                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):
                            if self.verbose > 0:
                                print('\n\nEpoch %05d: %s improved from %0.5f to %0.5f,'
                                      ' storing weights.\n'
                                      % (epoch + 1, self.monitor, self.best, current))
                            self.best = current
                            self.best_epochs = epoch + 1
                            self.best_weights = self.model.get_weights()
                        else:
                            if self.verbose > 0:
                                print('\n\nEpoch %05d: %s did not improve\n' % (epoch + 1, self.monitor))

                # model save path
                common_path = self.dict_path['save_model_path'] + self.model_summary['model_type'] + '-' + \
                              self.model_summary['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'

                self.current_value = current

                # 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 self.verbose > 0:
                    print('\n\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

                end = get_time.now()
                self.model_summary['time_delta'] = str(end - start)
                self.model_summary['selected_metrics'] = {self.monitor: np.float64(self.best)}

                # 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']
                set_best_model_json_path(path_for_setting_model_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_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)
                modeler_rest = LifecycleRestApi(LcConst.BACK_END_API_URL,
                                                LcConst.BACK_END_API_PORT,
                                                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):
                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'] = self.y_val.tolist()
                    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")

        def run_and_log_function(self, original, args, kwargs, unlogged_params):
            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=';')

            ''' save json for visualization '''
            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]
                kwargs['callbacks'] += [TrainHistoryCallbacks()]

            else:
                kwargs['callbacks'] = [csv_logger]
                kwargs['callbacks'] += [get_visual]
                kwargs['callbacks'] += [TrainHistoryCallbacks()]

            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 run_and_log_function(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))

    def _send_message(metric=None, current_value=None, message=None, thresholds=None, token=None, channel_id=None):
        if token is not None and channel_id is not None:
            if thresholds is not None:
                try:
                    if current_value >= thresholds:
                        msg = '[모델 학습 완료] ' + metric + '이 설정하신 thresolds: ' + str(thresholds) + '를 초과하였습니다.'
                        response = WebClient(token=token).chat_postMessage(channel=channel_id, text=msg)
                except SlackApiError as e:
                    assert e.response["error"]
            elif message is not None:
                try:
                    msg = message
                    response = WebClient(token=token).chat_postMessage(channel=channel_id, text=msg)
                except SlackApiError as e:
                    assert e.response["error"]
        else:
            pass

    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))
