"""Pretrained weight loading hooks."""
# Copyright (c) 2023, NVIDIA CORPORATION.  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.
#
# 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.
import logging
import re
import sys
import tensorflow as tf

__all__ = ["PretrainedWeightsLoadingHook"]

logger = logging.getLogger(__name__)
tf.logging.set_verbosity(tf.logging.WARN)


# pylint: disable=protected-access
# Currently variable_scope doesn't provide very good APIs to access
# all variables under scope and retrieve and check existing scopes.

def get_variable_full_name(var):
    """Returns the full name of a variable.

    For normal Variables, this is the same as the var.op.name.  For
    sliced or PartitionedVariables, this name is the same for all the
    slices/partitions. In both cases, this is normally the name used in
    a checkpoint file.
    Args:
    var: A `Variable` object.
    Returns:
    A string that is the full name.
    """
    if var._save_slice_info:
        return var._save_slice_info.full_name
    return var.op.name


def assign_from_checkpoint(model_path, var_list, ignore_missing_vars=False, remove_head=False):
    """Creates an operation to assign specific variables from a checkpoint.

    Args:
    model_path: The full path to the model checkpoint. To get latest checkpoint
      use `model_path = tf.train.latest_checkpoint(checkpoint_dir)`
    var_list: A list of (possibly partitioned) `Variable` objects or a
      dictionary mapping names in the checkpoint to the corresponding variables
      or list of variables to initialize from that checkpoint value. For
      partitioned Variables, the name in the checkpoint must be the full
      variable, not the name of the partitioned variable, eg. "my_var" rather
      than "my_var/part_4". If empty, returns no_op(), {}.
    ignore_missing_vars: Boolean, if True ignore variables missing in the
      checkpoint with a warning instead of failing.
    Returns:
    the restore_op and the feed_dict that need to be run to restore var_list.
    Raises:
    ValueError: If `ignore_missing_vars` is False and the checkpoint specified
        at `model_path` is missing one of the variables in `var_list`.
    """
    # Normalize var_list into a dictionary mapping names in the
    # checkpoint to the list of variables to initialize from that
    # checkpoint variable. Sliced (including partitioned) variables will
    # end up under the same key.
    grouped_vars = {}
    if isinstance(var_list, (tuple, list)):
        for var in var_list:
            ckpt_name = get_variable_full_name(var)
            if ckpt_name not in grouped_vars:
                grouped_vars[ckpt_name] = []
            grouped_vars[ckpt_name].append(var)
    else:
        var_list_items = list(var_list.items())
        # The last 4 graph nodes will correspond to the layer producing the
        # final output with channels equal to the num classes. Hence removing them
        # to do transfer learning.

        var_list_items = var_list_items[:-4]
        for var_list_item in var_list_items:
            ckpt_name, value = var_list_item
            if isinstance(value, (tuple, list)):
                grouped_vars[ckpt_name] = value
            else:
                grouped_vars[ckpt_name] = [value]

    # Read each checkpoint entry. Create a placeholder variable and
    # add the (possibly sliced) data from the checkpoint to the feed_dict.
    reader = tf.compat.v1.train.NewCheckpointReader(model_path)
    feed_dict = {}
    assign_ops = []
    for ckpt_name in grouped_vars:
        if not reader.has_tensor(ckpt_name):
            log_str = 'Checkpoint is missing variable [%s]' % ckpt_name
            if ignore_missing_vars:
                logger.debug(log_str)
                continue
            # @vpraveen: Removing the else coming from
            # pylint errors.
            raise ValueError(log_str)
        ckpt_value = reader.get_tensor(ckpt_name)
        if remove_head:
            restore_variables = grouped_vars[ckpt_name][:-1]
        else:
            restore_variables = grouped_vars[ckpt_name]
        for var in restore_variables:
            placeholder_tensor = tf.compat.v1.placeholder(
                dtype=var.dtype.base_dtype,
                shape=var.get_shape(),
                name='placeholder/' + var.op.name
            )
            assign_ops.append(var.assign(placeholder_tensor))

            if not var._save_slice_info:
                if var.get_shape() != ckpt_value.shape:
                    raise ValueError(
                        'Total size of new array must be unchanged for %s '
                        'lh_shape: [%s], rh_shape: [%s]' %
                        (ckpt_name, str(ckpt_value.shape), str(var.get_shape())))

                feed_dict[placeholder_tensor] = ckpt_value.reshape(ckpt_value.shape)

            else:
                slice_dims = zip(var._save_slice_info.var_offset,
                                 var._save_slice_info.var_shape)

                slice_dims = [(start, start + size) for (start, size) in slice_dims]
                slice_dims = [slice(*x) for x in slice_dims]

                slice_value = ckpt_value[slice_dims]
                slice_value = slice_value.reshape(var._save_slice_info.var_shape)

                feed_dict[placeholder_tensor] = slice_value

    print_op = tf.print(
        "[GPU] Restoring pretrained weights from: %s" % (
            model_path
        ),
        output_stream=sys.stdout
    )

    with tf.control_dependencies([print_op]):
        assign_op = tf.group(*assign_ops)

    return assign_op, feed_dict


def build_assigment_map(prefix=None, skip_variables_regex=None):
    """Generate assigment map for loading checkpoints."""

    all_vars = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, scope=prefix)

    if not prefix:
        prefix = ''

    assignment_map = {}

    for var in all_vars:

        var_name = var.name

        var_name_filter = (
            var_name[-11:] in "/Momentum:0" or
            var_name[-11:] in "/Adadelta:0" or
            var_name[-13:] in "/Adadelta_1:0" or
            var_name[-7:] in "/Adam:0" or
            var_name[-9:] in "/Adam_1:0" or
            var_name[-10:] in "/Adagrad:0" or
            var_name[-10:] in "/RMSProp:0" or
            var_name[-12:] in "/RMSProp_1:0" or
            var_name[-16:] in "/LARSOptimizer:0"
        )
        if var_name_filter:
            continue

        # Trim the index of the variable.
        if ':' in var_name:
            var_name = var_name[:var_name.rindex(':')]

        if skip_variables_regex and re.match(skip_variables_regex, var_name[len(prefix):]):
            continue

        assignment_map[var_name[len(prefix):]] = var
        # assignment_map[var_name] = var

    return assignment_map


class PretrainedWeightsLoadingHook(tf.estimator.SessionRunHook):
    """Hook for loading pretrained weights."""

    def __init__(self, prefix, checkpoint_path, skip_variables_regex=None, remove_head=False):
        """Initialize."""
        self._prefix = prefix
        self._checkpoint_path = checkpoint_path
        self._skip_variables_regex = skip_variables_regex
        self._is_initialized = False
        self._init_op = None
        self._init_feed_dict = None
        self.remove_head = remove_head

    def begin(self):
        """Begin."""
        vars_to_load = build_assigment_map(
            prefix=self._prefix,
            skip_variables_regex=self._skip_variables_regex
        )

        self._init_op, self._init_feed_dict = assign_from_checkpoint(
            model_path=self._checkpoint_path,
            var_list=vars_to_load,
            ignore_missing_vars=True,
            remove_head=self.remove_head
        )

    def after_create_session(self, session, coord=None):
        """Run session."""
        if not self._is_initialized:
            session.run(self._init_op, feed_dict=self._init_feed_dict)
            logger.info("Pretrained weights loaded with success...\n")

            self._is_initialized = True
