# Copyright 2019 The TensorFlow 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.
# ==============================================================================
"""Interface to run mask rcnn model in different distributed strategies."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import json
import os
from absl import logging
import numpy as np
import tensorflow.compat.v1 as tf
from tensorflow.compat.v1 import estimator as tf_estimator

from hyperparameters import params_dict
import evaluation


class DistributedExecuter(object):
  """Interface to run Mask RCNN model in TPUs/GPUs.

  Attributes:
    flags: FLAGS object passed from the user.
    model_params: Model configuration needed to run distribution strategy.
    model_fn: Model function to be passed to Estimator.
  """

  __metaclass__ = abc.ABCMeta

  def __init__(self, flags, model_params, model_fn):
    self._flags = flags
    self._model_params = model_params
    self._model_fn = model_fn
    self._tpu_job_name = flags.tpu_job_name

  @abc.abstractmethod
  def build_strategy_configuration(self):
    """Builds run configuration for distributed train/eval.

    Returns:
      RunConfig with distribution strategy configurations
      to pass to the constructor of TPUEstimator/Estimator.
    """

    NotImplementedError('Must be implmented in subclass')

  @abc.abstractmethod
  def build_model_parameters(self, unused_mode, unused_run_config):
    """Builds model parameters.

    Returns:
      A dictionary to pass into `model_fn` of Estimator/TPUEstimator.
    """

    NotImplementedError('Must be implmented in subclass')

  @abc.abstractmethod
  def build_mask_rcnn_estimator(self, params, run_config, mode):
    """Creates TPUEstimator/Estimator instance.

    Arguments:
      params: A dictionary to pass to Estimator `model_fn`.
      run_config: RunConfig instance specifying distribution strategy
        configurations.
      mode: Mode -- one of 'train` or `eval`.

    Returns:
      TFEstimator or TPUEstimator instance.
    """

    NotImplementedError('Must be implmented in subclass')

  def _save_config(self):
    """Save parameters to config files if model_dir is defined."""

    model_dir = self._flags.model_dir
    if model_dir is not None:
      if not tf.gfile.Exists(model_dir):
        tf.gfile.MakeDirs(model_dir)
      params_dict.save_params_dict_to_yaml(self._model_params,
                                           model_dir + '/params.yaml')

  def _write_summary(self, summary_writer, eval_results, predictions,
                     current_step):
    if not self._model_params.visualize_images_summary:
      predictions = None
    evaluation.write_summary(
        eval_results, summary_writer, current_step, predictions=predictions)

  def train(self,
            train_input_fn,
            run_eval_after_train=False,
            eval_input_fn=None):
    """Run distributed training on Mask RCNN model."""

    self._save_config()
    run_config = self.build_strategy_configuration()
    params = self.build_model_parameters('train', run_config)
    logging.info(params)
    train_estimator = self.build_mask_rcnn_estimator(params, run_config,
                                                     'train')
    if self._model_params.use_tpu:
      train_estimator.train(
          input_fn=train_input_fn, max_steps=self._model_params.total_steps)
    else:
      # As MirroredStrategy only supports `train_and_evaluate`, for training,
      # we pass dummy `eval_spec`.
      train_spec = tf_estimator.TrainSpec(
          input_fn=train_input_fn, max_steps=self._model_params.total_steps)
      eval_spec = tf_estimator.EvalSpec(input_fn=tf.data.Dataset)
      tf_estimator.train_and_evaluate(train_estimator, train_spec, eval_spec)

    eval_results = None
    if not run_eval_after_train:
      return eval_results

    if eval_input_fn is None:
      raise ValueError('Eval input_fn must be passed to conduct '
                       'evaluation after training.')

    eval_params = self.build_model_parameters('eval', run_config)
    eval_estimator = self.build_mask_rcnn_estimator(eval_params, run_config,
                                                    'eval')
    eval_results, predictions = evaluation.evaluate(
        eval_estimator, eval_input_fn, self._model_params.eval_samples,
        self._model_params.eval_batch_size, self._model_params.include_mask,
        self._model_params.val_json_file)

    output_dir = os.path.join(self._flags.model_dir, 'eval')
    tf.gfile.MakeDirs(output_dir)
    # Summary writer writes out eval metrics.
    summary_writer = tf.summary.FileWriter(output_dir)
    self._write_summary(summary_writer, eval_results, predictions,
                        self._model_params.total_steps)
    summary_writer.close()

    return eval_results

  def eval(self, eval_input_fn):
    """Run distributed eval on Mask RCNN model."""

    output_dir = os.path.join(self._flags.model_dir, 'eval')
    tf.gfile.MakeDirs(output_dir)

    # Summary writer writes out eval metrics.
    summary_writer = tf.summary.FileWriter(output_dir)
    run_config = self.build_strategy_configuration()
    eval_params = self.build_model_parameters('eval', run_config)
    eval_estimator = self.build_mask_rcnn_estimator(eval_params, run_config,
                                                    'eval')

    def _terminate_eval():
      logging.info('Terminating eval after %d seconds of '
                   'no checkpoints', self._flags.eval_timeout)
      return True

    eval_results = None
    # Run evaluation when there's a new checkpoint
    for ckpt in tf.train.checkpoints_iterator(
        self._flags.model_dir,
        min_interval_secs=self._flags.min_eval_interval,
        timeout=self._flags.eval_timeout,
        timeout_fn=_terminate_eval):
      # Terminate eval job when final checkpoint is reached
      current_step = int(os.path.basename(ckpt).split('-')[1])

      logging.info('Starting to evaluate.')
      try:
        eval_results, predictions = evaluation.evaluate(
            eval_estimator, eval_input_fn, self._model_params.eval_samples,
            self._model_params.eval_batch_size, self._model_params.include_mask,
            self._model_params.val_json_file)
        self._write_summary(summary_writer, eval_results, predictions,
                            current_step)

        if current_step >= self._model_params.total_steps:
          logging.info('Evaluation finished after training step %d',
                       current_step)
          break
      except tf.errors.NotFoundError:
        # Since the coordinator is on a different job than the TPU worker,
        # sometimes the TPU worker does not finish initializing until long after
        # the CPU job tells it to start evaluating. In this case, the checkpoint
        # file could have been deleted already.
        logging.info('Checkpoint %s no longer exists, skipping checkpoint',
                     ckpt)
    summary_writer.close()
    return eval_results

  def train_and_eval(self, train_input_fn, eval_input_fn):
    """Run distributed train and eval on Mask RCNN model."""

    self._save_config()
    output_dir = os.path.join(self._flags.model_dir, 'eval')
    tf.gfile.MakeDirs(output_dir)
    summary_writer = tf.summary.FileWriter(output_dir)

    run_config = self.build_strategy_configuration()
    train_params = self.build_model_parameters('train', run_config)
    eval_params = self.build_model_parameters('eval', run_config)
    train_estimator = self.build_mask_rcnn_estimator(train_params, run_config,
                                                     'train')
    eval_estimator = self.build_mask_rcnn_estimator(eval_params, run_config,
                                                    'eval')

    num_cycles = int(self._model_params.total_steps /
                     self._model_params.num_steps_per_eval)
    for cycle in range(num_cycles):
      logging.info('Start training cycle %d.', cycle)
      train_estimator.train(
          input_fn=train_input_fn, steps=self._model_params.num_steps_per_eval)

      logging.info('Start evaluation cycle %d.', cycle)
      eval_results, predictions = evaluation.evaluate(
          eval_estimator, eval_input_fn, self._model_params.eval_samples,
          self._model_params.eval_batch_size, self._model_params.include_mask,
          self._model_params.val_json_file)

      current_step = int(cycle * self._model_params.num_steps_per_eval)
      self._write_summary(summary_writer, eval_results, predictions,
                          current_step)

    logging.info('Starting training cycle %d.', num_cycles)
    train_estimator.train(
        input_fn=train_input_fn, max_steps=self._model_params.total_steps)
    eval_results, predictions = evaluation.evaluate(
        eval_estimator, eval_input_fn, self._model_params.eval_samples,
        self._model_params.eval_batch_size, self._model_params.include_mask,
        self._model_params.val_json_file)
    self._write_summary(summary_writer, eval_results, predictions,
                        self._model_params.total_steps)
    summary_writer.close()
    return eval_results


class TPUEstimatorExecuter(DistributedExecuter):
  """Interface that runs Mask RCNN model using TPUEstimator."""

  def build_strategy_configuration(self):
    """Retrieves model configuration for running tpu estimator."""

    if self._model_params.use_tpu:
      tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
          self._flags.tpu,
          zone=self._flags.tpu_zone,
          project=self._flags.gcp_project)
      tpu_grpc_url = tpu_cluster_resolver.get_master()
      tf.Session.reset(tpu_grpc_url)
    else:
      tpu_cluster_resolver = None

    num_cores = self._model_params.num_cores
    input_partition_dims = self._flags.input_partition_dims

    # The following is for spatial partitioning. `features` has one tensor while
    # `labels` has 4 + (`max_level` - `min_level` + 1) * 2 tensors. The input
    # partition is performed on `features` and all partitionable tensors of
    # `labels`, see the partition logic below.
    # Note: In the below code, TPUEstimator uses both `shard` and `replica`
    # (with the same meaning).
    if input_partition_dims:
      labels_partition_dims = {
          'gt_boxes': None,
          'gt_classes': None,
          'cropped_gt_masks': None,
      }
      # TODO(b/119617317): The Input Partition Logic. We partition only the
      # partition-able tensors. Spatial partition requires that the
      # to-be-partitioned tensors must have a dimension that is a multiple of
      # `partition_dims`. Depending on the `partition_dims` and the `image_size`
      # and the `max_level` in config, some high-level anchor labels (i.e.,
      # `cls_targets` and `box_targets`) cannot be partitioned. For example,
      # when `partition_dims` is [1, 4, 2, 1], image size is 1536, `max_level`
      # is 9, `cls_targets_8` has a shape of [batch_size, 6, 6, 9], which
      # cannot be partitioned (6 % 4 != 0). In this case, the level-8 and
      # level-9 target tensors are not partition-able, and the highest
      # partition-able level is 7.
      image_size = self._model_params.image_size
      for level in range(self._model_params.min_level,
                         self._model_params.max_level + 1):

        def _can_partition(spatial_dim):
          partitionable_index = np.where(spatial_dim %
                                         np.array(input_partition_dims) == 0)
          return len(partitionable_index[0]) == len(input_partition_dims)

        assert len(image_size) == 2
        spatial_dim = [d // (2**level) for d in image_size]
        if _can_partition(spatial_dim[0]) and _can_partition(spatial_dim[1]):
          labels_partition_dims['box_targets_%d' % level] = input_partition_dims
          labels_partition_dims['score_targets_%d' %
                                level] = input_partition_dims
        else:
          labels_partition_dims['box_targets_%d' % level] = None
          labels_partition_dims['score_targets_%d' % level] = None

      num_cores_per_replica = np.prod(input_partition_dims)
      transpose_input = self._model_params.transpose_input
      image_partition_dims = [input_partition_dims[i] for i in [1, 2, 3, 0]
                             ] if transpose_input else input_partition_dims

      features_partition_dims = {
          'images': image_partition_dims,
          'source_ids': None,
          'image_info': None,
      }
      input_partition_dims = [features_partition_dims, labels_partition_dims]
      num_shards = num_cores // num_cores_per_replica
    else:
      num_cores_per_replica = None
      input_partition_dims = None
      num_shards = num_cores

    tpu_config = tf_estimator.tpu.TPUConfig(
        self._model_params.iterations_per_loop,
        num_shards=num_shards,
        num_cores_per_replica=num_cores_per_replica,
        input_partition_dims=input_partition_dims,
        tpu_job_name=self._tpu_job_name,
        per_host_input_for_training=tf_estimator.tpu.InputPipelineConfig
        .PER_HOST_V2)
    run_config = tf_estimator.tpu.RunConfig(
        cluster=tpu_cluster_resolver,
        evaluation_master=self._flags.eval_master,
        model_dir=self._flags.model_dir,
        log_step_count_steps=self._model_params.iterations_per_loop,
        tpu_config=tpu_config,
    )
    return run_config

  def build_model_parameters(self, mode, run_config):
    assert mode in ('train', 'eval')
    params = dict(
        self._model_params.as_dict().items(),
        mode=self._flags.mode,
        model_dir=self._flags.model_dir,
        transpose_input=self._model_params.transpose_input,
        num_shards=run_config.tpu_config.num_shards,
        use_tpu=self._model_params.use_tpu,
        # Used by the host_call function.
        iterations_per_loop=self._model_params.iterations_per_loop)

    if mode == 'eval':
      params = dict(
          params,
          input_rand_hflip=False,
          is_training_bn=False,
          transpose_input=False)
    return params

  def build_mask_rcnn_estimator(self, params, run_config, unused_mode):
    estimator = tf_estimator.tpu.TPUEstimator(
        model_fn=self._model_fn,
        use_tpu=params['use_tpu'],
        train_batch_size=self._model_params.train_batch_size,
        eval_batch_size=self._model_params.eval_batch_size,
        predict_batch_size=self._model_params.eval_batch_size,
        config=run_config,
        params=params)
    return estimator


class MultiWorkerExecuter(DistributedExecuter):
  """Interface that runs Mask RCNN model using MultiWorkerMirroredStrategy."""

  @staticmethod
  def is_eval_task():
    return tf.distribute.cluster_resolver.TFConfigClusterResolver(
    ).task_type == 'evaluator'

  def build_strategy_configuration(self):
    """Retrieves model configuration for MultiWorkerMirroredStrategy."""

    worker_hosts = self._flags.worker_hosts

    if worker_hosts is not None:
      # Set TF_CONFIG environment variable
      worker_hosts = worker_hosts.split(',')
      task_index = self._flags.task_index
      os.environ['TF_CONFIG'] = json.dumps({
          'cluster': {
              'worker': worker_hosts
          },
          'task': {
              'type': 'worker',
              'index': task_index
          }
      })

    multiworker_strategy = (
        tf.distribute.experimental.MultiWorkerMirroredStrategy())
    run_config = tf_estimator.RunConfig(
        train_distribute=multiworker_strategy, model_dir=self._flags.model_dir)
    return run_config

  def build_model_parameters(self, mode, unused_config):
    """Builds model parameter to run in MultiWorkerMirroredStrategy."""

    assert mode in ('train', 'eval')
    batch_size = (
        self._model_params.train_batch_size
        if mode == 'train' else self._model_params.eval_batch_size)
    params = dict(
        self._model_params.as_dict().items(),
        use_tpu=False,
        mode=mode,
        model_dir=self._flags.model_dir,
        # For MultiWorkerMirroredStrategy, we use CPU for evaluation and
        # CPU only supports channel-last data format. As so, we do not
        # transpose input by default to make data format consistent.
        transpose_input=False,
        batch_size=batch_size,
        precision='float32')
    return params

  def build_mask_rcnn_estimator(self, params, run_config, mode):
    """Returns Mask Rcnn model running on MultiWorkerMirroredStrategy."""
    assert mode in ('train', 'eval')
    if mode == 'train':
      return tf_estimator.Estimator(
          model_fn=self._model_fn,
          model_dir=self._flags.model_dir,
          config=run_config,
          params=params)

    # Evaluation on multi-worker mirrored strategy is done in CPU for now
    # as only `train_and_evaluate` is supported and eval pipeline for
    # Mask RCNN model is uses `predict` API.
    cpu_run_config = tf_estimator.RunConfig(model_dir=self._flags.model_dir)
    return tf_estimator.Estimator(
        model_fn=self._model_fn,
        model_dir=self._flags.model_dir,
        config=cpu_run_config,
        params=params)

  def train_and_eval(self, train_input_fn, eval_input_fn):
    if self.is_eval_task():
      assert eval_input_fn is not None
      self.eval(eval_input_fn)
    else:
      self.train(train_input_fn)
