# Copyright 2020 The FedLearner Authors. All Rights Reserved.
#
# 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.

# coding: utf-8
# pylint: disable=protected-access

import os
import logging
import time
import tensorflow.compat.v1 as tf

from tensorflow.compat import as_str_any
from tensorflow.compat.v1.train import Optimizer
from tensorflow.compat.v1.estimator import ModeKeys
from tensorflow_estimator.python.estimator import model_fn as model_fn_lib

from fedlearner.common.mysql_client import DBClient
from fedlearner.common.summary_hook import SummaryHook
from fedlearner.trainer import patch  # pylint: disable=unused-import
from fedlearner.common import metrics
from fedlearner.data_join.common import get_kvstore_config

SYNC_PATH = '/sync/'
DATA_CHECKPOINT_INIT_VALUE = "_init_value"

class DataCheckpointSaverListener(tf.estimator.CheckpointSaverListener):
    def __init__(self, tm, appid):
        self._trainer_master = tm
        self._application_id = appid

    def begin(self):
        ckpt = tf.placeholder(tf.string, name="data_checkpoint_plhd")
        var_tmp = tf.Variable(DATA_CHECKPOINT_INIT_VALUE, \
                              name="data_checkpoint")
        self._ckpt_tensor = var_tmp.assign(ckpt)

    def before_save(self, session, global_step_value):
        logging.info('About to write a checkpoint at step %d', \
                global_step_value)
        data_checkpoint = self._trainer_master.get_data_block_checkpoint(
            self._application_id)
        #if empty block from checkpoint fetched due to exception or
        # master not ready, no need to save.
        if len(data_checkpoint) == 0:
            return
        res = session.run(self._ckpt_tensor, {"data_checkpoint_plhd:0":
                                        ",".join(data_checkpoint)})
        logging.info("data checkpoint saved result: %s", res)

class FLModel(object):
    def __init__(self, role, bridge, example_ids, exporting=False):
        self._role = role
        self._bridge = bridge
        self._example_ids = example_ids
        self._exporting = exporting

        self._train_ops = []
        self._recvs = []
        self._sends = []
        self._outputs = []

    @property
    def train_ops(self):
        return self._train_ops

    @property
    def sends(self):
        return [(n, t) for n, t, _ in self._sends]

    @property
    def recvs(self):
        return [(n, t) for n, t, _ in self._recvs]

    def verify_example_ids(self):
        tensor = tf.strings.to_hash_bucket_fast(self._example_ids, 2**31 - 1)
        if self._role == 'leader':
            self.send('_verify_example_ids', tensor)
        else:
            recv_tensor = self.recv('_verify_example_ids', tensor.dtype)
            op = tf.assert_equal(tensor, recv_tensor)
            self._train_ops.append(op)

    def send(self, name, tensor, require_grad=False):
        with tf.control_dependencies([self._example_ids]):
            op = self._bridge.send_op(name, tensor)
        self._train_ops.append(op)
        self._sends.append((name, tensor, require_grad))
        if require_grad:
            return self.recv(name + '_grad', tensor.dtype)
        return None

    def recv(self, name, dtype=tf.float32, require_grad=False):
        with tf.control_dependencies([self._example_ids]):
            tensor = self._bridge.receive_op(name, dtype)
        self._recvs.append((name, tensor, require_grad))
        return tensor

    def minimize(self,
                 optimizer,
                 loss,
                 global_step=None,
                 var_list=None,
                 gate_gradients=Optimizer.GATE_OP,
                 aggregation_method=None,
                 colocate_gradients_with_ops=False,
                 name=None,
                 grad_loss=None):
        recv_grads = [i for i in self._recvs if i[2]]

        if var_list is None:
            var_list = \
                tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) + \
                tf.get_collection(tf.GraphKeys.TRAINABLE_RESOURCE_VARIABLES)
        var_list = [v for _, v, _ in recv_grads] + var_list

        grads_and_vars = optimizer.compute_gradients(
            loss,
            var_list=var_list,
            gate_gradients=gate_gradients,
            aggregation_method=aggregation_method,
            colocate_gradients_with_ops=colocate_gradients_with_ops,
            grad_loss=grad_loss)

        send_grads = grads_and_vars[:len(recv_grads)]
        for (n, _, _), (grad, _) in zip(recv_grads, send_grads):
            if grad is not None:
                self.send(n + '_grad', grad)

        if grads_and_vars[len(recv_grads):]:
            train_op = optimizer.apply_gradients(
                grads_and_vars[len(recv_grads):],
                global_step=global_step,
                name=name)
        else:
            train_op = tf.no_op()

        return train_op

    def _append_summary_hook(self, training_hooks):
        if not training_hooks:
            training_hooks = []
        summary_hook = SummaryHook.get_hook()
        if summary_hook:
            training_hooks.append(summary_hook)
        return training_hooks

    def make_spec(self,
                  mode,
                  predictions=None,
                  loss=None,
                  train_op=None,
                  eval_metric_ops=None,
                  training_chief_hooks=None,
                  training_hooks=None,
                  evaluation_hooks=None,
                  prediction_hooks=None):
        if isinstance(predictions, tf.Tensor):
            predictions = {'output': predictions}
        if mode == ModeKeys.TRAIN:
            train_op = tf.group([train_op] + self._train_ops)
            training_hooks = self._append_summary_hook(training_hooks)
        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions,
            loss=loss,
            train_op=train_op,
            eval_metric_ops=eval_metric_ops,
            training_chief_hooks=training_chief_hooks,
            training_hooks=training_hooks,
            evaluation_hooks=evaluation_hooks,
            prediction_hooks=prediction_hooks)


class FLEstimator(object):
    def __init__(self,
                 model_fn,
                 bridge,
                 trainer_master,
                 role,
                 worker_rank=0,
                 application_id=None,
                 cluster_spec=None):
        self._model_fn = model_fn
        self._bridge = bridge
        self._trainer_master = trainer_master
        self._role = role
        self._worker_rank = worker_rank
        self._cluster_spec = cluster_spec
        self._application_id = application_id

    def _get_features_and_labels_from_input_fn(self, input_fn, mode):
        dataset = input_fn(self._bridge, self._trainer_master)
        features, labels = dataset.make_one_shot_iterator().get_next()
        return features, labels

    def _get_model_spec(self, features, labels, mode):
        model = FLModel(self._role, self._bridge,
                        features.get('example_id', None),
                        exporting=(mode == ModeKeys.PREDICT))
        spec = self._model_fn(model, features, labels, mode)
        return spec, model

    def _restore_datablock(self, blk_ids):
        # only chief worker restores from checkpoint.
        if self._worker_rank != 0 or blk_ids is None:
            return True
        block_id_str = as_str_any(blk_ids)
        block_ids = []
        if block_id_str != DATA_CHECKPOINT_INIT_VALUE:
            block_ids = block_id_str.split(",")
        logging.info("restore: %s", block_id_str)
        return self._trainer_master.restore_data_block_checkpoint(
            self._application_id, block_ids)

    def _cheif_barriar(self, is_chief=False, sync_times=300):
        worker_replicas = os.environ.get('REPLICA_NUM', 0)
        kvstore_type = os.environ.get('KVSTORE_TYPE', 'etcd')
        db_database, db_addr, db_username, db_password, _ = \
            get_kvstore_config(kvstore_type)
        kvstore_client = DBClient(db_database,
                                  db_addr,
                                  db_username,
                                  db_password,
                                  SYNC_PATH)
        sync_path = '%s/%s' % (os.environ['APPLICATION_ID'],
                               os.environ['WORKER_RANK'])
        logging.info('Creating a sync flag at %s', sync_path)
        kvstore_client.set_data(sync_path, "1")
        if is_chief:
            for _ in range(sync_times):
                sync_list = kvstore_client.get_prefix_kvs(
                    os.environ['APPLICATION_ID'])
                logging.info('Sync file pattern is: %s', sync_list)
                if len(sync_list) < worker_replicas:
                    logging.info('Count of ready workers is %d',
                                 len(sync_list))
                    time.sleep(6)
                else:
                    break

    def train(self,
              input_fn,
              checkpoint_path=None,
              save_checkpoint_steps=None,
              save_checkpoint_secs=None):
        if self._cluster_spec is not None:
            device_fn = tf.train.replica_device_setter(
                worker_device="/job:worker/task:%d" % self._worker_rank,
                merge_devices=True,
                cluster=self._cluster_spec)
            cluster_def = self._cluster_spec.as_cluster_def()
            local_address = self._cluster_spec.job_tasks('worker')[
                self._worker_rank]
            server = tf.train.Server(tf.train.ClusterSpec(
                {'local': {
                    0: local_address
                }}),
                job_name='local',
                task_index=0)
            target = 'grpc://' + local_address
        else:
            device_fn = None
            cluster_def = None
            target = None

        config = tf.ConfigProto(cluster_def=cluster_def)
        config.inter_op_parallelism_threads = 4
        config.intra_op_parallelism_threads = 4
        config.experimental.share_session_state_in_clusterspec_propagation \
            = True
        tf.config.set_soft_device_placement(False)

        with tf.Graph().as_default() as g:
            with tf.device(device_fn):
                features, labels = self._get_features_and_labels_from_input_fn(
                    input_fn, ModeKeys.TRAIN)
                spec, _ = self._get_model_spec(features, labels, ModeKeys.TRAIN)

            # Explicitly add a Saver
            if not tf.get_collection(tf.GraphKeys.SAVERS):
                saver = tf.train.Saver(
                    sharded=True,
                    defer_build=True,
                    save_relative_paths=True)  # Must set for portability
                tf.add_to_collection(tf.GraphKeys.SAVERS, saver)

            listener = DataCheckpointSaverListener(self._trainer_master,
                                                   self._application_id)
            saver_hook = tf.estimator.CheckpointSaverHook(
                checkpoint_path, save_secs=save_checkpoint_secs,
                save_steps=save_checkpoint_steps, listeners=[listener])
            self._bridge.connect()

            try:
                with tf.train.MonitoredTrainingSession(
                    master=target,
                    config=config,
                    is_chief=(self._worker_rank == 0),
                    chief_only_hooks=[saver_hook],
                    checkpoint_dir=checkpoint_path,
                    save_checkpoint_steps=save_checkpoint_steps,
                    save_checkpoint_secs=save_checkpoint_secs,
                    hooks=spec.training_hooks) as sess:
                    iter_id = 0

                    data_checkpoint_value = None
                    if hasattr(saver_hook, "data_checkpoint"):
                        data_checkpoint_value = saver_hook.data_checkpoint
                    if not self._restore_datablock(data_checkpoint_value):
                        raise ValueError("Restore data checkpoint error")

                    while not sess.should_stop():
                        self._bridge.start(iter_id)
                        logging.debug('after bridge start.')
                        start_time = time.time()
                        sess.run(spec.train_op, feed_dict={})
                        end_time = time.time()
                        metrics.emit_timer(
                            name="iter_timer",
                            value=end_time-start_time,
                            tags={})
                        logging.debug('after session run.')
                        self._bridge.commit()
                        logging.debug('after bridge commit.')
                        iter_id += 1
            finally:
                self._bridge.terminate()

        return self

    def evaluate(self,
                 input_fn,
                 checkpoint_path=None):
        if not tf.train.latest_checkpoint(checkpoint_path):
            raise ValueError(
                "Could not find trained model at %s" % checkpoint_path)

        with tf.Graph().as_default():
            features, labels = self._get_features_and_labels_from_input_fn(
                input_fn, ModeKeys.EVAL)
            spec, model = self._get_model_spec(features, labels, ModeKeys.EVAL)

            # Track the average loss in default
            eval_metric_ops = spec.eval_metric_ops or {}
            if model_fn_lib.LOSS_METRIC_KEY not in eval_metric_ops:
                loss_metric = tf.metrics.mean(spec.loss)
                eval_metric_ops[model_fn_lib.LOSS_METRIC_KEY] = loss_metric

            # Create the real eval op
            update_ops, eval_dict = _extract_metric_update_ops(eval_metric_ops)
            update_ops.extend(model._train_ops)
            eval_op = tf.group(*update_ops)

            # Also track the global step
            if tf.GraphKeys.GLOBAL_STEP in eval_dict:
                raise ValueError(
                    'Metric with name `global_step` is not allowed, because '
                    'Estimator already defines a default metric with the '
                    'same name.')
            eval_dict[tf.GraphKeys.GLOBAL_STEP] = \
                tf.train.get_or_create_global_step()

            # Prepare the session creator.
            scaffold = tf.train.Scaffold()
            session_creator = tf.train.ChiefSessionCreator(
                scaffold=scaffold,
                checkpoint_dir=checkpoint_path)

            # Prepare hooks
            all_hooks = list(spec.evaluation_hooks) or []
            final_ops_hook = tf.train.FinalOpsHook(eval_dict)
            all_hooks.append(final_ops_hook)

            # Evaluate over dataset
            self._bridge.connect()
            try:
                with tf.train.MonitoredSession(
                    session_creator=session_creator, hooks=all_hooks) as sess:
                    if not self._restore_datablock(DATA_CHECKPOINT_INIT_VALUE):
                        raise ValueError("Restore data checkpoint error")
                    iter_id = 0
                    while not sess.should_stop():
                        self._bridge.start(iter_id)
                        logging.debug('after bridge start.')
                        start_time = time.time()
                        sess.run(eval_op)
                        end_time = time.time()
                        metrics.emit_timer(
                            name="iter_timer",
                            value=end_time-start_time,
                            tags={})
                        logging.debug('after session run.')
                        self._bridge.commit()
                        logging.debug('after bridge commit.')
                        iter_id += 1
            finally:
                self._bridge.terminate()

            # Print result
            logging.info('Metrics for iteration %d: %s',
                iter_id, _dict_to_str(final_ops_hook.final_ops_values))
            return final_ops_hook.final_ops_values

    def export_saved_model(self,
                           export_dir_base,
                           serving_input_receiver_fn,
                           checkpoint_path=None):
        with tf.Graph().as_default():
            receiver = serving_input_receiver_fn()
            spec, model = self._get_model_spec(receiver.features, None,
                                               ModeKeys.PREDICT)
            assert not model.sends, "Exported model cannot send"
            assert not model.recvs, "Exported model cannot receive"

            with tf.Session() as sess:
                saver_for_restore = tf.train.Saver(sharded=True)
                saver_for_restore.restore(
                    sess, tf.train.latest_checkpoint(checkpoint_path))
                tf.saved_model.simple_save(sess, export_dir_base,
                                           receiver.receiver_tensors,
                                           spec.predictions, None)

        return export_dir_base


def _extract_metric_update_ops(eval_dict):
    """Separate update operations from metric value operations."""
    update_ops = []
    value_ops = {}
    # Sort metrics lexicographically so graph is identical every time.
    for name in sorted(eval_dict.keys()):
        metric_tensor, update_op = eval_dict[name]
        value_ops[name] = metric_tensor
        update_ops.append(update_op)
    return update_ops, value_ops


def _dict_to_str(dictionary):
    """Get a `str` representation of a `dict`.

    Args:
        dictionary: The `dict` to be represented as `str`.

    Returns:
        A `str` representing the `dictionary`.
    """
    return ', '.join('%s = %s' % (k, v)
                     for k, v in sorted(dictionary.items())
                     if not isinstance(v, bytes))

tensorflow.compat.v1.train.replica_device_setter tensorflow.compat.v1.Graph tensorflow.compat.v1.train.FinalOpsHook tensorflow.compat.v1.saved_model.simple_save tensorflow.compat.v1.train.ClusterSpec tensorflow.compat.v1.strings.to_hash_bucket_fast tensorflow.compat.v1.train.MonitoredTrainingSession tensorflow.compat.as_str_any tensorflow.compat.v1.train.ChiefSessionCreator tensorflow.compat.v1.device tensorflow.compat.v1.add_to_collection tensorflow.compat.v1.estimator.CheckpointSaverHook tensorflow.compat.v1.placeholder tensorflow.compat.v1.train.MonitoredSession tensorflow.compat.v1.train.get_or_create_global_step tensorflow.compat.v1.metrics.mean tensorflow.compat.v1.estimator.EstimatorSpec tensorflow.compat.v1.train.latest_checkpoint tensorflow.compat.v1.train.Scaffold tensorflow.compat.v1.train.Saver tensorflow.compat.v1.ConfigProto tensorflow.compat.v1.Session tensorflow.compat.v1.control_dependencies tensorflow.compat.v1.get_collection tensorflow.compat.v1.group tensorflow.compat.v1.assert_equal tensorflow.compat.v1.no_op tensorflow.compat.v1.config.set_soft_device_placement tensorflow.compat.v1.Variable