# Copyright 2020 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.
# ==============================================================================
r"""Utils for EfficientNet models for Keras.

Write weights from  ckpt file as in original repo
(https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet)
to h5 file for keras implementation of the models.

Usage:

# use checkpoint efficientnet-b0/model.ckpt (can be downloaded from
# https://storage.googleapis.com/cloud-tpu-checkpoints/
#     efficientnet/ckptsaug/efficientnet-b0.tar.gz)
# to update weight without top layers, saving to efficientnetb0_notop.h5
python efficientnet_weight_update_util.py --model b0 --notop \
    --ckpt efficientnet-b0/model.ckpt --o efficientnetb0_notop.h5

# use checkpoint noisy_student_efficientnet-b3/model.ckpt (providing
# improved result for b3, can be downloaded from
# https://storage.googleapis.com/cloud-tpu-checkpoints/
#     efficientnet/noisystudent/noisy_student_efficientnet-b3.tar.gz)
# to update weight with top layers, saving to efficientnetb3_new.h5
python efficientnet_weight_update_util.py --model b3 --notop \
    --ckpt noisy_student_efficientnet-b3/model.ckpt --o efficientnetb3_new.h5
"""

import tensorflow.compat.v2 as tf

import argparse
import warnings
from tensorflow.keras.applications import efficientnet


def write_ckpt_to_h5(path_h5, path_ckpt, keras_model, use_ema=True):
  """Map the weights in checkpoint file (tf) to h5 file (keras).

  Args:
    path_h5: str, path to output hdf5 file to write weights loaded from ckpt
      files.
    path_ckpt: str, path to the ckpt files (e.g. 'efficientnet-b0/model.ckpt')
      that records efficientnet weights from original repo
      https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet
    keras_model: keras model, built from keras.applications efficientnet
      functions (e.g. EfficientNetB0)
    use_ema: Bool, whether to use ExponentialMovingAverage result or not
  """
  model_name_keras = keras_model.name
  model_name_tf = model_name_keras.replace('efficientnet', 'efficientnet-')

  keras_weight_names = [w.name for w in keras_model.weights]
  tf_weight_names = get_variable_names_from_ckpt(path_ckpt)

  keras_blocks = get_keras_blocks(keras_weight_names)
  tf_blocks = get_tf_blocks(tf_weight_names)

  print('check variables match in each block')
  for keras_block, tf_block in zip(keras_blocks, tf_blocks):
    check_match(keras_block, tf_block, keras_weight_names, tf_weight_names,
                model_name_tf)
    print('{} and {} match.'.format(tf_block, keras_block))

  block_mapping = {x[0]: x[1] for x in zip(keras_blocks, tf_blocks)}

  changed_weights = 0
  for w in keras_model.weights:
    if 'block' in w.name:
      # example: 'block1a_dwconv/depthwise_kernel:0' -> 'block1a'
      keras_block = w.name.split('/')[0].split('_')[0]
      tf_block = block_mapping[keras_block]
      tf_name = keras_name_to_tf_name_block(
          w.name,
          keras_block=keras_block,
          tf_block=tf_block,
          use_ema=use_ema,
          model_name_tf=model_name_tf)
    elif any([x in w.name for x in ['stem', 'top', 'predictions', 'probs']]):
      tf_name = keras_name_to_tf_name_stem_top(
          w.name, use_ema=use_ema, model_name_tf=model_name_tf)
    elif 'normalization' in w.name:
      print('skipping variable {}: normalization is a layer'
            'in keras implementation, but preprocessing in '
            'TF implementation.'.format(w.name))
      continue
    else:
      raise ValueError('{} failed to parse.'.format(w.name))

    try:
      w_tf = tf.train.load_variable(path_ckpt, tf_name)
      if (w.value().numpy() != w_tf).any():
        w.assign(w_tf)
        changed_weights += 1
    except ValueError as e:
      if any([x in w.name for x in ['top', 'predictions', 'probs']]):
        warnings.warn('Fail to load top layer variable {}'
                      'from {} because of {}.'.format(w.name, tf_name, e))
      else:
        raise ValueError('Fail to load {} from {}'.format(w.name, tf_name))

  total_weights = len(keras_model.weights)
  print('{}/{} weights updated'.format(changed_weights, total_weights))
  keras_model.save_weights(path_h5)


def get_variable_names_from_ckpt(path_ckpt, use_ema=True):
  """Get list of tensor names from checkpoint.

  Args:
    path_ckpt: str, path to the ckpt files
    use_ema: Bool, whether to use ExponentialMovingAverage result or not.
  Returns:
    List of variable names from checkpoint.
  """
  v_all = tf.train.list_variables(path_ckpt)

  # keep name only
  v_name_all = [x[0] for x in v_all]

  if use_ema:
    v_name_all = [x for x in v_name_all if 'ExponentialMovingAverage' in x]
  else:
    v_name_all = [x for x in v_name_all if 'ExponentialMovingAverage' not in x]

  # remove util variables used for RMSprop
  v_name_all = [x for x in v_name_all if 'RMS' not in x]
  return v_name_all


def get_tf_blocks(tf_weight_names):
  """Extract the block names from list of full weight names."""
  # Example: 'efficientnet-b0/blocks_0/conv2d/kernel' -> 'blocks_0'
  tf_blocks = {x.split('/')[1] for x in tf_weight_names if 'block' in x}
  # sort by number
  tf_blocks = sorted(tf_blocks, key=lambda x: int(x.split('_')[1]))
  return tf_blocks


def get_keras_blocks(keras_weight_names):
  """Extract the block names from list of full weight names."""
  # example: 'block1a_dwconv/depthwise_kernel:0' -> 'block1a'
  keras_blocks = {x.split('_')[0] for x in keras_weight_names if 'block' in x}
  return sorted(keras_blocks)


def keras_name_to_tf_name_stem_top(keras_name,
                                   use_ema=True,
                                   model_name_tf='efficientnet-b0'):
  """Mapping name in h5 to ckpt that is in stem or top (head).

  we map name keras_name that points to a weight in h5 file
  to a name of weight in ckpt file.

  Args:
    keras_name: str, the name of weight in the h5 file of keras implementation
    use_ema: Bool, use the ExponentialMovingAverage resuolt in ckpt or not
    model_name_tf: str, the name of model in ckpt.

  Returns:
    String for the name of weight as in ckpt file.

  Raises:
    KeyError: if we cannot parse the keras_name.
  """
  if use_ema:
    ema = '/ExponentialMovingAverage'
  else:
    ema = ''

  stem_top_dict = {
      'probs/bias:0': '{}/head/dense/bias{}',
      'probs/kernel:0': '{}/head/dense/kernel{}',
      'predictions/bias:0': '{}/head/dense/bias{}',
      'predictions/kernel:0': '{}/head/dense/kernel{}',
      'stem_conv/kernel:0': '{}/stem/conv2d/kernel{}',
      'top_conv/kernel:0': '{}/head/conv2d/kernel{}',
  }
  for x in stem_top_dict:
    stem_top_dict[x] = stem_top_dict[x].format(model_name_tf, ema)

  # stem batch normalization
  for bn_weights in ['beta', 'gamma', 'moving_mean', 'moving_variance']:
    tf_name = '{}/stem/tpu_batch_normalization/{}{}'.format(
        model_name_tf, bn_weights, ema)
    stem_top_dict['stem_bn/{}:0'.format(bn_weights)] = tf_name

  # top / head batch normalization
  for bn_weights in ['beta', 'gamma', 'moving_mean', 'moving_variance']:
    tf_name = '{}/head/tpu_batch_normalization/{}{}'.format(
        model_name_tf, bn_weights, ema)
    stem_top_dict['top_bn/{}:0'.format(bn_weights)] = tf_name

  if keras_name in stem_top_dict:
    return stem_top_dict[keras_name]
  raise KeyError('{} from h5 file cannot be parsed'.format(keras_name))


def keras_name_to_tf_name_block(keras_name,
                                keras_block='block1a',
                                tf_block='blocks_0',
                                use_ema=True,
                                model_name_tf='efficientnet-b0'):
  """Mapping name in h5 to ckpt that belongs to a block.

  we map name keras_name that points to a weight in h5 file
  to a name of weight in ckpt file.

  Args:
    keras_name: str, the name of weight in the h5 file of keras implementation
    keras_block: str, the block name for keras implementation (e.g. 'block1a')
    tf_block: str, the block name for tf implementation (e.g. 'blocks_0')
    use_ema: Bool, use the ExponentialMovingAverage resuolt in ckpt or not
    model_name_tf: str, the name of model in ckpt.

  Returns:
    String for the name of weight as in ckpt file.

  Raises:
    ValueError if keras_block does not show up in keras_name
  """

  if keras_block not in keras_name:
    raise ValueError('block name {} not found in {}'.format(
        keras_block, keras_name))

  # all blocks in the first group will not have expand conv and bn
  is_first_blocks = (keras_block[5] == '1')

  tf_name = [model_name_tf, tf_block]

  # depthwide conv
  if 'dwconv' in keras_name:
    tf_name.append('depthwise_conv2d')
    tf_name.append('depthwise_kernel')

  # conv layers
  if is_first_blocks:
    # first blocks only have one conv2d
    if 'project_conv' in keras_name:
      tf_name.append('conv2d')
      tf_name.append('kernel')
  else:
    if 'project_conv' in keras_name:
      tf_name.append('conv2d_1')
      tf_name.append('kernel')
    elif 'expand_conv' in keras_name:
      tf_name.append('conv2d')
      tf_name.append('kernel')

  # squeeze expansion layers
  if '_se_' in keras_name:
    if 'reduce' in keras_name:
      tf_name.append('se/conv2d')
    elif 'expand' in keras_name:
      tf_name.append('se/conv2d_1')

    if 'kernel' in keras_name:
      tf_name.append('kernel')
    elif 'bias' in keras_name:
      tf_name.append('bias')

  # batch normalization layers
  if 'bn' in keras_name:
    if is_first_blocks:
      if 'project' in keras_name:
        tf_name.append('tpu_batch_normalization_1')
      else:
        tf_name.append('tpu_batch_normalization')
    else:
      if 'project' in keras_name:
        tf_name.append('tpu_batch_normalization_2')
      elif 'expand' in keras_name:
        tf_name.append('tpu_batch_normalization')
      else:
        tf_name.append('tpu_batch_normalization_1')

    for x in ['moving_mean', 'moving_variance', 'beta', 'gamma']:
      if x in keras_name:
        tf_name.append(x)
  if use_ema:
    tf_name.append('ExponentialMovingAverage')
  return '/'.join(tf_name)


def check_match(keras_block, tf_block, keras_weight_names, tf_weight_names,
                model_name_tf):
  """Check if the weights in h5 and ckpt match.

  we match each name from keras_weight_names that is in keras_block
  and check if there is 1-1 correspondence to names from tf_weight_names
  that is in tf_block

  Args:
    keras_block: str, the block name for keras implementation (e.g. 'block1a')
    tf_block: str, the block name for tf implementation (e.g. 'blocks_0')
    keras_weight_names: list of str, weight names in keras implementation
    tf_weight_names: list of str, weight names in tf implementation
    model_name_tf: str, the name of model in ckpt.
  """
  names_from_keras = set()
  for x in keras_weight_names:
    if keras_block in x:
      y = keras_name_to_tf_name_block(
          x,
          keras_block=keras_block,
          tf_block=tf_block,
          model_name_tf=model_name_tf)
      names_from_keras.add(y)

  names_from_tf = set()
  for x in tf_weight_names:
    if tf_block in x and x.split('/')[1].endswith(tf_block):
      names_from_tf.add(x)

  names_missing = names_from_keras - names_from_tf
  if names_missing:
    raise ValueError('{} variables not found in checkpoint file: {}'.format(
        len(names_missing), names_missing))

  names_unused = names_from_tf - names_from_keras
  if names_unused:
    warnings.warn('{} variables from checkpoint file are not used: {}'.format(
        len(names_unused), names_unused))


if __name__ == '__main__':
  arg_to_model = {
      'b0': efficientnet.EfficientNetB0,
      'b1': efficientnet.EfficientNetB1,
      'b2': efficientnet.EfficientNetB2,
      'b3': efficientnet.EfficientNetB3,
      'b4': efficientnet.EfficientNetB4,
      'b5': efficientnet.EfficientNetB5,
      'b6': efficientnet.EfficientNetB6,
      'b7': efficientnet.EfficientNetB7
  }

  p = argparse.ArgumentParser(description='write weights from checkpoint to h5')
  p.add_argument(
      '--model',
      required=True,
      type=str,
      help='name of efficient model',
      choices=arg_to_model.keys())
  p.add_argument(
      '--notop',
      action='store_true',
      help='do not include top layers',
      default=False)
  p.add_argument('--ckpt', required=True, type=str, help='checkpoint path')
  p.add_argument(
      '--output', '-o', required=True, type=str, help='output (h5) file path')
  args = p.parse_args()

  include_top = not args.notop

  model = arg_to_model[args.model](include_top=include_top)
  write_ckpt_to_h5(args.output, args.ckpt, keras_model=model)
