# 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.
"""TAO Core hooks.

Hooks, like callbacks, are called at specific intervals. They typically inherit from a base object
and should override its methods.
"""

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

from collections import deque
import glob
import itertools
import logging
import os
import re
import signal

import keras
import tensorflow as tf
from tensorflow.python.keras import backend as backend_python

logger = logging.getLogger(__name__)


class StartAndStepOpHook(tf.estimator.SessionRunHook):
    """Hook to extend calls to ``MonitoredSession.run()``.

    Given a list of ops, this hook will execute this op once when the training session is created,
    and after that execute it with every following training step.
    Using this hook in combination with a buffer, you can preallocate one sample at the start,
    and asynchronously add one sample to the buffer with each step thereafter.
    """

    def __init__(self, ops):
        """Initialize the hook.

        Args:
            ops: A list of tensors, to be executed once when the session is created, and
                every time a run is performed after that.
        """
        self._ops = ops

    def before_run(self, run_context):
        """Called before each call to run().

        Run the ops each run.

        Args:
            run_context: A `SessionRunContext` object.
        Returns:
            A `SessionRunArgs` object.
        """
        return tf.estimator.SessionRunArgs(self._ops)

    def after_create_session(self, session, coord):
        """Called when new TensorFlow session is created.

        Runs the ops once.

        Args:
            session: A TensorFlow Session that has been created.
            coord: A Coordinator object which keeps track of all threads.
        """
        session.run(self._ops)


class TFCheckpointListener(tf.estimator.CheckpointSaverListener):
    """Listener hook to call upon any methods when a tensorflow file gets saved.

    This hook is basically a shell to incorporate a callback.

    For callbacks related to Keras HDF5 files, use KerasCheckpointListener.
    """

    def __init__(self, checkpoint_dir, after_save_callable, cleanup=None):
        """Initializes a `TFCheckpointListener`.

        Args:
            checkpoint_dir (str): The checkpoint directory where the checkpoint files will be in.
            after_save_callable ((ckpt_base (str), ckpt_files (list(str)), step (int)) -> None):
                The callable to take in all the checkpoint files. `ckpt_base` is the file with base
                `.ckpt` extension while `ckpt_files` are all the files generated by tensorflow.
            cleanup ((step (int)) -> None): Optional. Method that is called at the end of the
                Tensorflow session. Default None.
        """
        self._checkpoint_dir = checkpoint_dir
        self._after_save_callable = after_save_callable
        self._cleanup = cleanup

    def after_save(self, session, global_step_value):
        """Called right after saving the tensorflow checkpoint. This will call the callable.

        Args:
            session: the current instance of type `tf.Session`.
            global_step_value (int): the current global step.
        """
        all_files = os.listdir(self._checkpoint_dir)
        extension = ".ckpt-{}".format(global_step_value)
        ckpt_files = list(filter(lambda f: extension in f, all_files))
        ckpt_files_with_path = list(
            map(lambda s: os.path.join(self._checkpoint_dir, s), ckpt_files)
        )
        ckpt_base = os.path.splitext(ckpt_files_with_path[0])[0]
        self._after_save_callable(ckpt_base, ckpt_files_with_path, global_step_value)

    def end(self, session, global_step_value):
        """Run at the end of the session."""
        if self._cleanup:
            self._cleanup(global_step_value)


class KerasCheckpointListener(tf.estimator.CheckpointSaverListener):
    """Listener hook to save keras model snapshots.

    Usage of this checkpoint listener toggle Keras variable initialization off during the
    lifetime of the session.

    Note that this object should be used in conjunction with the `KerasModelHook`.
    """

    def __init__(
        self,
        model,
        checkpoint_dir,
        after_save_callable=None,
        max_to_keep=None,
        prefix="model",
    ):
        """Initialize a `KerasCheckpointListener`.

        Args:
            model (Keras.models.Models): a instance of a ``Keras.models.Model`` object.
            checkpoint_dir (str): base directory for the checkpoint files.
            after_save_callable ((filename (str), step (int)) -> None): Optional. A function to
                call after the files are saved for any necessary post processing with the files.
                Default None.
            max_to_keep (int): Optional. Maximum number of Keras checkpoints to keep. Keeps the
                latest `max_to_keep` checkpoints. Default None.
            prefix (str): Prefix to add to the checkpoints' filenames. Default "model".
        """
        if os.environ.get("TF_KERAS"):
            self.K = tf.keras
            self._previous_MANUAL_VAR_INIT = backend_python._MANUAL_VAR_INIT
        else:
            self.K = keras
            self._previous_MANUAL_VAR_INIT = (
                keras.backend.tensorflow_backend._MANUAL_VAR_INIT
            )

        self._model = model
        self._checkpoint_dir = checkpoint_dir
        self._after_save_callable = after_save_callable
        self._max_to_keep = max_to_keep
        self._latest_checkpoints = deque(maxlen=self._max_to_keep)
        self._prefix = prefix

    def begin(self):
        """Called after starting the session."""
        pattern = r"^%s.keras-(\d+)\.hdf5$" % re.escape(
            os.path.join(self._checkpoint_dir, self._prefix)
        )
        compiled = re.compile(pattern)

        def extract_model_number(filename):
            s = compiled.findall(filename)
            return int(s[0]) if s else -1, filename

        filenames = glob.glob(os.path.join(self._checkpoint_dir, "*.hdf5"))
        # Weed out filenames that do not match the pattern.
        filenames = [
            filename for filename in filenames if compiled.match(filename) is not None
        ]
        sorted_filenames = sorted(filenames, key=extract_model_number)
        self._latest_checkpoints.extend(sorted_filenames)

    def after_save(self, session, global_step_value):
        """Called right after saving the tensorflow checkpoint.

        Args:
            session: the current instance of type `tf.Session`.
            global_step_value (int): the current global step.
        """
        self.K.backend.set_session(session)
        self.K.backend.manual_variable_initialization(True)
        filename = os.path.join(
            self._checkpoint_dir, "%s.keras-%s.hdf5" % (self._prefix, global_step_value)
        )
        self._model.save(filename, overwrite=True)

        if self._after_save_callable:
            self._after_save_callable(filename, global_step_value)

        self._cleanup(filename)

    def _cleanup(self, filename):
        # Clean up old checkpoints if need be.
        if (
            self._max_to_keep is not None
            and len(self._latest_checkpoints) >= self._max_to_keep
        ):
            # First, get the name of the latest checkpoint to remove.
            old_checkpoint = self._latest_checkpoints.popleft()
            if os.path.exists(old_checkpoint):
                os.remove(old_checkpoint)

        # Now, add the new checkpoint.
        if filename not in self._latest_checkpoints:
            self._latest_checkpoints.append(filename)
        else:
            logger.info("Overwritten Keras model: {}.".format(filename))

    def end(self, session, global_step_value):
        """Run at the end of the session, reset the old variale initialization setting."""
        self.K.backend.manual_variable_initialization(self._previous_MANUAL_VAR_INIT)


class KerasModelHook(tf.estimator.SessionRunHook):
    """Hook to extend calls to ``MonitoredSession.run()``.

    Extracts relevant information and ops from a Keras model. Extracts the update ops from a keras
    model, that are for example relevant and set when using batch_norm.

    Note that there are two major intended use cases for this hook:
        * During training, layers that have non-trainable, but updatable weights (e.g. batch norm)
          need this hook to make sure these updates are run.
        * When loading a keras hdf5 from disk (e.g. as pretrained weights for finetuning, or
          simply for inference), the weights' values in the hdf5 need to be copied over into
          the TensorFlow graph, otherwise they would be reinitialized from scratch by running
          each tf.Variable's `initializer`, which is typically random or null, and produce
          garbage results.
    """

    def __init__(self, models, ignore_initialized_values=False):
        """Initialize the hook.

        Args:
            models (list): A list of `keras.Model` used to extract relevant information from.
            A single `keras.Model` is also supported for backward compatibility.
            Update ops are run for all of the `keras.Model` with each session step.
            ignore_initialized_values (bool): If True, Keras variable values set in Keras
                background session are not copied to Tensorflow session at session creation. Set
                this True when restoring session from a Tensorflow checkpoint.
        """
        self.models = models if isinstance(models, list) else [models]
        self._updates = [update for model in self.models for update in model.updates]
        if os.environ.get("TF_KERAS"):
            self.K = tf.keras
            self._previous_MANUAL_VAR_INIT = backend_python._MANUAL_VAR_INIT
        else:
            self.K = keras
            self._previous_MANUAL_VAR_INIT = (
                keras.backend.tensorflow_backend._MANUAL_VAR_INIT
            )
        self._ignore_initialized_values = ignore_initialized_values
        self._variables_initialized = []

    def begin(self):
        """Called once before using the session.

        Assignment operations are added to each tensorflow weight variable. This needs to be done
        before the graph is finalized. The actual running of assignment operations is done later
        in `after_create_session()`.
        """
        self.K.backend.manual_variable_initialization(True)
        for model in self.models:
            for x in model.weights:
                tf_dtype = tf.as_dtype(x.dtype.name.split("_")[0])
                assign_placeholder = tf.compat.v1.placeholder(tf_dtype)
                assign_op = x.assign(assign_placeholder)
                x._assign_placeholder = assign_placeholder
                x._assign_op = assign_op
                self._variables_initialized.append(
                    tf.compat.v1.is_variable_initialized(x)
                )

    def end(self, session):
        """Run at the end of the session, clears the closed keras session so it can be reused."""
        self.K.backend.clear_session()
        # Reset the old variable initialization setting
        self.K.backend.manual_variable_initialization(self._previous_MANUAL_VAR_INIT)

    def after_create_session(self, session, coord):
        """Called when new TensorFlow session is created.

        Use the assignment operations created in `begin()` and load in all initialized weights from
        the Keras session over to our Tensorflow session.

        Args:
            session: A TensorFlow Session that has been created.
            coord: A Coordinator object which keeps track of all threads.
        """
        if not self._ignore_initialized_values:
            vars_all = list(itertools.chain(*[m.weights for m in self.models]))
            vars_is_initialized = self.K.backend.get_session().run(
                self._variables_initialized
            )
            vars_to_initialize = list(
                itertools.compress(data=vars_all, selectors=vars_is_initialized)
            )
            assign_weights = self.K.backend.get_session().run(vars_to_initialize)
            assign_ops = [v._assign_op for v in vars_to_initialize]
            assign_phs = [v._assign_placeholder for v in vars_to_initialize]
            session.run(assign_ops, feed_dict=dict(zip(assign_phs, assign_weights)))
        else:
            # Otherwise TF spits out a bunch of crap.
            # Note that this looks like a potential can of worms. No mention of this
            # `mark_used` method can be found in TensorFlow's documentation.
            try:
                for init_flag in self._variables_initialized:
                    init_flag.mark_used()
            except Exception:
                pass
        self.K.backend.set_session(session)

    def before_run(self, run_context):
        """Called before each call to run().

        Run the ops each run.

        Args:
            run_context: A `SessionRunContext` object.
        Returns:
            A `SessionRunArgs` object.
        """
        return tf.estimator.SessionRunArgs(self._updates)


class SignalHandlerHook(tf.estimator.SessionRunHook):
    """Hook to allow clean shutdown (incl. checkpointing) when receiving SIGUSR1.

    Catch SIGUSR1 during training. When the step is finished, request session to stop, which causes
    the session to call end() in all hooks. This will trigger checkpointing, if a hook is
    configured for creating checkpoints.
    """

    def __init__(self):
        """Initialize the hook."""
        self._signal_caught = False
        self._previous_usr1_handler = None
        self._main_process_pid = -1

    def begin(self):
        """Register the signal handler."""
        # Remember the previously configured handler.
        self._previous_usr1_handler = signal.getsignal(signal.SIGUSR1)
        # begin() is always called by the main process, get its PID.
        self._main_process_pid = os.getpid()
        # Assign a new handler.
        signal.signal(signal.SIGUSR1, self._handle_signal)

    def after_run(self, run_context, run_values):
        """If self._signal_caught is set, end the session."""
        if self._signal_caught:
            logger.info("SignalHandlerHook requests session to stop.")
            run_context.request_stop()

    def end(self, session):
        """Assign the previously configured signal handler after session ends."""
        signal.signal(signal.SIGUSR1, self._previous_usr1_handler)

    def _handle_signal(self, signum, frame):
        """Flip self._signal_caught flag, when a signal is received."""
        this_process_pid = os.getpid()
        # Flip self._signal_caught flag only in the main process.
        # In any other process, ignore SIGUSR1.
        if this_process_pid == self._main_process_pid:
            logger.info("SignalHandlerHook caught signal {}.".format(signum))
            self._signal_caught = True
