# Copyright 2016 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.
# ==============================================================================
"""Ops to use variables as resources."""

# pylint: disable=g-bad-name
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import gen_resource_variable_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resources
# go/tf-wildcard-import
# pylint: disable=wildcard-import
from tensorflow.python.ops.gen_resource_variable_ops import *
# pylint: enable=wildcard-import


def _register_variable_read(read, collections, trainable):
  """Helper function to put a read from a variable in the collections."""
  if collections is None:
    collections = []
  if (trainable and
      ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES not in collections):
    collections = (list(collections) +
                   [ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES])
    ops.add_to_collections(collections, read)


class ResourceVariable(object):
  """Variable based on resource handles.

  TODO(apassos): fill this out explaining the semantics and Variable
  compatibility when the API has settled more.

  """

  def __init__(self,
               initial_value=None,
               name=None,
               trainable=True,
               collections=None,
               dtype=None,
               shape=None):
    """Creates a variable.

    Args:
      initial_value: An `Output` or Python object convertible to an `Output`
        representing the initial value of this variable.
      name: The name of this variable. Automatically uniquified.
      trainable: Whether the global read of this variable will be used for
        training.
      collections: Additional collections to which the `read` operation for
        this variable is to be added. Defaults to [].
      dtype: The type of this variable. Can be omitted if it can be deduced
        from the initial_value. If different from the type of the initial
        value it will be cast to this type.
      shape: The shape of this variable. Only specify if there is no initial
        value but shape inference is desired.
    """
    if initial_value is not None:
      initial_value = ops.convert_to_tensor(initial_value)
    if dtype is None:
      assert initial_value is not None, ("Trying to create a resource variable "
                                         "with no dtype or initial value. At"
                                         " least one of these must be set.")
      dtype = initial_value.dtype
    elif initial_value is not None:
      initial_value = math_ops.cast(initial_value, dtype)
    if shape is None:
      if initial_value is not None:
        shape = initial_value.get_shape().as_proto()
      else:
        shape = tensor_shape.unknown_shape()
    else:
      shape = tensor_shape.as_shape(shape)

    self._dtype = dtype
    with ops.name_scope(name, "Variable", [initial_value]) as name:
      self._handle = gen_resource_variable_ops.var_handle_op(shared_name=name,
                                                             name=name,
                                                             dtype=dtype,
                                                             shape=shape)

      with ops.name_scope("IsInitialized"):
        self._is_initialized_op = (
            gen_resource_variable_ops.var_is_initialized_op(self._handle))
      if initial_value is not None:
        with ops.name_scope("Create"):
          self._initialize_op = gen_resource_variable_ops.create_variable_op(
              self._handle, initial_value)
        resources.register_resource(self._handle,
                                    self._initialize_op,
                                    self._is_initialized_op)

      with ops.name_scope("Read"):
        self._value = gen_resource_variable_ops.read_variable_op(
            self._handle, dtype=self._dtype)
      _register_variable_read(
          self._value, trainable=trainable, collections=collections)

  @property
  def dtype(self):
    """The dtype of this variable."""
    return self._dtype

  @property
  def create(self):
    """The op responsible for initializing this variable."""
    return self._initialize_op

  @property
  def handle(self):
    """The handle by which this variable can be accessed."""
    return self._handle

  @property
  def value(self):
    """A cached operation which reads the value of this variable."""
    return self._value

  @property
  def op(self):
    """The op which reads the value of this variable."""
    return self._value.op

  def eval(self, session=None):
    """Evaluates and returns the value of this variable."""
    return self._value.eval(session=session)

  def read_value(self, collections=None, trainable=True):
    """Constructs an op which reads the value of this variable.

    Should be used when there are multiple reads, or when it is desirable to
    read the value only after some condition is true.

    Args:
     collections: any collections in which this operation should be inserted.
     trainable: whether this read is to be used for training.

    Returns:
     the read operation.
    """
    with ops.name_scope("Read"):
      value = gen_resource_variable_ops.read_variable_op(
          self._handle, dtype=self._dtype)
    _register_variable_read(value, collections=collections, trainable=trainable)
    return value

  def sparse_read(self, indices, collections=None, trainable=True, name=None):
    with ops.name_scope("Gather" if name is None else name):
      value = gen_resource_variable_ops.resource_gather(
          self._handle, indices, dtype=self._dtype)
    _register_variable_read(value, collections=collections, trainable=trainable)
    return value


# pylint: disable=unused-argument
def _dense_var_to_tensor(var, dtype=None, name=None, as_ref=False):
  if dtype is not None and dtype != var.value.dtype:
    print("trying to switch the dtype to ", dtype, " from ", var.value.dtype)
    return NotImplemented
  return var.value

# Register a conversion function which reads the value of the variable,
# allowing instances of the class to be used as tensors.
ops.register_tensor_conversion_function(ResourceVariable, _dense_var_to_tensor)
