# Copyright 2018 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.

"""TensorBoard helper routine for TF op evaluator.

Requires TensorFlow.
"""

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

import threading


class PersistentOpEvaluator(object):
  """Evaluate a fixed TensorFlow graph repeatedly, safely, efficiently.

  Extend this class to create a particular kind of op evaluator, like an
  image encoder. In `initialize_graph`, create an appropriate TensorFlow
  graph with placeholder inputs. In `run`, evaluate this graph and
  return its result. This class will manage a singleton graph and
  session to preserve memory usage, and will ensure that this graph and
  session do not interfere with other concurrent sessions.

  A subclass of this class offers a threadsafe, highly parallel Python
  entry point for evaluating a particular TensorFlow graph.

  Example usage:

      class FluxCapacitanceEvaluator(PersistentOpEvaluator):
        \"\"\"Compute the flux capacitance required for a system.

        Arguments:
          x: Available power input, as a `float`, in jigawatts.

        Returns:
          A `float`, in nanofarads.
        \"\"\"

        def initialize_graph(self):
          self._placeholder = tf.placeholder(some_dtype)
          self._op = some_op(self._placeholder)

        def run(self, x):
          return self._op.eval(feed_dict: {self._placeholder: x})

      evaluate_flux_capacitance = FluxCapacitanceEvaluator()

      for x in xs:
        evaluate_flux_capacitance(x)
  """

  def __init__(self):
    super(PersistentOpEvaluator, self).__init__()
    self._session = None
    self._initialization_lock = threading.Lock()

  def _lazily_initialize(self):
    """Initialize the graph and session, if this has not yet been done."""
    # TODO(nickfelt): remove on-demand imports once dep situation is fixed.
    import tensorflow.compat.v1 as tf
    with self._initialization_lock:
      if self._session:
        return
      graph = tf.Graph()
      with graph.as_default():
        self.initialize_graph()
      # Don't reserve GPU because libpng can't run on GPU.
      config = tf.ConfigProto(device_count={'GPU': 0})
      self._session = tf.Session(graph=graph, config=config)

  def initialize_graph(self):
    """Create the TensorFlow graph needed to compute this operation.

    This should write ops to the default graph and return `None`.
    """
    raise NotImplementedError('Subclasses must implement "initialize_graph".')

  def run(self, *args, **kwargs):
    """Evaluate the ops with the given input.

    When this function is called, the default session will have the
    graph defined by a previous call to `initialize_graph`. This
    function should evaluate any ops necessary to compute the result of
    the query for the given *args and **kwargs, likely returning the
    result of a call to `some_op.eval(...)`.
    """
    raise NotImplementedError('Subclasses must implement "run".')

  def __call__(self, *args, **kwargs):
    self._lazily_initialize()
    with self._session.as_default():
      return self.run(*args, **kwargs)
