# Copyright 2015 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.
# ==============================================================================

# pylint: disable=invalid-name
"""Save and restore variables."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import os.path
import re
import time
import uuid

import numpy as np
import six

from google.protobuf import text_format

from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.client import session
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import errors
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_io_ops
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import training_util
from tensorflow.python.training.checkpoint_state_pb2 import CheckpointState
from tensorflow.python.util import compat


# Op names which identify variable reads which should be saved.
_VARIABLE_OPS = set(["Variable",
                     "VariableV2",
                     "AutoReloadVariable",
                     "VarHandleOp",
                     "ReadVariableOp"])


def _set_cpu0(device_string):
  """Creates a new device string based on `device_string` but using /CPU:0.

  If the device is already on /CPU:0, this is a no-op.

  Args:
    device_string: A device string.

  Returns:
    A device string.
  """
  parsed_device = pydev.DeviceSpec.from_string(device_string)
  parsed_device.device_type = "CPU"
  parsed_device.device_index = 0
  return parsed_device.to_string()


class BaseSaverBuilder(object):
  """Base class for Savers.

  Can be extended to create different Ops.
  """

  class SaveSpec(object):
    """Class used to describe tensor slices that need to be saved."""

    def __init__(self, tensor, slice_spec, name):
      """Creates a `SaveSpec` object.

      Args:
        tensor: the tensor to save or callable that produces a tensor to save.
        slice
_spec: the slice to be saved. See `Variable.SaveSliceInfo`.
        name: the name to save the tensor under.
      """
      self._tensor = tensor
      self.slice_spec = slice_spec
      self.name = name

    @property
    def tensor(self):
      return self._tensor() if callable(self._tensor) else self._tensor

  class SaveableObject(object):
    """Base class for saving and restoring saveable objects."""

    def __
init__(self, op, specs, name):
      """Creates a `SaveableObject` object.

      Args:
        op: the "producer" object that this class wraps; it produces a list of
          tensors to save.  E.g., a "Variable" object saving its backing tensor.
        specs: a list of SaveSpec, each element of which describes one tensor to
          save under this object.
        name: the name to save the object under.
      """
      self.op = op
      self.specs = specs
      self.name = name
      # The device of this saveable. All tensors must be on the same device.
      self.device = specs[0].tensor.device

    def restore(self, restored_tensors, restored_shapes):
      """Restores this object from 'restored_tensors'.

      Args:
        restored_tensors: the tensors that were loaded from a checkpoint
        restored_shapes: the shapes this object should conform to after
          restore, or None.

      Returns:
        An operation that restores the state of the object.

      Raises:
        ValueError: If the object cannot be restored using the provided
          parameters.
      """
      # pylint: disable=unused-argument
      raise ValueError("Calling an abstract method.")

  class VariableSaveable(SaveableObject):
    """SaveableObject implementation that handles Variables."""

    def __init__(self, var, slice_spec, name):
      spec = BaseSaverBuilder.SaveSpec(var, slice_spec, name)
      super(BaseSaverBuilder.VariableSaveable, self).__init__(var, [spec], name)

    def restore(self, restored_tensors, restored_shapes):
      restored_tensor = restored_tensors[0]
      if restored_shapes is not None:
        restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
      return state_ops.assign(
          self.op,
          restored_tensor,
          validate_shape=restored_shapes is None and
          self.op.get_shape().is_fully_defined())

  class ResourceVariableSaveable(SaveableObject):
    """SaveableObject implementation that handles ResourceVariables."""

    def __init__(self, var, slice_spec, name):
      self._var_device = var.device
      self._var_shape = var.shape
      if isinstance(var, ops.Tensor):
        self.handle_op = var.op.inputs[0]
        tensor = var
      elif isinstance(var, resource_variable_ops.ResourceVariable):

        def _read_variable_closure(v):
          def f():
            with ops.device(v.device):
              x = v.read_value()
            with ops.device("/device:CPU:0"):
              return array_ops.identity(x)
          return f

        self.handle_op = var.handle
        tensor = _read_variable_closure(var)
      else:
        raise ValueError(
            "Saveable is neither a resource variable nor a read operation."
            " Got: %s" % repr(var))
      spec = BaseSaverBuilder.SaveSpec(tensor, slice_spec, name)
      super(BaseSaverBuilder.ResourceVariableSaveable, self).__init__(
          var, [spec], name)

    def restore(self, restored_tensors, restored_shapes):
      restored_tensor = restored_tensors[0]
      if restored_shapes is not None:
        restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
      # Copy the restored tensor to the variable's device.
      with ops.device(self._var_device):
        restored_tensor = array_ops.identity(restored_tensor)
      return resource_variable_ops.shape_safe_assign_variable_handle(
          self.handle_op, self._var_shape, restored_tensor)

  def __init__(self, write_version=saver_pb2.SaverDef.V2):
    self._write_version = write_version

  def save_op(self, filename_tensor, saveables):
    """Create an Op to save 'saveables'.

    This is intended to be overridden by subclasses that want to generate
    different Ops.

    Args:
      filename_tensor: String Tensor.
      saveables: A list of BaseSaverBuilder.SaveableObject objects.

    Returns:
      An Operation that save the variables.

    Raises:
      RuntimeError: (implementation detail) if "self._write_version" is an
        unexpected value.
    """
    # pylint: disable=protected-access
    tensor_names = []
    tensors = []
    tensor_slices = []
    for saveable in saveables:
      for spec in saveable.specs:
        tensor_names.append(spec.name)
        tensors.append(spec.tensor)
        tensor_slices.append(spec.slice_spec)
    if self._write_version == saver_pb2.SaverDef.V1:
      return io_ops._save(
          filename=filename_tensor,
          tensor_names=tensor_names,
          tensors=tensors,
          tensor_slices=tensor_slices)
    el