# coding=utf-8
# Copyright 2022 The Deeplab2 Authors.
#
# 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.

"""MOAT: This file contains the implementation of MOAT [1].

[1] MOAT: Alternating Mobile Convolution and Attention
    Brings Strong Vision Models,
    arXiv: 2210.01820.
      Chenglin Yang, Siyuan Qiao, Qihang Yu, Xiaoding Yuan,
      Yukun Zhu, Alan Yuille, Hartwig Adam, Liang-Chieh Chen.
"""

import copy
import re
from typing import Optional

import tensorflow as tf
from absl import logging
from deeplab2.model.layers.moat_blocks import MBConvBlock
from deeplab2.model.layers.moat_blocks import MOATBlock
from deeplab2.utils.hparam_configs import Config
from deeplab2.utils.hparam_configs import create_config_from_dict

# This handles the invalid name scope of tf.keras.sequential
# used in stem layers.
_STEM_LAYER_NAME_SCOPE = 'moat/stem/'

# This is for loading the exponential moving average of variables
# in the checkpoint.
_EMA_VARIABLE_NAME_POSTFIX = '/ExponentialMovingAverage'

# The position embedding size at stride 16 and 32.
# The input size changes, but the number of learnable parameters of position
# embedding does not change. The position embeddings are interpolated for
# different input sizes.
_STRIDE_16_POSITION_EMBEDDING_SIZE = 14
_STRIDE_32_POSITION_EMBEDDING_SIZE = 7


class MOAT(tf.keras.Model):
    """MOAT backbone."""

    def _retrieve_config(self, config):
        """Retrieves the config of MOAT.

        Args:
          config: A dictionary containing the following keys.
            -stem_size: A list of integers, the list length is the number of stem
              layers, the list element is the output channels of the stem layer.
            -block_type: A list of strings, the list length is the stage number.
              The list element is the block type for a stage. The supported block
              types are 'mbconv' and 'moat'.
            -num_blocks: A list of integers, the list length is the stage number.
              The list element is the number of blocks in a stage.
            -hidden_size: A list of integers, the list length is the stage number.
              The list element is the output channels of the blocks in a stage.
            -stage_stride: A list of intergers, the stride of first block in
              each stage.
            -expansion_rate: An integer, expansion rate of MBConv.
            -se_ratio: An integer, expansion ratio of SE in MBConv block.
            -head_size: An integer, feature channels per head in Attention.
            -window_size: A list of list of integers, specifying the window size for
              each stage. The length should be the number of stages. For example:
              [None, None, [14, 16], [7, 8]] means the last two stages uses
              14 by 16 and 7 by 8 windows, respectively.
            -position_embedding_size: A list of integers, specifying the position
              embedding sizes for all stages. If the feature maps have larger sizes,
              the position embeddings will be interpolated to match them.
            -use_checkpointing_for_attention: A boolean, specifying whether to use
              checkpointing for attention.
            -global_attention_at_end_of_moat_stage: A boolean, specifying whether
              to use global attention for the last block if the stage consists of
              moat blocks.
            -relative_position_embedding_type: A string, type of relative position
              embedding in Attention. Currently, only '2d_multi_head' is supported.
              If None, no relative position embedding will be used.
            -ln_epsilon: A float, epsilon for layer normalization in Attention.
            -pool_size: An integer, kernel size for pooling in shortcut branch in
              MBConv block and MOAT block.
            -survival_prob: A float, 1 - drop_path_rate.
            -kernel_initializer: Initializer for the kernel weights matrix.
            -bias_initializer: Initializer for the bias vector.
            -name: A string, model name.
            -build_classification_head_with_class_num: An integer, number of
              classes. If None, no classification head will be built.

        Returns:
          A Config class: hparams_config.Config.

        Raises:
          ValueError: If the lengths of block_type, num_blocks and hidden_size are
            not the same.
          ValueError: If the element of block_type is not one of ['mbconv', 'moat'].
        """

        required_keys = ['stem_size', 'block_type', 'num_blocks', 'hidden_size']
        optional_keys = {
            'stage_stride': [2, 2, 2, 2],
            'expansion_rate': 4,
            'se_ratio': 0.25,
            'head_size': 32,
            'window_size': [None, None, [14, 14], [7, 7]],
            'position_embedding_size': [
                None, None,
                _STRIDE_16_POSITION_EMBEDDING_SIZE,
                _STRIDE_32_POSITION_EMBEDDING_SIZE],
            'use_checkpointing_for_attention': False,
            'global_attention_at_end_of_moat_stage': False,
            'relative_position_embedding_type': '2d_multi_head',
            'ln_epsilon': 1e-5,
            'pool_size': 2,
            'survival_prob': None,
            'kernel_initializer': tf.random_normal_initializer(stddev=0.02),
            'bias_initializer': tf.zeros_initializer,
            'build_classification_head_with_class_num': None,
        }
        config = create_config_from_dict(config, required_keys, optional_keys)

        stage_number = len(config.block_type)
        if (len(config.num_blocks) != stage_number or
                len(config.hidden_size) != stage_number):
            raise ValueError('The lengths of block_type, num_blocks and hidden_size ',
                             'should be the same.')
        return config

    def _local_config(self, config, idx, exclude_regex=None):
        """Gets stage-wise config from backbone-wise config."""

        config = copy.deepcopy(config)
        for key in config.__dict__:
            if isinstance(config[key], (list, tuple)):
                if exclude_regex is None or not re.search(exclude_regex, key):
                    config[key] = config[key][idx]
        return config

    def __init__(self, **config):
        super().__init__(name='moat')
        self._config = self._retrieve_config(config)

    def _build_stem(self):
        stem_layers = []
        for i in range(len(self._config.stem_size)):
            conv_layer = tf.keras.layers.Conv2D(
                filters=self._config.stem_size[i],
                kernel_size=3,
                strides=2 if i == 0 else 1,
                padding='same',
                kernel_initializer=self._config.kernel_initializer,
                bias_initializer=self._config.bias_initializer,
                use_bias=True,
                name='conv_{}'.format(i))
            stem_layers.append(conv_layer)
            if i < len(self._config.stem_size) - 1:
                stem_layers.append(self._config.norm_class(name='norm_{}'.format(i)))
                stem_layers.append(tf.keras.layers.Activation(
                    self._config.activation, name='act_{}'.format(i)))
        # The name scope of tf.keras.Sequential is invalid, see error handling
        # in the part of loading checkpoints in function get_model.
        self._stem = tf.keras.Sequential(
            layers=stem_layers,
            name='stem')

    def _build_block(self, local_block_config):
        if local_block_config.block_type == 'mbconv':
            block = MBConvBlock(**local_block_config)
        elif local_block_config.block_type == 'moat':
            block = MOATBlock(**local_block_config)
        else:
            raise ValueError('Unsupported block_type: {}'.format(
                local_block_config.block_type))
        return block

    def _build_classification_head(self):
        self._final_layer_norm = tf.keras.layers.LayerNormalization(
            epsilon=self._config.ln_epsilon,
            name='final_layer_norm')
        self._logits_head = tf.keras.layers.Conv2D(
            filters=self._config.build_classification_head_with_class_num,
            kernel_size=1,
            strides=1,
            kernel_initializer=self._config.kernel_initializer,
            bias_initializer=self._config.bias_initializer,
            padding='same',
            use_bias=True,
            name='logits_head')

    def _adjust_survival_rate(self, local_block_config,
                              block_id, total_num_blocks):
        survival_prob = self._config.survival_prob
        if survival_prob is not None:
            drop_rate = 1.0 - survival_prob
            survival_prob = 1.0 - drop_rate * block_id / total_num_blocks
            local_block_config = local_block_config.replace(
                survival_prob=survival_prob)
        return local_block_config

    def build(self, input_shape) -> None:
        norm_class = tf.keras.layers.experimental.SyncBatchNormalization
        self._config.norm_class = norm_class
        self._config.activation = tf.nn.gelu

        self._build_stem()

        self._blocks = []
        total_num_blocks = sum(self._config.num_blocks)

        for stage_id in range(len(self._config.block_type)):
            stage_config = self._local_config(self._config, stage_id, '^stem.*')
            stage_blocks = []

            for local_block_id in range(stage_config.num_blocks):
                local_block_config = copy.deepcopy(stage_config)

                block_stride = 1
                if local_block_id == 0:
                    block_stride = self._config.stage_stride[stage_id]
                local_block_config = local_block_config.replace(
                    block_stride=block_stride)

                block_id = sum(self._config.num_blocks[:stage_id]) + local_block_id
                local_block_config = self._adjust_survival_rate(
                    local_block_config,
                    block_id, total_num_blocks)

                block_name = 'block_{:0>2d}_{:0>2d}'.format(stage_id, local_block_id)
                local_block_config.name = block_name

                if (local_block_id == stage_config.num_blocks - 1 and
                        self._config.block_type[stage_id] == 'moat' and
                        self._config.global_attention_at_end_of_moat_stage):
                    local_block_config.window_size = None

                block = self._build_block(local_block_config)
                stage_blocks.append(block)

            self._blocks.append(stage_blocks)

        if self._config.build_classification_head_with_class_num is not None:
            self._build_classification_head()

    def call(self, inputs, training=False, mask=None):
        endpoints = {}

        output = self._stem(inputs, training=training)
        endpoints['stage1'] = output
        endpoints['res1'] = self._config.activation(output)

        for stage_id, stage_blocks in enumerate(self._blocks):
            for block in stage_blocks:
                output = block(output, training=training)
            endpoints['stage{}'.format(stage_id + 2)] = output
            endpoints['res{}'.format(stage_id + 2)] = self._config.activation(output)

        if self._config.build_classification_head_with_class_num is None:
            return endpoints
        else:
            reduce_axes = list(range(1, output.shape.rank - 1))
            output = tf.reduce_mean(output, axis=reduce_axes, keepdims=True)
            output = self._final_layer_norm(output)
            output = self._logits_head(output, training=training)
            logits = tf.squeeze(output, axis=[1, 2])
            return logits


moat0_config = Config(
    stem_size=[64, 64],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 3, 7, 2],
    hidden_size=[96, 192, 384, 768],
)

moat1_config = Config(
    stem_size=[64, 64],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 6, 14, 2],
    hidden_size=[96, 192, 384, 768],
)

moat2_config = Config(
    stem_size=[128, 128],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 6, 14, 2],
    hidden_size=[128, 256, 512, 1024],
)

moat3_config = Config(
    stem_size=[160, 160],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 12, 28, 2],
    hidden_size=[160, 320, 640, 1280],
)

moat4_config = Config(
    stem_size=[256, 256],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 12, 28, 2],
    hidden_size=[256, 512, 1024, 2048],
)

tiny_moat0_config = Config(
    stem_size=[32, 32],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 3, 7, 2],
    hidden_size=[32, 64, 128, 256],
)

tiny_moat1_config = Config(
    stem_size=[40, 40],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 3, 7, 2],
    hidden_size=[40, 80, 160, 320],
)

tiny_moat2_config = Config(
    stem_size=[56, 56],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 3, 7, 2],
    hidden_size=[56, 112, 224, 448],
)

tiny_moat3_config = Config(
    stem_size=[80, 80],
    block_type=['mbconv', 'mbconv', 'moat', 'moat'],
    num_blocks=[2, 3, 7, 2],
    hidden_size=[80, 160, 320, 640],
)

no_relative_pe = Config(
    relative_position_embedding_type=None,
)


def get_model(
        name: str,
        input_shape,
        window_size=None,
        survival_rate: Optional[float] = None,
        pool_size: Optional[int] = 3,
        override_config=None,
        pretrained_weights_path: Optional[str] = None,
        remove_position_embedding: Optional[bool] = None,
        return_config: Optional[bool] = False,
        strict_loading: Optional[bool] = False,
        use_checkpointing_for_attention: Optional[bool] = False,
        global_attention_at_end_of_moat_stage: Optional[bool] = False,
) -> ...:
    """Gets MOAT according to name.

    Args:
      name: A string, the model name. It must be the key of the global dictionary
        _CHECKPOINTS_PATH.
      input_shape: A list of integers with length 3, the input shape to
        initialize the model.
      window_size: A list of lists of integers, specifying the window size for
        each stage. The length should be the number of stages. For example:
        [None, None, [14, 16], [7, 8]] means the last two stages uses
        14 by 16 and 7 by 8 windows, respectively. Global attention will be used
        if it is None.
      survival_rate: A float, this equals 1 - drop path rate.
      pool_size: An integer, kernel size for pooling in shortcut branch. For
        classification, pool size 2x2 saves model flops. For downstream tasks,
        pool size 3x3 is preferred for better feature alignments.
      override_config: A dictionary overriding the model config.
      pretrained_weights_path: A string, specifying the path to load the
        pretrained checkpoint. If None, no pretrained checkpoint will be loaded.
      remove_position_embedding: A boolean, specifying whether to remove position
        embedding. If None, the usage of position embedding will depend on the
        model name. This flag can enable the model that is pretrained with
        position embedding to be finetuned without position embedding in
        downstream tasks.
      return_config: A boolean, whether to return config or model.
      strict_loading: A boolean, if True, all variables in the checkpoint must
        be found in the model.
      use_checkpointing_for_attention: A boolean, specifying whether to use
        checkpointing for attention.
      global_attention_at_end_of_moat_stage: A boolean, specifying whether to use
        global attention for the last block if the stage consists of moat blocks.

    Returns:
      config: Config to build moat. This is returned when return_config is True.
      moat: moat model. This is returned when return_config is False.

    Raises:
      ValueError: If the length of input_shape does not equal 3.
    """

    if len(input_shape) != 3:
        raise ValueError('The input shape should be a list of length 3.')

    config = _get_moat_config_via_model_name(name)
    config.window_size = window_size
    config.pool_size = pool_size
    config.use_checkpointing_for_attention = use_checkpointing_for_attention
    config.global_attention_at_end_of_moat_stage = (
        global_attention_at_end_of_moat_stage)
    if survival_rate is not None:
        config.survival_prob = survival_rate
    if override_config is not None:
        config.update(override_config)
    if remove_position_embedding:
        config.update(no_relative_pe)
    if return_config:
        return config

    moat = MOAT(**config)
    moat(tf.keras.Input(shape=input_shape))
    if pretrained_weights_path:
        moat = _load_moat_pretrained_checkpoint(
            moat, pretrained_weights_path, strict_loading)
    return moat


def _get_moat_config_via_model_name(name):
    """Gets the moat config according to name."""

    if name in ['moat0_pretrain_224_1k', 'moat0_pretrain_224_no_pe_1k']:
        config = copy.deepcopy(moat0_config)
        config.survival_prob = 0.8
    elif name in ['moat1_pretrain_224_1k', 'moat1_pretrain_224_no_pe_1k']:
        config = copy.deepcopy(moat1_config)
        config.survival_prob = 0.7
    elif name in ['moat2_finetune_384_22k', 'moat2_finetune_384_no_pe_22k']:
        config = copy.deepcopy(moat2_config)
        config.survival_prob = 0.7
    elif name in ['moat3_finetune_512_22k', 'moat3_finetune_512_no_pe_22k']:
        config = copy.deepcopy(moat3_config)
        config.survival_prob = 0.4
    elif name in ['moat4_finetune_512_22k', 'moat4_finetune_512_no_pe_22k']:
        config = copy.deepcopy(moat4_config)
        config.survival_prob = 0.3
    elif name in ['tiny_moat0_pretrain_256_1k',
                  'tiny_moat0_pretrain_256_no_pe_1k']:
        config = copy.deepcopy(tiny_moat0_config)
        config.survival_prob = 1.0
    elif name in ['tiny_moat1_pretrain_256_1k',
                  'tiny_moat1_pretrain_256_no_pe_1k']:
        config = copy.deepcopy(tiny_moat1_config)
        config.survival_prob = 1.0
    elif name in ['tiny_moat2_pretrain_256_1k',
                  'tiny_moat2_pretrain_256_no_pe_1k']:
        config = copy.deepcopy(tiny_moat2_config)
        config.survival_prob = 1.0
    elif name in ['tiny_moat3_pretrain_256_1k',
                  'tiny_moat3_pretrain_256_no_pe_1k']:
        config = copy.deepcopy(tiny_moat3_config)
        config.survival_prob = 0.9
    else:
        raise ValueError('Not supported moat name %s' % name)
    if 'no_pe' in name:
        config.update(no_relative_pe)
    return config


def _load_moat_pretrained_checkpoint(
        moat, path, strict_loading=False):
    """Loads the TF1 pretrained weights for MOAT."""

    checkpoint_reader = tf.train.load_checkpoint(path)
    variable_to_shape_map = checkpoint_reader.get_variable_to_shape_map()

    model_var_name = sorted([var.name for var in moat.trainable_variables])
    ckpt_var_name = list(sorted(variable_to_shape_map.keys()))

    # This for loop ensures all moat variables can be found in the checkpoint.
    for var in moat.trainable_variables:
        name_to_find = var.name

        # We change the TF2 variable name for loading TF1 checkpoint.
        name_to_find = name_to_find.replace(':0', '')

        if name_to_find not in ckpt_var_name:
            if _STEM_LAYER_NAME_SCOPE + name_to_find in ckpt_var_name:
                # This handles the invalid name scope of tf.keras.sequential.
                logging.info('Add name scope %s to the variable: %s',
                             _STEM_LAYER_NAME_SCOPE,
                             name_to_find)
                name_to_find = _STEM_LAYER_NAME_SCOPE + name_to_find
            else:
                raise ValueError('Variable name %s is not found in the checkpoint'
                                 % var.name)

        if var.shape != variable_to_shape_map[name_to_find]:
            raise ValueError('Inconsistent shape for var: %s ' % var.name,
                             'should be %s' % var.shape,
                             'but found %s' % variable_to_shape_map[name_to_find])

        var_mean_before_loading = tf.reduce_mean(var)
        var.assign(tf.Variable(checkpoint_reader.get_tensor(name_to_find)))
        var_mean_after_loading = tf.reduce_mean(var)
        if tf.math.equal(var_mean_before_loading,
                         var_mean_after_loading).numpy():
            raise ValueError('Loading error for the variable: %s' % name_to_find)

    logging.info('All variables in moat are found in the checkpoint.')

    # This for loop finds all checkpoint variables that are not found in moat.
    for var_name in variable_to_shape_map.keys():
        name_in_ckpt = var_name.replace(_EMA_VARIABLE_NAME_POSTFIX, '')
        name_in_ckpt = name_in_ckpt + ':0'

        if (name_in_ckpt not in model_var_name and
                name_in_ckpt.replace(_STEM_LAYER_NAME_SCOPE, '')
                not in model_var_name):
            # This handles the invalid name scope of tf.keras.sequential.
            if strict_loading:
                raise ValueError('In ckpt but not in model_var_name', name_in_ckpt)
            else:
                logging.info('In ckpt but not in model_var_name: %s',
                             name_in_ckpt)

    if strict_loading:
        logging.info('All variables in checkpoint are found in moat.')
    return moat
