import copy
from typing import Optional
import weakref

from tensorflow.core.protobuf import struct_pb2
from tensorflow.python.distribute import device_util
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import packed_distributed_variable as packed
from tensorflow.python.distribute import reduce_util
from tensorflow.python.distribute import values_util
from tensorflow.python.eager import context
from tensorflow.python.eager import record
from tensorflow.python.framework import composite_tensor
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_conversion_registry
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.framework import type_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.trackable import base as trackable
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.types import core
from tensorflow.python.types import distribute as ds_types
from tensorflow.python.types import trace

def _on_write_update_replica(var, update_fn, value, **kwargs):
  """Updates variables with ON_WRITE synchronization in replica context."""
  if var.aggregation == vs.VariableAggregation.NONE:
    return update_fn(var._get_on_device_or_primary(), value, **kwargs)  # pylint: disable=protected-access

  if not distribute_lib.get_strategy().extended._use_merge_call():  # pylint: disable=protected-access
    # Don't allow MEAN with non float dtype, since it may cause unexpected
    # precision loss. Python3 and NumPy automatically upcast integers to
    # float in division, but we should always preserve the type.
    if var.aggregation == vs.VariableAggregation.MEAN and (
        not var.dtype.is_floating) and tensor_util.is_tf_type(value):
      raise ValueError(
          "Cannot update non-float variables with "
          "tf.VariableAggregation.MEAN aggregation in replica context. "
          "Either change the variable dtype to float or update it in "
          "cross-replica context.")

    aggregated_value = apply_aggregation_replica_context(
        value, var.aggregation, var)
    values_util.mark_as_unsaveable()

    return distribute_lib.get_replica_context()._update(  # pylint: disable=protected-access
        var,
        update_fn,
        args=(aggregated_value,),
        kwargs=kwargs,
        group=True)

  else:

    def merge_fn(strategy, value, **kwargs):
      """Aggregate values and update all variables in cross replica context."""
      # Don't allow MEAN with non float dtype, since it may cause unexpected
      # precision loss. Python3 and NumPy automatically upcast integers to
      # float in division, but we should always preserve the type.
      #
      # Note that to be backward compatible we allow the case when the value
      # is *always* the same on each replica. I.E. value is not a
      # PerReplica. Refer to regroup() to see how values are grouped.
      if var.aggregation == vs.VariableAggregation.MEAN and (
          not var.dtype.is_floating) and isinstance(value, PerReplica):
        raise ValueError(
            "Cannot update non-float variables with "
            "tf.VariableAggregation.MEAN aggregation in replica context. "
            "Either change the variable dtype to float or update it in "
            "cross-replica context.")

      assert strategy == var.distribute_strategy
      v = values_util.apply_aggregation(strategy, value, var.aggregation, var)
      return var._update_cross_replica(update_fn, v, **kwargs)  # pylint: disable=protected-access

    return distribute_lib.get_replica_context().merge_call(
        merge_fn, args=(value,), kwargs=kwargs)

def apply_aggregation_replica_context(value, aggregation, destinations):
  """Aggregate `value` to `destinations` as specified by `aggregation`."""
  # if it is a python literal, return without aggregation
  if isinstance(value, DistributedValues):
    raise TypeError(
        "Cannot use DistributedValues to update variables in replica context.")
  if not tensor_util.is_tf_type(value):
    return value

  if aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
    # Switch to cross-replica context to broadcast
    def merge_fn(strategy, value):
      return strategy.extended.broadcast_to(
          strategy.experimental_local_results(value)[0],
          destinations=destinations)

    return distribute_lib.get_replica_context().merge_call(
        merge_fn, args=(value,))

  else:
    reduce_op = reduce_util.ReduceOp.from_variable_aggregation(aggregation)
    aggregated_value = distribute_lib.get_strategy(  # pylint: disable=protected-access
    ).extended._replica_ctx_all_reduce(reduce_op, value)
    return aggregated_value

class DistributedValues(ds_types.DistributedValues):
  """Base class for representing distributed values."""

  def __init__(self, values):
    """Should only be called by subclass __init__."""
    self._values = tuple(values)

  def _get(self):
    """Returns the value for the current device or raises a ValueError."""
    replica_id = values_util.get_current_replica_id_as_int()
    if replica_id is None:
      return self._get_cross_replica()
    else:
      return self._values[replica_id]

  def _get_cross_replica(self):
    raise NotImplementedError(
        "DistributedValues._get_cross_replica should be implemented by "
        "sub-classes which support cross-replica accesses. "
        f"Type name is {type(self)}"
    )

  def _get_on_device_or_primary(self):
    """Returns value in same replica or device if possible, else the _primary."""
    replica_id = values_util.get_current_replica_id_as_int()
    if replica_id is None:
      # Try to find a value on the current device.
      current_device = device_util.canonicalize(device_util.current())
      for value in self._values:
        if device_util.canonicalize(value.device) == current_device:
          return value
      return self._primary
    else:
      return self._values[replica_id]

  @property
  def _primary(self):
    """Returns a representative component."""
    return self._values[0]

  @property
  def _devices(self):
    return tuple(v.device for v in self._values)

  def __str__(self):
    debug_str = ",\n".join(
        "  %d: %s" % (i, v) for i, v in enumerate(self._values))
    return "%s:{\n%s\n}" % (self.__class__.__name__, debug_str)

  def __repr__(self):
    debug_repr = ",\n".join(
        "  %d: %r" % (i, v) for i, v in enumerate(self._values))
    return "%s:{\n%s\n}" % (self.__class__.__name__, debug_repr)

# NOTE(josh11b,apassos): It would be great if we could inspect the values this was
# initialized with and use that to generate the overloaded operators here.
# Unfortunately, Python's rules for special methods don't allow this, see
# https://docs.python.org/3/reference/datamodel.html#special-method-names
# "if a class defines a method named __getitem__(), and x is an instance of
# this class, then x[i] is roughly equivalent to type(x).__getitem__(x, i)."
# In particular, these special methods don't go through __getattr__, and
# it will only use those methods if they are defined in the class, not the
# object.
class DistributedDelegate(DistributedValues):
  """A map from device to values; acts as the same type as the values."""

  def __getattr__(self, name):
    # The '_use_resource_variables' and the attrs starts with '_self' are used
    # for restoring the saved_model proto, and '_attribute_sentinel' is used for
    # Layer tracking. At the point these attrs are queried, the variable has not
    # been initialized. Thus it should not query those of the underlying
    # components.
    if name.startswith("_self_") or name in ("_use_resource_variables",
                                             "_attribute_sentinel",
                                             "_distributed_container"):
      return super(DistributedDelegate, self).__getattr__(name)

    # This allows copy.copy(DistributedDelegate). When copying an object,
    # copy.copy doesn't invoke its __init__ method, instead it makes a new
    # empty object, then copies the attributes over. copy.copy looks for
    # attributes like "__getstate__" in case the object implements its custom
    # copying. Since DistributedDelegate doesn't have those attributes defined,
    # __getattr__ will be invoked, which tries to access "_values" attributes,
    # but that doesn't exist either because this is an empty object, and again
    # __getattr__ is invoked, leading to an infinite recursion.
    if name == "_values":
      raise AttributeError()

    # TODO(priyag): This needs to be made robust against pitfalls from mix use
    # __getattr__ and @property. See b/120402273.
    return getattr(self._get(), name)

  @property
  def values(self):
    """Returns the per replica values."""
    return self._values

  def _get_as_operand(self):
    """Returns the value for operations for the current device.

    Some implementations, e.g. `TPUMirroredVariable`, are not able to return the
    value type within a replica context. They can, however, return a value that
    can be used by the operations below.
    """
    return self._get()

  # pylint: disable=multiple-statements
  def __add__(self, o):
    return self._get_as_operand() + o

  def __radd__(self, o):
    return o + self._get_as_operand()

  def __sub__(self, o):
    return self._get_as_operand() - o

  def __rsub__(self, o):
    return o - self._get_as_operand()

  def __mul__(self, o):
    return self._get_as_operand() * o

  def __rmul__(self, o):
    return o * self._get_as_operand()

  def __truediv__(self, o):
    return self._get_as_operand() / o

  def __rtruediv__(self, o):
    return o / self._get_as_operand()

  def __floordiv__(self, o):
    return self._get_as_operand() // o

  def __rfloordiv__(self, o):
    return o // self._get_as_operand()

  def __mod__(self, o):
    return self._get_as_operand() % o

  def __rmod__(self, o):
    return o % self._get_as_operand()

  def __lt__(self, o):
    return self._get_as_operand() < o

  def __le__(self, o):
    return self._get_as_operand() <= o

  def __gt__(self, o):
    return self._get_as_operand() > o

  def __ge__(self, o):
    return self._get_as_operand() >= o

  def __and__(self, o):
    return self._get_as_operand() & o

  def __rand__(self, o):
    return o & self._get_as_operand()

  def __or__(self, o):
    return self._get_as_operand() | o

  def __ror__(self, o):
    return o | self._get_as_operand()

  def __xor__(self, o):
    return self._get_as_operand() ^ o

  def __rxor__(self, o):
    return o ^ self._get_as_operand()

  def __getitem__(self, o):
    return self._get_as_operand()[o]

  def __pow__(self, o, modulo=None):
    return pow(self._get_as_operand(), o, modulo)

  def __rpow__(self, o):
    return pow(o, self._get_as_operand())

  def __invert__(self):
    return ~self._get_as_operand()

  def __neg__(self):
    return -self._get_as_operand()

  def __abs__(self):
    return abs(self._get_as_operand())

  def __div__(self, o):
    try:
      return self._get_as_operand().__div__(o)
    except AttributeError:
      # See https://docs.python.org/3/library/constants.html#NotImplemented
      return NotImplemented

  def __rdiv__(self, o):
    try:
      return self._get_as_operand().__rdiv__(o)
    except AttributeError:
      # See https://docs.python.org/3/library/constants.html#NotImplemented
      return NotImplemented

  def __matmul__(self, o):
    try:
      return self._get_as_operand().__matmul__(o)
    except AttributeError:
      # See https://docs.python.org/3/library/constants.html#NotImplemented
      return NotImplemented

  def __rmatmul__(self, o):
    try:
      return self._get_as_operand().__rmatmul__(o)
    except AttributeError:
      # See https://docs.python.org/3/library/constants.html#NotImplemented
      return NotImplemented

  # TODO(josh11b): Even more operator overloads.

class PerReplica(DistributedValues, composite_tensor.CompositeTensor,
                 ds_types.PerReplica):
  """Holds a map from replica to unsynchronized values."""

  @property
  def _type_spec(self):
    return PerReplicaSpec(
        *(type_spec.type_spec_from_value(v) for v in self._values))

  @property
  def values(self):
    """Returns the per replica values."""
    return self._values

def _per_replica_to_tensor(var, dtype=None, name=None, as_ref=False):
  """Converts a `PerReplica` to a `Tensor`."""
  del name
  if dtype is not None and not dtype.is_compatible_with(var.dtype):
    raise ValueError(
        "Incompatible type conversion requested to type {!r} for variable "
        "of type {!r}".format(dtype.name, var.dtype.name))
  if as_ref:
    raise NotImplementedError(
        "PerReplica doesn't support being used as a reference.")
  if (distribute_lib.in_cross_replica_context() or
      not distribute_lib.has_strategy()):
    raise ValueError("It looks like you are using a PerReplica object while "
                     "not inside a replica context, which is not supported. "
                     "Try running your op or function inside a replica context "
                     "by using `strategy.run`")
  else:
    replica_id = values_util.get_current_replica_id_as_int()
    return var.values[replica_id]

# Register a conversion function to provide a useful error message when users
# try to use PerReplica values in the wrong contexts
tensor_conversion_registry.register_tensor_conversion_function(
    PerReplica, _per_replica_to_tensor)

class PerReplicaSpec(type_spec.TypeSpec):
  """Type specification for a `PerReplica`."""

  __slots__ = ["_value_specs"]

  value_type = property(lambda self: PerReplica)

  def __init__(self, *value_specs):
    self._value_specs = tuple(value_specs)

  def _serialize(self):
    return self._value_specs

  @property
  def _component_specs(self):
    return self._value_specs

  def _to_components(self, value):
    replica_context = distribute_lib.get_replica_context()
    if replica_context is not None and replica_context.num_replicas_in_sync > 1:
      raise ValueError(
          "Flattening a PerReplica to components is not supported in replica "
          "context.")
    return value._values  # pylint: disable=protected-access

  def _from_components(self, tensor_list):
    return PerReplica(tensor_list)

nested_structure_coder.register_codec(
    nested_structure_coder.BuiltInTypeSpecCodec(
        PerReplicaSpec, struct_pb2.TypeSpecProto.PER_REPLICA_SPEC
    )
)

# Note that unlike PerReplica, Mirrored values inherit from
# DistributedDelegate and so can be used directly in cross-replica mode.
# TODO(tomhennigan) Should this extend CompositeTensor?
class Mirrored(DistributedDelegate, ds_types.Mirrored):
  """Holds a map from replica to values which are kept in sync."""

  def _get_cross_replica(self):
    return self._get_on_device_or_primary()

  def _as_graph_element(self):
    obj = self._get()
    conv_fn = getattr(obj, "_as_graph_element", None)
    if conv_fn and callable(conv_fn):
      return conv_fn()
    return obj

  def _is_mirrored(self):
    return True

class DistributedVarOp(object):
  """A class that looks like `tf.Operation`."""

  def __init__(self, name, graph, traceback, typ):
    self.name = name
    self.graph = graph
    self.traceback = traceback
    self.type = typ

  def __eq__(self, o):
    if not isinstance(o, self.__class__):
      raise NotImplementedError
    return (self.name == o.name and self.graph == o.graph and
            self.traceback == o.traceback and self.type == o.type)

  def __hash__(self):
    return hash((self.name, self.graph, tuple(self.traceback), self.type))

# TODO(b/209081027): Remove this once Variable is a CompositeTensor.
class DistributedVariableTraceType(trace.TraceType):
  """TraceType of DistributedVariable objects."""

  def __init__(self, distributed_variable):
    self.distributed_variable = distributed_variable
    self.components = (tuple(distributed_variable.shape.as_list()),
                       distributed_variable.dtype)

  def is_subtype_of(self, other):
    return self == other

  def most_specific_common_supertype(self, others):
    return self if all(self == other for other in others) else None

  def placeholder_value(self, placeholder_context=None):
    return self.distributed_variable

  def to_tensors(self, value):
    return []

  def cast(self, value, _):
    return value

  def __hash__(self) -> int:
    return hash(self.components)

  def __eq__(self, other) -> bool:
    if not isinstance(other, DistributedVariableTraceType):
      return False

    return self.components == other.components

class DistributedVariable(DistributedDelegate, variables_lib.Variable,
                          core.Tensor):
  """Holds a map from replica to variables."""

  def __init__(self, strategy, values, aggregation, var_policy=None):
    if (aggregation == variables_lib.VariableAggregation.MEAN and
        not values[0].dtype.is_floating):
      raise ValueError(
          "creating distributed tf.Variable with aggregation=MEAN and a "
          "non-floating dtype is not supported, please use a different "
          "aggregation or dtype")
    self._distribute_strategy = strategy
    self._aggregation = aggregation
    super(DistributedVariable, self).__init__(values)
    self._common_name = self._primary.name.split(":")[0]
    # Use a weakref to make it easy to map from the contained values
    # to the container without introducing a reference cycle.
    for v in values:
      # ResourceVariable is a CompositeTensor. Attributes added to
      # CompositeTensors will get lost through tf.nest packing and unpacking.
      if isinstance(v, composite_tensor.CompositeTensor) and hasattr(
          v, "handle"):
        v.handle._distributed_container = weakref.ref(self)  # pylint: disable=protected-access
      else:
        v._distributed_container = weakref.ref(self)  # pylint: disable=protected-access

    # Packed variable is used to reduce the overhead of function execution.
    # For a DistributedVariable, only one variable handle is captured into a
    # function graph. It's only supported in eager mode.
    if ops.executing_eagerly_outside_functions() and getattr(
        strategy, "_enable_packed_variable_in_eager_mode", False):
      name = "%s/packed/" % self._common_name
      if hasattr(values[0], "_vars"):
        # Handle when the resource variables are "nested" underneath another
        # layer of values, e.g., TPUReplicatedVariable, by packing all them
        # together and pushing the packed var down a level
        # pylint: disable=protected-access
        packed_var = packed.PackedDistributedVariable(
            sum((value._vars for value in values), []), name=name)
        for value in values:
          value._packed_var = packed_var
        self._packed_var = None
        # pylint: enable=protected-access
      else:
        self._packed_var = packed.PackedDistributedVariable(values, name=name)
    else:
      self._packed_var = None

    # tf.keras keeps track of variables initialized using this attribute. When
    # tf.keras gets the default session, it initializes all uninitialized vars.
    # We need to make _keras_initialized a member of DistributedVariable because
    # without this it will use `__getattr__` which will delegate to a component
    # variable.
    self._keras_initialized = False
    # Typically, a `DistributedVariable`'s initializer is composed of the
    # initializers of the components variables. However, in some cases, such as
    # when restoring from a checkpoint, we may set the _initializer_op
    # property on the entire `DistributedVariable`.
    self._initializer_op = None
    # Set a VariablePolicy which decides how we replicate/aggregate the given
    # variable.
    self._policy = var_policy

  def __deepcopy__(self, memo):
    """Perform a deepcopy of the `DistributedVariable`.

    Unlike the deepcopy of a regular tf.Variable, this keeps the original
    strategy and devices of the `DistributedVariable`.  To avoid confusion
    with the behavior of deepcopy on a regular `Variable` (which does
    copy into new devices), we only allow a deepcopy of a `DistributedVariable`
    within its originating strategy scope.

    Args:
      memo: The memoization object for `deepcopy`.

    Returns:
      A deep copy of the current `DistributedVariable`.

    Raises:
      RuntimeError: If trying to deepcopy into a different strategy.
    """
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      new_values = []

      for value in self._values:
        with ops.device(value.device):
          new_values.append(copy.deepcopy(value, memo))

    copied_variable = type(self)(
        strategy=self._distribute_strategy,
        values=new_values,
        aggregation=self._aggregation,
        var_policy=copy.deepcopy(self._policy, memo))

    memo[id(self)] = copied_variable

    return copied_variable

  def _use_packed_variable(self):
    # Don't use packed variable when under a SaveContext to avoid explicit
    # device placement on variable consuming ops.
    return self._packed_var is not None and (
        not values_util.is_saving_non_distributed())

  def is_initialized(self, name=None):
    """Identifies if all the component variables are initialized.

    Args:
      name: Name of the final `logical_and` op.

    Returns:
      The op that evaluates to True or False depending on if all the
      component variables are initialized.
    """
    if values_util.is_saving_non_distributed():
      return self._primary.is_initialized()
    if self._use_packed_variable():
      return self._packed_var.is_initialized()
    result = self._primary.is_initialized()
    # We iterate through the list of values except the last one to allow us to
    # name the final `logical_and` op the same name that is passed by the user
    # to the `is_initialized` op. For distributed variables, the
    # `is_initialized` op is a `logical_and` op.
    for v in self._values[1:-1]:
      result = math_ops.logical_and(result, v.is_initialized())
    result = math_ops.logical_and(
        result, self._values[-1].is_initialized(), name=name)
    return result

  @property
  def initializer(self):
    if values_util.is_saving_non_distributed():
      return self._primary.initializer
    if self._initializer_op:
      init_op = self._initializer_op
    else:
      # return grouped ops of all the var initializations of component values of
      # the mirrored variable
      init_op = control_flow_ops.group(
          tuple(v.initializer for v in self._values))
    return init_op

  def initialized_value(self):
    return self._get_on_device_or_primary().initialized_value()

  def _is_mirrored(self):
    return (self._policy is not None) and (self._policy._is_mirrored())  # pylint: disable=protected-access

  @property
  def initial_value(self):
    return self._get_on_device_or_primary().initial_value

  @property
  def constraint(self):
    return self._primary.constraint

  @property
  def graph(self):
    return self._primary.graph

  @property
  def _shared_name(self):
    return self._common_name

  @property
  def _unique_id(self):
    return self._primary._unique_id  # pylint: disable=protected-access

  @property
  def _graph_key(self):
    """Lets Optimizers know which graph this variable is from."""
    return self._primary._graph_key  # pylint: disable=protected-access

  @property
  def name(self):
    return self._primary.name

  @property
  def dtype(self):
    return self._primary.dtype

  @property
  def shape(self):
    return self._primary.shape

  @property
  def synchronization(self):
    return self._primary.synchronization

  @property
  def aggregation(self):
    return self._aggregation

  @property
  def _packed_variable(self):
    if self._use_packed_variable():
      return self._packed_var
    return None

  @property
  def handle(self):
    if values_util.is_saving_non_distributed():
      return self._primary.handle
    replica_id = values_util.get_current_replica_id_as_int()
    if replica_id is None:
      raise ValueError(
          "DistributedVariable.handle is not available outside the replica "
          "context or a `tf.distribute.Strategy.update()` call.")
    else:
      if self._use_packed_variable():
        return self._packed_var.handle
      return self._values[replica_id].handle

  def eval(self, session=None):
    return self._get_on_device_or_primary().eval(session)

  @property
  def _save_slice_info(self):
    return self._primary._save_slice_info  # pylint: disable=protected-access

  def _get_save_slice_info(self):
    return self._primary._get_save_slice_info()  # pylint: disable=protected-access

  def _set_save_slice_info(self, save_slice_info):
    for v in self._values:
      v._set_save_slice_info(save_slice_info)  # pylint: disable=protected-access

  @property
  def device(self):
    return self._get_on_device_or_primary().device

  @property
  def trainable(self):
    return self._primary.trainable

  @property
  def distribute_strategy(self):
    return self._distribute_strategy

  def get_shape(self) -> tensor_shape.TensorShape:
    return self._primary.get_shape()

  def to_proto(self, export_scope=None):
    return self._primary.to_proto(export_scope=export_scope)

  @property
  def op(self) -> ops.Operation:
    if values_util.is_saving_non_distributed():
      return self._primary.op
    # We want cross-replica code that does some var.op.X calls
    # to work (even if the current device isn't in self._devices), but
    # other uses of var.op in a cross-replica context to fail.
    if distribute_lib.in_cross_replica_context():
      return DistributedVarOp(self._primary.op.name, self._primary.op.graph,
                              self._primary.op.traceback, self._primary.op.type)
    return self._get().op

  @property
  def _in_graph_mode(self):
    return self._primary._in_graph_mode  # pylint: disable=protected-access

  def _get_replica(self, replica_id):
    """Returns the value on a device with the given replica_id."""
    value = self._values[replica_id]
    if self._use_packed_variable():
      return self._packed_var.on_device(value.device)
    else:
      return value

  def _get(self):
    """Returns the value for the current device or raises a ValueError."""
    if values_util.is_saving_non_distributed():
      return self._primary
    replica_id = values_util.get_current_replica_id_as_int()
    if replica_id is None:
      return self._get_cross_replica()
    else:
      return self._get_replica(replica_id)

  def _get_on_device_or_primary(self):
    """Returns value in same replica or device if possible, else the _primary."""
    if values_util.is_saving_non_distributed():
      return self._primary
    replica_id = values_util.get_current_replica_id_as_int()
    if replica_id is None:
      # Try to find a value on the current device.
      current_device = device_util.canonicalize(device_util.current())
      for i, value in enumerate(self._values):
        if device_util.canonicalize(value.device) == current_device:
          return self._get_replica(i)
      return self._get_replica(0)
    else:
      return self._get_replica(replica_id)

  def read_value(self):
    if values_util.is_saving_non_distributed():
      return self._primary.read_value()
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      return array_ops.identity(self._get())

  def value(self):
    if values_util.is_saving_non_distributed():
      return self._primary.value()
    if self._policy:
      return self._policy.value(self)
    return self._get_on_device_or_primary().value()

  def numpy(self):
    if context.executing_eagerly():
      return self.read_value().numpy()
    else:
      raise NotImplementedError("DistributedVariable.numpy() is only available "
                                "when eager execution is enabled.")

  def assign_sub(self, value, use_locking=False, name=None, read_value=True):
    if values_util.is_saving_non_distributed():
      return self._primary.assign_sub(value, use_locking, name, read_value)
    if self._policy:
      return self._policy.assign_sub(
          self,
          value,
          use_locking=use_locking,
          name=name,
          read_value=read_value)
    return values_util.on_write_assign_sub(
        self, value, use_locking=use_locking, name=name, read_value=read_value)

  def assign_add(self, value, use_locking=False, name=None, read_value=True):
    if values_util.is_saving_non_distributed():
      return self._primary.assign_add(value, use_locking, name, read_value)
    if self._policy:
      return self._policy.assign_add(
          self,
          value,
          use_locking=use_locking,
          name=name,
          read_value=read_value)
    return values_util.on_write_assign_add(
        self, value, use_locking=use_locking, name=name, read_value=read_value)

  def assign(self, value, use_locking=False, name=None, read_value=True):
    if values_util.is_saving_non_distributed():
      return self._primary.assign(value, use_locking, name, read_value)
    if self._policy:
      return self._policy.assign(
          self,
          value,
          use_locking=use_locking,
          name=name,
          read_value=read_value)
    return values_util.on_write_assign(
        self, value, use_locking=use_locking, name=name, read_value=read_value)

  def scatter_sub(self, sparse_delta, use_locking=False, name=None):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_sub(sparse_delta, use_locking, name)
    if self._policy:
      return self._policy.scatter_sub(
          self, sparse_delta, use_locking=use_locking, name=name)
    return values_util.scatter_sub(
        self, sparse_delta, use_locking=use_locking, name=name)

  def scatter_add(self, sparse_delta, use_locking=False, name=None):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_add(sparse_delta, use_locking, name)
    if self._policy:
      return self._policy.scatter_add(
          self, sparse_delta, use_locking=use_locking, name=name)
    return values_util.scatter_add(
        self, sparse_delta, use_locking=use_locking, name=name)

  def scatter_mul(self, sparse_delta, use_locking=False, name=None):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_mul(sparse_delta, use_locking, name)
    if self._policy:
      return self._policy.scatter_mul(
          self, sparse_delta, use_locking=use_locking, name=name)
    return values_util.scatter_mul(
        self, sparse_delta, use_locking=use_locking, name=name)

  def scatter_div(self, sparse_delta, use_locking=False, name=None):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_div(sparse_delta, use_locking, name)
    if self._policy:
      return self._policy.scatter_div(
          self, sparse_delta, use_locking=use_locking, name=name)
    return values_util.scatter_div(
        self, sparse_delta, use_locking=use_locking, name=name)

  def scatter_min(self, sparse_delta, use_locking=False, name=None):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_min(sparse_delta, use_locking, name)
    if self._policy:
      return self._policy.scatter_min(
          self, sparse_delta, use_locking=use_locking, name=name)
    return values_util.scatter_min(
        self, sparse_delta, use_locking=use_locking, name=name)

  def scatter_max(self, sparse_delta, use_locking=False, name=None):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_max(sparse_delta, use_locking, name)
    if self._policy:
      return self._policy.scatter_max(
          self, sparse_delta, use_locking=use_locking, name=name)
    return values_util.scatter_max(
        self, sparse_delta, use_locking=use_locking, name=name)

  def scatter_update(self, sparse_delta, use_locking=False, name=None):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_update(sparse_delta, use_locking, name)
    if self._policy:
      return self._policy.scatter_update(
          self, sparse_delta, use_locking=use_locking, name=name)
    return values_util.scatter_update(
        self, sparse_delta, use_locking=use_locking, name=name)

  def __tf_tracing_type__(self, _):
    return DistributedVariableTraceType(self)

  def _gather_saveables_for_checkpoint(self):
    """Overrides Trackable method.

    This allows both name-based and object-based save and restore of
    DistributedVariables.

    Returns:
      A dictionary mapping attribute names to `SaveableObject` factories.
    """

    def _saveable_factory(name=self._common_name):
      return _DistributedVariableSaveable(self, self._primary, name)

    return {trackable.VARIABLE_VALUE_KEY: _saveable_factory}

  def _as_graph_element(self):
    if values_util.is_saving_non_distributed():
      return self._primary._as_graph_element()  # pylint: disable=protected-access
    if self._policy:
      return self._policy._as_graph_element(self)  # pylint: disable=protected-access

    raise NotImplementedError(
        "DistributedVariable._as_graph_element requires a valid "
        "VariablePolicy. Please set the policy via the `var_policy` argument "
        "in the constructor, or override this method in sub-classes which "
        "support cross-replica accesses. "
        f"Type name is {type(self)}"
    )

  def _get_cross_replica(self):
    if values_util.is_saving_non_distributed():
      return self._primary
    if self._policy:
      return self._policy._get_cross_replica(self)  # pylint: disable=protected-access

    raise NotImplementedError(
        "DistributedVariable._get_cross_replica requires a valid "
        "VariablePolicy. Please set the policy via the `var_policy` argument "
        "in the constructor, or override this method in sub-classes which "
        "support cross-replica accesses. "
        f"Type name is {type(self)}"
    )

  def _update_cross_replica(self, update_fn, value, **kwargs):
    """Applies updates across replicas.

    Args:
      update_fn: A callable to pass to `strategy.extended.update` to update the
        variable. It should has the same signature as `Variable.assign()`.
      value: value to be passed to `update_fn`.
      **kwargs: remaining arguments to `update_fn`.

    Returns:
      Updated variable or `tf.Operation`.
    """
    values_util.mark_as_unsaveable()
    return self.distribute_strategy.extended.update(
        self, update_fn, args=(value,), kwargs=kwargs, group=True)

  def _update_replica(self, update_fn, value, **kwargs):
    """Applies updates in one replica.

    Args:
      update_fn: A callable to update the variable. It should has the same
        signature as `Variable.assign()`.
      value: value to be passed to `update_fn`.
      **kwargs: remaining arguments to `update_fn`.

    Returns:
      Updated variable or `tf.Operation`.
    """
    if self._policy:
      return self._policy._update_replica(self, update_fn, value, **kwargs)  # pylint: disable=protected-access
    raise NotImplementedError(
        "DistributedVariable._update_replica requires a valid VariablePolicy. "
        "Please set the policy via the `var_policy` argument in the "
        "constructor, or override this method in sub-classes which support "
        "cross-replica accesses. "
        f"Type name is {type(self)}"
    )

  def _update(self, update_fn, value, **kwargs):
    """Applies updates depending on the context.

    The method calls `_update_replica` in replica context,
    `_update_cross_replica` in cross replica context, and `update_fn` in update
    context.

    If `read_value` is True, the method returns the updated Variable. If
    `read_value` is False, the method returns the update `tf.Operation`.

    Args:
      update_fn: A callable to pass to `strategy.extended.update` to update the
        variable. It should have the same signature as `Variable.assign()`.
      value: value to be passed to `update_fn`.
      **kwargs: keyword arguments to `update_fn`.

    Returns:
      Updated variable or `tf.Operation`.

    """
    if values_util.is_saving_non_distributed():
      return update_fn(self._primary, value, **kwargs)
    with distribute_lib.enter_or_assert_strategy(self.distribute_strategy):
      if distribute_lib.in_cross_replica_context():
        update_replica_id = distribute_lib.get_update_replica_id()
        if update_replica_id is not None:
          replica_value = self._get_replica(update_replica_id)
          return update_fn(replica_value, value, **kwargs)
        return self._update_cross_replica(update_fn, value, **kwargs)
      else:
        values_util.assert_replica_context(self.distribute_strategy)
        return self._update_replica(update_fn, value, **kwargs)

  def _should_act_as_resource_variable(self):
    """Pass resource_variable_ops.is_resource_variable check."""
    pass

  def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
    """Converts a variable to a tensor."""
    if values_util.is_saving_non_distributed():
      return ops.convert_to_tensor(
          self._primary, dtype=dtype, name=name, as_ref=as_ref)
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      return ops.convert_to_tensor(
          self._get(), dtype=dtype, name=name, as_ref=as_ref)

  def __tf_tensor__(self,
                    dtype: Optional[dtypes.DType] = None,
                    name: Optional[str] = None) -> tensor_lib.Tensor:
    return self._dense_var_to_tensor(dtype, name)

  def _export_to_saved_model_graph(self,
                                   object_map=None,
                                   tensor_map=None,
                                   options=None,
                                   **kwargs):
    # Initialize for self._primary first, so that obj_map[self._primary] and
    # resource_map[self._primary.handle] contain mapped values.
    resource_list = self._primary._export_to_saved_model_graph(  # pylint:disable=protected-access
        object_map=object_map,
        tensor_map=tensor_map,
        options=options,
        **kwargs)
    for v in [v for v in self._values if v != self._primary]:
      if (options.experimental_variable_policy  # pylint:disable=protected-access
          ._expand_distributed_variables()):
        resource_list.extend(
            v._export_to_saved_model_graph(  # pylint:disable=protected-access
                object_map=object_map,
                tensor_map=tensor_map,
                options=options,
                **kwargs))  # pylint:disable=protected-access
      else:
        object_map[v] = object_map[self._primary]
        tensor_map[v.handle] = tensor_map[self._primary.handle]
        resource_list.append(v.handle)
    object_map[self] = object_map[self._primary]
    tensor_map[self] = tensor_map[self._primary.handle]
    resource_list.append(self)
    if self._packed_var is not None:
      tensor_map[self._packed_var.packed_handle] = tensor_map[
          self._primary.handle]
      resource_list.append(self._packed_var.packed_handle)
    return resource_list

  def _copy_trackable_to_cpu(self, object_map):
    """For implementing `Trackable`."""
    if self not in object_map:
      # If not populated, initialize the cpu copy first.
      op_device = pydev.DeviceSpec.from_string(self.device).replace(
          device_type="CPU", device_index=0).to_string()
      with ops.device(op_device):
        new_var = resource_variable_ops.UninitializedVariable(
            trainable=self.trainable,
            shape=self.shape,
            dtype=self.dtype,
            name=self._shared_name,
            distribute_strategy=self._distribute_strategy,
            aggregation=self._aggregation)  # pylint: disable=protected-access
      object_map[self] = new_var

    # Then copy value of self to the copy.
    destination_var = object_map[self]
    with ops.device(destination_var.device):
      destination_var.assign(self.read_value())

  def _write_object_proto(self, proto, options):
    """Update a SavedObject proto for the caller.

    If a DistributedVariable object supports this method, it will be called when
    saving with a pre-built `SavedObject` proto representing the object, plus an
    instance of `SaveOptions`. This method is then free to modify that proto
    instance.

    `DistributedVariable` with `AUTO` or `ON_WRITE` synchronization optionally
    write out information about their components to the
    `experimental_distributed_variable_components` field of a
    `SavedVariable` (depending on the `SaveOptions` variable policy).

    Args:
      proto: A pre-built `SavedObject` proto for this object. It is assumed this
        will be a `SavedVariable` instance.
      options: A `SaveOptions` instance.
    """
    resource_variable_ops.write_object_proto_for_resource_variable(
        self, proto, options)
    # Set protos in the saved model such that distributed variables are
    # correctly restored on COMPOSITE devices (otherwise, task:0/TPU:0).
    values_util.write_object_proto(self, proto, options)

  @property
  def is_distributed_variable(self):
    return True

  def __tf_experimental_restore_capture__(
      self, concrete_function, internal_capture):
    graph = concrete_function.graph
    # Add given distributed variable to captures with given placeholder.
    graph.replace_capture(self, internal_capture)
    record.record_operation(
        "captured_value", [internal_capture], [self],
        backward_function=lambda x: [x],
        forward_function=lambda x: [x])
    return self

# We extend from `saveable_object.SaveableObject` instead of
# `saveable_object_util.ResourceVariableSaveable` since we need to read the
# value of ONREAD variables when saving. `SaveableObject` provides a way to
# specify the function to run to get the value of the variable or tensor at
# saving time. We can use this for both ON_READ and ON_WRITE variables.
# TODO(b/164586507): Consolidate ON_WRITE and ON_READ saving/restoring logic
# if possible.
class _DistributedVariableSaveable(saveable_object.SaveableObject):
  """Class for defining how to restore a DistributedVariable."""

  def __init__(self, distributed_variable, primary_variable, name):
    self._distributed_variable = distributed_variable
    if not self._distributed_variable._policy:
      raise ValueError(
          "The VariablePolicy of the argument `distributed_variable` must be "
          "set to create a _DistributedVariableSaveable. Please set it via "
          "the `var_policy` argument in the constructor of DistributedVariable."
      )
    tensor, spec = distributed_variable._policy.get_saveable(
        distributed_variable, primary_variable, name)
    super(_DistributedVariableSaveable, self).__init__(tensor, spec, name)

  def restore(self, restored_tensors, restored_shapes):
    """Restore the same value into all variables."""
    tensor, = restored_tensors
    return self._distributed_variable._policy.get_restore_ops(  # pylint: disable=protected-access
        self._distributed_variable, tensor)

class _MirroredSaveable(saveable_object.SaveableObject):
  """Class for defining how to restore a MirroredVariable."""

  def __init__(self, mirrored_variable, primary_variable, name):
    self._mirrored_variable = mirrored_variable
    tensor, spec = values_util.get_on_write_saveable(self._mirrored_variable,
                                                     primary_variable, name)
    super(_MirroredSaveable, self).__init__(tensor, spec, name)

  def restore(self, restored_tensors, restored_shapes):
    """Restore the same value into all variables."""
    tensor, = restored_tensors
    return values_util.get_on_write_restore_ops(self._mirrored_variable, tensor)

class MirroredVariable(DistributedVariable, Mirrored):
  """Holds a map from replica to variables whose values are kept in sync."""

  def _is_mirrored(self):
    return Mirrored._is_mirrored(self)  # Use correct parent class.

  def _update_replica(self, update_fn, value, **kwargs):
    return _on_write_update_replica(self, update_fn, value, **kwargs)

  def scatter_min(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_min(*args, **kwargs)
    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
        self._aggregation != vs.VariableAggregation.NONE):
      raise NotImplementedError(
          values_util.scatter_error_msg.format(
              op_name="scatter_min", aggregation=self._aggregation))
    return super(MirroredVariable, self).scatter_min(*args, **kwargs)

  def scatter_max(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_max(*args, **kwargs)
    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
        self._aggregation != vs.VariableAggregation.NONE):
      raise NotImplementedError(
          values_util.scatter_error_msg.format(
              op_name="scatter_max", aggregation=self._aggregation))
    return super(MirroredVariable, self).scatter_max(*args, **kwargs)

  def scatter_update(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_update(*args, **kwargs)
    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
        self._aggregation != vs.VariableAggregation.NONE):
      raise NotImplementedError(
          values_util.scatter_error_msg.format(
              op_name="scatter_update", aggregation=self._aggregation))
    return super(MirroredVariable, self).scatter_update(*args, **kwargs)

  def _get_cross_replica(self):
    # Return identity, to avoid directly exposing the variable to the user and
    # allowing it to be modified by mistake.
    return array_ops.identity(Mirrored._get_cross_replica(self))

  def _as_graph_element(self):
    return self._get_on_device_or_primary()._as_graph_element()  # pylint: disable=protected-access

  def _gather_saveables_for_checkpoint(self):
    """Overrides Trackable method.

    This allows both name-based and object-based save and restore of
    MirroredVariables.

    Returns:
      A dictionary mapping attribute names to `SaveableObject` factories.
    """

    def _saveable_factory(name=self._common_name):
      return _MirroredSaveable(self, self._primary, name)

    return {trackable.VARIABLE_VALUE_KEY: _saveable_factory}

  def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
    """Converts a variable to a tensor."""
    # TODO(b/154017756): Make _dense_var_to_tensor consistent between ON_READ
    # and ON_WRITE.
    # Try to avoid assignments to and other mutations of MirroredVariable
    # state except through a DistributionStrategy.extended.update() or any of
    # the `assign*` and `scatter*` calls.
    if as_ref:
      # A TF 1.x case where the variable is a boolean variable and used like:
      # tf.cond(v, true_fn, false_fn).
      raise ValueError(
          "You may be using variable created under distribute strategy in TF "
          "1.x control flows. Try explicitly converting the variable to Tensor "
          "using variable.read_value(), or switch to TF 2.x.")
    return ops.convert_to_tensor(
        self._get(), dtype=dtype, name=name, as_ref=as_ref)

class _SyncOnReadSaveable(saveable_object.SaveableObject):
  """Class for defining how to restore a SyncOnReadVariable."""

  def __init__(self, sync_on_read_variable, name):
    self._sync_on_read_variable = sync_on_read_variable
    tensor, spec = values_util.get_on_read_saveable(
        sync_on_read_variable, sync_on_read_variable._primary, name)

    super(_SyncOnReadSaveable, self).__init__(tensor, spec, name)

  def restore(self, restored_tensors, restored_shapes):
    """Restore the same value into all variables."""
    tensor, = restored_tensors
    return values_util.get_on_read_restore_ops(
        self._sync_on_read_variable, tensor,
        self._sync_on_read_variable.aggregation)

class SyncOnReadVariable(DistributedVariable):
  """Holds a map from replica to variables whose values are reduced on save."""

  def _update_replica(self, update_fn, value, **kwargs):
    return update_fn(self._get_on_device_or_primary(), value, **kwargs)

  def _get(self):
    """Returns the value of SyncOnReadVariable based on surrounding context.

    If called under a non-default replica-context, returns the corresponding
    variable on that replica.
    If called under default replica-context or cross-replica context, returns
    the synced value.
    """
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      return super(SyncOnReadVariable, self)._get()

  # TODO(b/154017756): Make assign behaivor in cross replica context consistent
  # with MirroredVariable.
  def assign_sub(self, value, use_locking=False, name=None, read_value=True):
    if values_util.is_saving_non_distributed():
      return self._primary.assign_sub(value, use_locking, name, read_value)
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        values_util.mark_as_unsaveable()
        return values_util.on_read_assign_sub_cross_replica(
            self, value, read_value=read_value)
      else:
        return super(SyncOnReadVariable,
                     self).assign_sub(value, use_locking, name, read_value)

  def assign_add(self, value, use_locking=False, name=None, read_value=True):
    if values_util.is_saving_non_distributed():
      return self._primary.assign_add(value, use_locking, name, read_value)
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        values_util.mark_as_unsaveable()
        return values_util.on_read_assign_add_cross_replica(
            self, value, read_value=read_value)
      else:
        return super(SyncOnReadVariable,
                     self).assign_add(value, use_locking, name, read_value)

  def assign(self, value, use_locking=False, name=None, read_value=True):
    if values_util.is_saving_non_distributed():
      return self._primary.assign(value, use_locking, name, read_value)
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        values_util.mark_as_unsaveable()
        return values_util.on_read_assign_cross_replica(
            self, value, read_value=read_value)
      else:
        return super(SyncOnReadVariable, self).assign(value, use_locking, name,
                                                      read_value)

  def _scatter_not_implemented(self, method):
    raise NotImplementedError(
        f"Variables with `synchronization=ON_READ` doesn't support `{method}`")

  def scatter_sub(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_sub(*args, **kwargs)
    self._scatter_not_implemented("scatter_sub")

  def scatter_add(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_add(*args, **kwargs)
    self._scatter_not_implemented("scatter_add")

  def scatter_mul(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_mul(*args, **kwargs)
    self._scatter_not_implemented("scatter_mul")

  def scatter_div(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_div(*args, **kwargs)
    self._scatter_not_implemented("scatter_div")

  def scatter_min(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_min(*args, **kwargs)
    self._scatter_not_implemented("scatter_min")

  def scatter_max(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_max(*args, **kwargs)
    self._scatter_not_implemented("scatter_max")

  def scatter_update(self, *args, **kwargs):
    if values_util.is_saving_non_distributed():
      return self._primary.scatter_update(*args, **kwargs)
    self._scatter_not_implemented("scatter_update")

  def value(self):
    if distribute_lib.in_variable_sync_on_read_context():
      raise NotImplementedError(
          "call `variable.value()` inside variable_sync_on_read_context is not "
          "supported")
    if values_util.is_saving_non_distributed():
      return self._primary.value()
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
          return self._get_replica(0).value()
        return self._get_cross_replica()
      else:
        # _get_on_device_or_primary() returns a Variable.
        return self._get_on_device_or_primary().value()

  def read_value(self):
    if distribute_lib.in_variable_sync_on_read_context():
      raise NotImplementedError(
          "call `variable.read_value()` inside variable_sync_on_read_context is"
          " not supported")
    return super().read_value()

  def _get_cross_replica(self):
    if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
      # Consider returning a tensor value here to make the return value of
      # _get_cross_replica consistent.
      return self._get_replica(0)
    if self._aggregation == vs.VariableAggregation.SUM:
      values_util.mark_as_unsaveable()
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      return self._distribute_strategy.reduce(
          reduce_util.ReduceOp.from_variable_aggregation(self._aggregation),
          self,
          axis=None)

  def _as_graph_element(self):
    if values_util.is_saving_non_distributed():
      return self._primary._as_graph_element()  # pylint: disable=protected-access
    # pylint: disable=protected-access
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      if distribute_lib.in_cross_replica_context():
        return ops.convert_to_tensor(self._get_cross_replica())
    return self._get()._as_graph_element()

  def _gather_saveables_for_checkpoint(self):
    """Overrides Trackable method.

    This allows both name-based and object-based save and restore of
    `SyncOnReadVariable`s.

    Returns:
      A dictionary mapping attribute names to `SaveableObject` factories.
    """

    def _saveable_factory(name=self._common_name):
      return _SyncOnReadSaveable(self, name)

    return {trackable.VARIABLE_VALUE_KEY: _saveable_factory}

  def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
    """Converts a SyncOnReadVariable to a tensor."""
    if values_util.is_saving_non_distributed():
      return ops.convert_to_tensor(
          self._primary, dtype=dtype, name=name, as_ref=as_ref)
    with distribute_lib.enter_or_assert_strategy(self._distribute_strategy):
      replica_context = distribute_lib.get_replica_context()
      if (replica_context is not None and
          distribute_lib.in_variable_sync_on_read_context()):
        if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
          return ops.convert_to_tensor(
              self._get_replica(0), dtype=dtype, name=name, as_ref=as_ref)
        if self._aggregation == vs.VariableAggregation.SUM:
          values_util.mark_as_unsaveable()
        # pylint: disable=protected-access
        reduced = (
            replica_context.strategy.extended._replica_ctx_all_reduce(
                reduce_util.ReduceOp.from_variable_aggregation(
                    self._aggregation),
                self._get().read_value()))
        return ops.convert_to_tensor(
            reduced, dtype=dtype, name=name, as_ref=as_ref)

      return ops.convert_to_tensor(
          self._get(), dtype=dtype, name=name, as_ref=as_ref)

# Register a conversion functions which reads the value of the variable,
# allowing instances of the class to be used as tensors.
# DistributedVariable
def _tensor_conversion_distributed_var(var,
                                       dtype=None,
                                       name=None,
                                       as_ref=False):
  return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access

tensor_conversion_registry.register_tensor_conversion_function(
    DistributedVariable, _tensor_conversion_distributed_var)

# MirroredVariables
def _tensor_conversion_mirrored(var, dtype=None, name=None, as_ref=False):
  return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access

tensor_conversion_registry.register_tensor_conversion_function(
    MirroredVariable, _tensor_conversion_mirrored)

# Mirrored Values
def _tensor_conversion_mirrored_val(value, dtype=None, name=None, as_ref=False):
  return ops.convert_to_tensor(
      value._get(), dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access

tensor_conversion_registry.register_tensor_conversion_function(
    Mirrored, _tensor_conversion_mirrored_val)

# SyncOnReadVariables
def _tensor_conversion_sync_on_read(var, dtype=None, name=None, as_ref=False):
  return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref)  # pylint: disable=protected-access

tensor_conversion_registry.register_tensor_conversion_function(
    SyncOnReadVariable, _tensor_conversion_sync_on_read)

class VariablePolicy(object):
  """Policy defining synchronization and aggregation of a distributed variable.

  Given `synchronization` and `aggregation` parameters set on a `tf.Variable`
  during variable creation within `tf.distribute` scope, `tf.distribute` creates
  an appropriate policy object and assigns it to the distributed variable. All
  variable operations are delegated to the respective policy object.
  """

  def __init__(self, aggregation):
    self._aggregation = aggregation

  def value(self):
    raise NotImplementedError(
        "VariablePolicy.value should be overridden by sub-classes. "
        f"Type name is {type(self)}"
    )

  def _is_mirrored(self):
    raise NotImplementedError(
        "VariablePolicy._is_mirrored should be overridden by sub-classes. "
        f"Type name is {type(self)}"
    )

  def _as_graph_element(self, _):
    raise NotImplementedError(
        "VariablePolicy._as_graph_element should be overridden by sub-classes. "
        f"Type name is {type(self)}"
    )

  def _get_cross_replica(self, var):
    raise NotImplementedError(
        "VariablePolicy._get_cross_replica should be overridden by"
        f" sub-classes. Type name is {type(self)}"
    )

  def _update_replica(self, var, update_fn, value, **kwargs):
    raise NotImplementedError(
        "VariablePolicy._update_replica should be overridden by sub-classes. "
        f"Type name is {type(self)}"
    )

class OnReadPolicy(VariablePolicy):
  """Policy defined for `tf.VariableSynchronization.ON_READ` synchronization.

  This policy is created when `synchronization` is set to
  `tf.VariableSynchronization.ON_READ` and `aggregation` is set to any of the
  values allowed by the `tf.VariableAggregation` enum such as `NONE`, `SUM`,
  `MEAN` or `ONLY_FIRST_REPLICA`when creating a `tf.Variable` in `tf.distribute`
  scope.
  """

  def _is_mirrored(self):
    return False

  def value(self, var):
    with distribute_lib.enter_or_assert_strategy(var.distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
          return var._get_replica(0).value()  # pylint: disable=protected-access
        return var._get_cross_replica()  # pylint: disable=protected-access
      else:
        return var._get_on_device_or_primary().value()  # pylint: disable=protected-access

  def _as_graph_element(self, var):
    with distribute_lib.enter_or_assert_strategy(var.distribute_strategy):
      if distribute_lib.in_cross_replica_context():
        return ops.convert_to_tensor(var._get_cross_replica())  # pylint: disable=protected-access
    return var._get()._as_graph_element()  # pylint: disable=protected-access

  def _get_cross_replica(self, var):
    if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA:
      return var._get_replica(0)  # pylint: disable=protected-access
    if self._aggregation == vs.VariableAggregation.SUM:
      values_util.mark_as_unsaveable()
    with distribute_lib.enter_or_assert_strategy(var.distribute_strategy):
      return var.distribute_strategy.reduce(
          reduce_util.ReduceOp.from_variable_aggregation(self._aggregation),
          var,
          axis=None)

  def _update_replica(self, var, update_fn, value, **kwargs):
    return update_fn(var._get_on_device_or_primary(), value, **kwargs)  # pylint: disable=protected-access

  def _scatter_not_implemented(self, method):
    raise NotImplementedError(f"ON_READ variables doesn't support `{method}` "
                              "in cross replica context")

  def assign_sub(self,
                 var,
                 value,
                 use_locking=False,
                 name=None,
                 read_value=True):
    """Subtracts a value from this variable."""
    with distribute_lib.enter_or_assert_strategy(var.distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        values_util.mark_as_unsaveable()
        return values_util.on_read_assign_sub_cross_replica(
            var, value, read_value=read_value)
      else:
        return values_util.on_write_assign_sub(
            var,
            value,
            use_locking=use_locking,
            name=name,
            read_value=read_value)

  def assign_add(self,
                 var,
                 value,
                 use_locking=False,
                 name=None,
                 read_value=True):
    """Adds a value to this variable."""
    with distribute_lib.enter_or_assert_strategy(var.distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        values_util.mark_as_unsaveable()
        return values_util.on_read_assign_add_cross_replica(
            var, value, read_value=read_value)
      else:
        return values_util.on_write_assign_add(
            var,
            value,
            use_locking=use_locking,
            name=name,
            read_value=read_value)

  def assign(self, var, value, use_locking=False, name=None, read_value=True):
    with distribute_lib.enter_or_assert_strategy(var.distribute_strategy):
      if (distribute_lib.in_cross_replica_context() and
          not values_util.in_replica_update_context()):
        values_util.mark_as_unsaveable()
        return values_util.on_read_assign_cross_replica(
            var, value, read_value=read_value)
      else:
        return values_util.on_write_assign(
            var,
            value,
            use_locking=use_locking,
            name=name,
            read_value=read_value)

  def scatter_sub(self, *args, **kwargs):
    del args, kwargs
    self._scatter_not_implemented("scatter_sub")

  def scatter_add(self, *args, **kwargs):
    del args, kwargs
    self._scatter_not_implemented("scatter_add")

  def scatter_mul(self, *args, **kwargs):
    del args, kwargs
    self._scatter_not_implemented("scatter_mul")

  def scatter_div(self, *args, **kwargs):
    del args, kwargs
    self._scatter_not_implemented("scatter_div")

  def scatter_min(self, *args, **kwargs):
    del args, kwargs
    self._scatter_not_implemented("scatter_min")

  def scatter_max(self, *args, **kwargs):
    del args, kwargs
    self._scatter_not_implemented("scatter_max")

  def scatter_update(self, *args, **kwargs):
    del args, kwargs
    self._scatter_not_implemented("scatter_update")

  def get_saveable(self, var, primary_var, name):
    """Create a saveable object for the given variable."""
    return values_util.get_on_read_saveable(var, primary_var, name)

  def get_restore_ops(self, var, tensor):
    """Restore the same value into all variables."""
    return values_util.get_on_read_restore_ops(var, tensor, self._aggregation)

class OnWritePolicy(VariablePolicy):
  """Policy defined for `tf.VariableSynchronization.ON_WRITE` synchronization.

  This policy is created when the following `synchronization` and `aggregation`
  parameters are specified when creating a `tf.Variable` in `tf.distribute`
  scope and `synchronization` is equal to `tf.VariableSynchronization.ON_WRITE`
  or `tf.VariableSynchronization.AUTO`.
  """

  def _is_mirrored(self):
    return True

  def value(self, var):
    return var._get_on_device_or_primary().value()  # pylint: disable=protected-access

  def _as_graph_element(self, var):
    return var._get_on_device_or_primary()._as_graph_element()  # pylint: disable=protected-access

  def _get_cross_replica(self, var):
    # Return identity, to avoid directly exposing the variable to the user and
    # allowing it to be modified by mistake.
    return array_ops.identity(var._get_on_device_or_primary())  # pylint: disable=protected-access

  def _update_replica(self, var, update_fn, value, **kwargs):
    if var.aggregation == variables_lib.VariableAggregation.NONE:
      return update_fn(var._get_on_device_or_primary(), value, **kwargs)  # pylint: disable=protected-access
    return _on_write_update_replica(var, update_fn, value, **kwargs)

  def assign(self, var, value, use_locking=False, name=None, read_value=True):
    return values_util.on_write_assign(
        var, value, use_locking=use_locking, name=name, read_value=read_value)

  def assign_add(self,
                 var,
                 value,
                 use_locking=False,
                 name=None,
                 read_value=True):
    return values_util.on_write_assign_add(
        var, value, use_locking=use_locking, name=name, read_value=read_value)

  def assign_sub(self,
                 var,
                 value,
                 use_locking=False,
                 name=None,
                 read_value=True):
    return values_util.on_write_assign_sub(
        var, value, use_locking=use_locking, name=name, read_value=read_value)

  def scatter_sub(self, var, sparse_delta, use_locking=False, name=None):
    return values_util.scatter_sub(
        var, sparse_delta, use_locking=use_locking, name=name)

  def scatter_add(self, var, sparse_delta, use_locking=False, name=None):
    return values_util.scatter_add(
        var, sparse_delta, use_locking=use_locking, name=name)

  def scatter_mul(self, var, sparse_delta, use_locking=False, name=None):
    return values_util.scatter_mul(
        var, sparse_delta, use_locking=use_locking, name=name)

  def scatter_div(self, var, sparse_delta, use_locking=False, name=None):
    return values_util.scatter_div(
        var, sparse_delta, use_locking=use_locking, name=name)

  def scatter_min(self, var, sparse_delta, use_locking=False, name=None):
    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
        self._aggregation != vs.VariableAggregation.NONE):
      raise NotImplementedError(
          values_util.scatter_error_msg.format(
              op_name="scatter_min", aggregation=self._aggregation))
    return values_util.scatter_min(
        var, sparse_delta, use_locking=use_locking, name=name)

  def scatter_max(self, var, sparse_delta, use_locking=False, name=None):
    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
        self._aggregation != vs.VariableAggregation.NONE):
      raise NotImplementedError(
          values_util.scatter_error_msg.format(
              op_name="scatter_max", aggregation=self._aggregation))
    return values_util.scatter_max(
        var, sparse_delta, use_locking=use_locking, name=name)

  def scatter_update(self, var, sparse_delta, use_locking=False, name=None):
    if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and
        self._aggregation != vs.VariableAggregation.NONE):
      raise NotImplementedError(
          values_util.scatter_error_msg.format(
              op_name="scatter_update", aggregation=self._aggregation))
    return values_util.scatter_update(
        var, sparse_delta, use_locking=use_locking, name=name)

  def get_saveable(self, var, primary_var, name):
    """Saveable ops for AUTO variables."""
    return values_util.get_on_write_saveable(var, primary_var, name)

  def get_restore_ops(self, var, tensor):
    return values_util.get_on_write_restore_ops(var, tensor)

class PerWorkerResource():
  """A per-worker CapturableResource class for non-ParameterServer strategy.

  Resources that populate `host_to_resources` should be instances of classes
  subclassing CapturableResource, although currently it's only used and tested
  for StaticHashTable with TPUStrategy.
  """

  def __init__(self, strategy, host_to_resources):
    distribute_lib.distribution_strategy_input_api_counter.get_cell(
        "PerWorkerResource", "TPUDistributedLookupTable").increase_by(1)
    self._strategy = strategy
    self._host_to_resources = host_to_resources

  def __getattribute__(self, name):
    if name not in ("__init__", "__getattribute__", "_host_to_resources",
                    "_strategy", "local_resource"):
      return getattr(self.local_resource(), name)
    return super(PerWorkerResource, self).__getattribute__(name)

  def __setattr__(self, name, value):
    if name not in ("_strategy", "_host_to_resources"):
      return setattr(self.local_resource(), name, value)
    return super(PerWorkerResource, self).__setattr__(name, value)

  def local_resource(self):
    """Returns the resource on the local worker."""
    current_device = device_util.canonicalize(device_util.current())
    host_device = device_util.canonicalize(
        device_util.get_host_for_device(current_device))
    return self._host_to_resources.get(
        host_device,
        self._host_to_resources[next(iter(self._host_to_resources))])


from django.shortcuts import render,redirect, get_object_or_404
from .models import Contato
from .forms import ContatoForm

def listar_contatos(request):
    busca = request.GET.get('pesquisa', None)

    if busca:
        contatos = Contato.objects.all()
        contatos = contatos.filter(nome__icontains=busca)

    else:
        contatos = Contato.objects.all()

    return render(request, 'lista.html',{'contatos': contatos})

def criar_contato(request):
    form = ContatoForm(request.POST or None, request.FILES or None)

    if form.is_valid():
        form.save()
        return redirect('listar_contatos')

    else:
        return render(request, "novo_contato.html", {'form': form})

def alterar_contato(request, id):
    contato = get_object_or_404(Contato, pk=id)
    form = ContatoForm(request.POST or None, request.FILES or None, instance=contato)

    if form.is_valid():
        form.save()
        return redirect('listar_contatos')

    return render(request, 'novo_contato.html', {'form': form})

def deletar_contato(request, id):
    contato = get_object_or_404(Contato, pk=id)

    if request.method == 'POST':
        contato.delete()
        return redirect('listar_contatos')

    return render(request, 'deletar_contato.html', {'contato': contato})


import numpy as np
import scipy.cluster.vq as cluster
from operations import gen_blobs, measure_time

if __name__ == '__main__':
    t_time = measure_time()
    print("Type,Obs,Feats,Clusters,Time")
    for obs_power in range(4, 11):
        num_obs = 2**obs_power
        for feats_power in range(1, 11):
            for clust_power in range(1, obs_power):
                features = 2**feats_power
                k = 2**clust_power
                observations, labels = gen_blobs(num_obs, features, k)
                for _ in range(50):
                    try:
                        time = measure_time()
                        centroids, labels = cluster.kmeans2(observations, k)
                        time = measure_time() - time
                    except:
                        continue

                    print("scipy,%d,%d,%d,%.9f" %(num_obs, features, k, time))

    total_time = measure_time() - t_time
    print(total_time)


import torch
import numpy as np
from torchvision import transforms
from torch.utils.data import Dataset
from PIL import Image, ImageFile
from ..utils import custom_transforms as tr

VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]

VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']

class VOCSegmentation(Dataset):
    def __init__(self, args, dataset_root, split):
        """crop_size: (h, w)"""
        self.split = split
        self.dataset_root = dataset_root
        self.args = args
        #* Here Use the SDB Aug Dataset.
        self.id_file = '%s/ImageSets/Segmentation/%s' % (dataset_root, 'train_aug.txt' if self.split=='train' else 'val.txt')
        with open(self.id_file, 'r') as f:
            self.img_ids = f.read().split() # 拆分成一个个名字组成list

        # Display stats
        print('Number of images in {}: {:d}'.format(split, len(self.img_ids)))

    def __len__(self):
        return len(self.img_ids)

    def __getitem__(self, idx):
        self.img_id = self.img_ids[idx]
        _image, _target = self._read_voc_images()
        sample = {'image': _image, 'label': _target}
        if self.split == "train":
            return self._transform_train(sample)
        elif self.split == 'val':
            return self._transform_val(sample)

    def _read_voc_images(self):
        image= Image.open('%s/JPEGImages/%s.jpg' % (self.dataset_root, self.img_id)).convert("RGB")
        mask_image = Image.open('%s/SegmentationClassAugRaw/%s.png' % (self.dataset_root, self.img_id))
        # mask_image = Image.open('%s/SegmentationClass/%s.png' % (self.dataset_root, self.img_id)).convert("RGB")  #* Here Use the SDB Aug Dataset.
        return image, mask_image # PIL image 0-255

    def _transform_train(self, sample):
        composed_transforms = transforms.Compose([
            tr.RandomHorizontalFlip(),
            tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size),
            tr.RandomGaussianBlur(),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])

        return composed_transforms(sample)

    def _transform_val(self, sample):

        composed_transforms = transforms.Compose([
            tr.FixScaleCrop(crop_size = self.args.crop_size),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])
        return composed_transforms(sample)

    def _label2colormap(self, target):
        """Convert label image to color image for visualization."""
        mask = np.zeros((target.shape[0], target.shape[1], 3), dtype=np.uint8)
        for i in range(len(VOC_COLORMAP)):
            mask[target == i] = VOC_COLORMAP[i]
        mask = Image.fromarray(mask)

def _voc_label_indices( mask_colormap ):
    """Convert the colormap image to label indices."""
    colormap2label = np.zeros(256**3, dtype=np.uint8) # torch.Size([16777216])
    for i, colormap in enumerate(VOC_COLORMAP):
        # 每个通道的进制是256，这样可以保证每个 rgb 对应一个下标 i
        colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
    mask_colormap = np.array(mask_colormap.convert("RGB")).astype('int32')
    idx = ((mask_colormap[:, :, 0] * 256 + mask_colormap[:, :, 1]) * 256 + mask_colormap[:, :, 2])
    return colormap2label[idx] # colormap 映射 到colormaplabel中计算的下标

if __name__ == '__main__':
    #TODO Test the model.
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt
    import argparse

    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    args.data.BASE_SIZE = 513
    args.data.CROP_SIZE = 513
    dataset_root = '/home/yangqi/dataspace/VOCdevkit/VOC2012'
    voc_train = VOCSegmentation(args, dataset_root, split='train')

    dataloader = DataLoader(voc_train, batch_size=5, shuffle=True, num_workers=0)

    for ii, sample in enumerate(dataloader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()
            segmap = sample['label'].numpy()
            segmap = segmap[jj]
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0])
            img_tmp *= (0.229, 0.224, 0.225)
            img_tmp += (0.485, 0.456, 0.406)
            img_tmp *= 255.0
            img_tmp = img_tmp.astype(np.uint8)
            plt.figure()
            plt.title('display')
            plt.subplot(211)
            plt.imshow(img_tmp)
            plt.subplot(212)
            plt.imshow(segmap)

        if ii == 1:
            break

    plt.show(block=True)


# -*- coding: utf-8 -*-
'''
The epoch
The epoch is the point where the time starts.
The point is taken as the 0.0 hours of January 1st
of the current year.
For Unix,the epoch is 0.0 hours of January 1st 1970
It is possible to measure the time in seconds since
the epoch using time() function of 'time' module
'''

'''
Program 1 : A python program to measure the time in
seconds since the epoch
'''
import time
epoch = time.time()
print(epoch)

'''
Program 2 : A python program to get date and time
from the epoch time

tm_year - 4 digit year like 2019
tm_mon - range [1,12]
tm_mday - range [1,31]
tm_hour - range [0,23]
tm_min - range [0,59]
tm_sec - range [0,61] including leap seconds
tm_wday - range [0,6] Monday is 0
tm_yday - range [1,366]
tm_isdst - [0,1 or -1] 0 - no DST,1 - DST is ineffect
-1= not known
tm_zone timezone name
'''
import time
epoch = time.time()
t = time.localtime(epoch)

#Retrieve the date from the structure t
d = t.tm_mday
m = t.tm_mon
y = t.tm_year
print('Current date is: %d - %d - %d' %(d,m,y))

#Retrive the time from the strurcture t
h = t.tm_hour
m = t.tm_min
s = t.tm_sec

print('Current time is: %d : %d : %d' %(h,m,s))

'''
Program 3 : A python program to convert epoch time
into corresponding date and time

The current date and time as shown in our computer
system can be known using the following
ctime() function of 'time' module
now() method of 'datetime' class of 'datetime' module
today() method of 'datetime' class of 'datetime' module
'''
import time
epoch = time.time()
t = time.ctime(epoch)
print('Current date and time ',t)

'''
Program 4 : A python program to know the current date
and time using ctime() function
'''

import time
t = time.ctime()
print('Current date and time is ',t)

'''
Program 5 : A python program to know the local date and time.

now () methods returns the object that contains date
and time information in any timezone.
By default the local time zone in india is IST
if the timezone is not passed
'''

from datetime import *
now = datetime.now()
print('current date time ' ,now)
print('Date now : {}/{}/{}'.format(now.day,now.month,now.year))
print('Time now : {}:{}:{}'.format(now.hour,now.minute,now.second))

'''
Program 6 : A python program to know today's date and time
'''

from datetime import *
# today() of timetime class gives date and time
tdm = datetime.today()
print("Today's date and time = ",tdm)

#today() of the date class gives the date only
td = date.today()
print("Today's date = ",td)


import os
import sys

from sqlalchemy import engine_from_config

from pyramid.paster import (
    get_appsettings,
    setup_logging,
    )

from pyramid.scripts.common import parse_vars
from itertools import takewhile

def usage(argv):
    cmd = os.path.basename(argv[0])
    print('usage: %s <config_uri> [var=value, [var2=value] django management commands\n'
          '(example: "%s development.ini check")' % (cmd, cmd))
    sys.exit(1)

def main(argv=sys.argv):
    if len(argv) < 2:
        usage(argv)

    config_uri = argv[1]
    setup_logging(config_uri)

    var_strings = list(takewhile(lambda x: '=' in x, argv[2:]))
    options = parse_vars(var_strings)
    settings = get_appsettings(config_uri, options=options)

    del argv[1:2 + len(var_strings)]

    os.environ.setdefault("DJANGO_SETTINGS_MODULE", settings['django.settings'])
    sys.path.insert(0, settings['django.project'])
    sys.argv[0] = os.path.join(settings['django.project'], 'manage.py')

    from django.core.management import execute_from_command_line
    execute_from_command_line(sys.argv)


import pandas as pd
import torch
from typing import Tuple

RAW_DATA = "data/ES_continuous_adjusted_1min.txt"

def _load_data() -> pd.DataFrame:

    return pd.read_csv(
        RAW_DATA,
        engine="pyarrow",
        header=None,
        names=["time", "open", "high", "low", "close", "volume"]
    )

def get_univariate_dataset(column, limit=None) -> pd.Series:

    data = _load_data()

    data = data[column]

    if limit is not None:
        data = data[:limit]

    return data

def get_multivariate_dataset(columns, limit=None) -> pd.Series:

    data = _load_data()

    data = data[columns]

    if limit is not None:
        data = data[:limit]

    return data

def train_test_split(
        data, train_ratio=0.8,
        val_ratio=0.1
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

    num_data = data.shape[0]

    train_start_index = 0
    train_end_index = int(train_ratio * num_data)

    val_start_index = int(train_ratio * num_data)
    val_end_index = int((train_ratio + val_ratio) * num_data)

    test_start_index = int((train_ratio + val_ratio) * num_data)

    print(
        f'Dataset shape: {data.shape}\n',
        f'Train:         {train_start_index}:{train_end_index}\n',
        f'Val:           {val_start_index}:{val_end_index}\n',
        f'Test:          {test_start_index}:{num_data}\n'
    )

    train_data = data[train_start_index:train_end_index]
    val_data = data[val_start_index:val_end_index]
    test_data = data[test_start_index:num_data]

    #return (
    #    torch.tensor(train_data.values, dtype=torch.float32),
    #    torch.tensor(val_data.values, dtype=torch.float32),
    #    torch.tensor(test_data.values, dtype=torch.float32)
    #)

    return (
        train_data,
        val_data,
        test_data
    )


import datetime as dt
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.email_operator import EmailOperator
from airflow.models.baseoperator import chain

with DAG('dummy_dag',
    description='An Airflow DAG running dbt and Great Expectations tasks',
    schedule_interval=None,
    start_date=dt.datetime(2022, 7, 21),
    catchup=False,
    render_template_as_native_obj=True
    ) as dag:

    send_email_1 = EmailOperator(
        task_id='send_email_1',
        to='user1@gmail.com',
        subject='Airflow Alert',
        html_content=f"""
Hi user1, \nData Validated and Transformed Successfully!
 """,
        dag=dag
        )

    send_email_2 = EmailOperator(
        task_id='send_email_2',
        to='user2@gmail.com',
        subject='Airflow Alert',
        html_content=f"""
Hi user2, \nData Validated and Transformed Successfully!
 """,
        dag=dag
        )

    send_email_3 = EmailOperator(
        task_id='send_email_3',
        to='user3@gmail.com',
        subject='Airflow Alert',
        html_content=f"""
Hi user3, \nData Validated and Transformed Successfully!
 """,
        dag=dag
        )

    begin = DummyOperator(task_id="begin")
    end = DummyOperator(task_id="end")

    chain(
        begin,
        send_email_1,
        send_email_2,
        send_email_3,
        end,
    )


import usb  # This is pyusb

class uDMX():
    SingleChannelModeFlag = 1
    MultiChannelModeFlag = 2

    def __init__(self):
        self.vid = 0x16c0
        self.pid = 0x05dc
        self.dev = None
        self.channel_mode = "single"  # single = 1 or multi = 2

        self.dev = usb.core.find(idVendor=self.vid, idProduct=self.pid)
        if self.dev is None:
            raise ResourceWarning("uDMX device was not found")

        self.bmRequestType = usb.util.CTRL_TYPE_VENDOR | usb.util.CTRL_RECIPIENT_DEVICE | usb.util.CTRL_OUT

    def setDMX(self, channel, value):
        channel = channel - 1
        if type(value) is int:
            n = self.dev.ctrl_transfer(self.bmRequestType, self.SingleChannelModeFlag, wValue=value, wIndex=channel,
                                       data_or_wLength=1)
        elif type(value) is list:
            channel_values = bytearray(value)
            n = self.dev.ctrl_transfer(self.bmRequestType, self.MultiChannelModeFlag, wValue=len(channel_values), wIndex=channel, data_or_wLength=channel_values)
        else:
            raise ValueError("The value can be an integer or a list.")


import os

import nltk
from nltk.tokenize import word_tokenize

class FileReader:
    def __init__(self, file_name):
        self.file_name = file_name
        file = None
        self.file_lines = []
        self.file_content = ''
        self.line_count = None
        self.word_count = None
        try:
            file = open(file_name, 'r', encoding='utf8')
            self.file_lines = file.readlines()
            self.file_content = ''.join(self.file_lines)
            self.line_count = len(self.file_lines)
            self.word_count = len(
                word_tokenize(self.file_content, language='russian')
            )
        except FileNotFoundError:
            print('ERROR: file does not exist')
        file.close()

    def read(self):
        return self.file_contentЫ

    def write_line(self, content):
        try:
            file = open(self.file_name)
            file.write(content)
        except FileNotFoundError:
            print('ERROR: file does not exist')

    def count_lines(self):
        return [self.line_count, self.word_count]

    def __add__(self, other):
        new_name = self.file_name[:-4] + '_add_' + other.file_name[:-4] + '.txt'
        new_content = self.file_lines + other.file_lines

        try:
            new_file = open(new_name, 'w', encoding='utf8')
            new_file.writelines(new_content)
            new_file.close()
        except FileExistsError:
            print('ERROR: file already exist')

        new_file_reader = FileReader(new_name)
        return new_file_reader

    def __str__(self):
        return os.getcwd() + self.file_name


import math

import numpy as np
import matplotlib.pyplot as plt

from data.data_set import OurSet
from env.environment import ConflictEnv
from parameters import *

def draw(env, expert_set):
    train_num_list = [int(info['id']) for info in env.train_set]
    test_num_list = [int(info['id']) for info in env.test_set]
    _, train_action_list = expert_set.get_next_batch(train_num_list)
    _, test_action_list = expert_set.get_next_batch(test_num_list)
    total_action_list = list(train_action_list) + list(test_action_list)

    fig, axes = plt.subplots(3, 1, sharex=True, figsize=(6, 6))
    axes[0].hist(train_action_list, bins=bins)
    axes[0].xaxis.set_visible(False)
    ax_twin = axes[0].twinx()
    ax_twin.plot(*stat(train_action_list), 'red', label='train')
    ax_twin.ticklabel_format(style='sci', scilimits=(0, 1), axis='y')
    ax_twin.xaxis.set_visible(False)
    ax_twin.legend(fontsize=14)

    axes[1].hist(test_action_list, bins=bins)
    axes[1].xaxis.set_visible(False)
    ax_twin = axes[1].twinx()
    ax_twin.plot(*stat(test_action_list), 'green', label='test')
    ax_twin.ticklabel_format(style='sci', scilimits=(0, 1), axis='y')
    ax_twin.xaxis.set_visible(False)
    ax_twin.legend(fontsize=14)

    axes[2].hist(total_action_list, bins=bins)
    axes[2].xaxis.set_visible(False)
    ax_twin = axes[2].twinx()
    ax_twin.plot(*stat(total_action_list), 'blue', label='total')
    ax_twin.ticklabel_format(style='sci', scilimits=(0, 1), axis='y')
    ax_twin.xaxis.set_visible(False)
    ax_twin.legend(fontsize=14)

    fig.text(0.5, 0.05, 'Action Index', ha='center', fontsize=16)
    fig.text(0.05, 0.5, 'Count', va='center', rotation='vertical', fontsize=16)
    fig.text(0.95, 0.5, 'Frequency', va='center', rotation='vertical', fontsize=16)
    plt.show()

def stat(lst):
    length = len(lst)

    result = {i: 0 for i in range(bins)}
    for v in lst:
        assert v in result.keys()
        result[v] += 1

    x, y = [], []
    for key, value in result.items():
        x.append(key)
        y.append(value / length)
    return x, y

def kl_divergence(p, q):
    p = np.array(stat(p))
    q = np.array(stat(q))
    return np.sum(np.where(p != 0, p * np.log(p / (q+1e-8)), 0))

def get_acr_mse(int_num_list, int_act, exp_num_list, exp_act):
    acr_list = []
    mse_list = []
    for num, act in zip(int_num_list, int_act):
        e_act = exp_act[exp_num_list.index(num)]
        delta = act - e_act
        acr_list.append(int(delta == 0))
        mse_list.append(delta * delta)
    return np.mean(acr_list), math.sqrt(np.mean(mse_list))

def _analysis(path, prefix, start=0, end=10001, delta=1000):
    traj_data = np.load(expert_path)
    print(list(traj_data.keys()))
    expert_num = list(traj_data['name'])
    expert_act = traj_data['action']

    train_idx, test_idx = None, None
    expert_train_act, expert_test_act = None, None
    train_acr_list, test_acr_list, total_acr_list = [], [], []
    train_mse_list, test_mse_list, total_mse_list = [], [], []
    kl_train_lst, kl_test_lst = [], []
    sr_train_lst, sr_test_lst = [], []
    steps = list(range(start, end, delta))
    ad_lst = []
    for step in steps:
        print(path + prefix + '_train_{}.npz'.format(step))
        train_data = np.load(path + prefix + '_train_{}.npz'.format(step))
        test_data = np.load(path + prefix + '_test_{}.npz'.format(step))
        train_num, test_num = train_data['name'], test_data['name']
        train_act, test_act = train_data['action'], test_data['action']

        if train_idx is None:
            train_idx = [expert_num.index(num) for num in train_num]
            test_idx = [expert_num.index(num) for num in test_num]
            expert_train_act = expert_act[train_idx]
            expert_test_act = expert_act[test_idx]

        # AD
        kl_train = kl_divergence(expert_train_act, train_act)
        kl_train_lst.append(kl_train)
        kl_test = kl_divergence(expert_test_act, test_act)
        kl_test_lst.append(kl_test)
        ad_lst.append([stat(train_act), stat(test_act)])
        # SR
        sr_train_lst.append(np.mean(train_data['result']))
        sr_test_lst.append(np.mean(test_data['result']))
        # ACR, MSE
        acr, mse = get_acr_mse(train_num, train_act, expert_num, expert_act)
        train_acr_list.append(acr)
        train_mse_list.append(mse)
        print(step, acr, mse, end=' ')
        acr, mse = get_acr_mse(test_num, test_act, expert_num, expert_act)
        test_acr_list.append(acr)
        test_mse_list.append(mse)
        print(acr, mse)

    fig, axes = plt.subplots(2, 3)

    ax00 = axes[0, 0]
    ax10 = axes[1, 0]
    ax00.plot(*stat(expert_train_act), label='expert_train')
    ax10.plot(*stat(expert_test_act), label='expert_test')

    for i, [train_ad, test_ad] in enumerate(ad_lst):
        step = steps[i]
        if step not in [1000, 10000, end-1, ]:
            continue
        ax00.plot(*train_ad, label='train_{}'.format(step))
        ax10.plot(*test_ad, label='test_{}'.format(step))

    ax01 = axes[0, 1]
    ax11 = axes[1, 1]
    ax01.plot(steps, kl_train_lst, label='train_kl')
    ax01.plot(steps, kl_test_lst, label='test_kl')
    ax11.plot(steps, sr_train_lst, label='train_sr')
    ax11.plot(steps, sr_test_lst, label='test_sr')

    ax02 = axes[0, 2]
    ax12 = axes[1, 2]
    ax02.plot(steps, train_acr_list, label='train_acr')
    ax02.plot(steps, test_acr_list, label='test_acr')
    ax12.plot(steps, train_mse_list, label='train_mse')
    ax12.plot(steps, test_mse_list, label='test_mse')

    [ax.legend() for ax in [ax00, ax10, ax01, ax11, ax02, ax12]]
    plt.show()

if __name__ == '__main__':
    kwargs = {
        "start": 0,
        "end": 100001,
        "delta": 5000
    }

    # draw(
    #     env=ConflictEnv(limit_path=limit_path, size=size, ratio=ratio),
    #     expert_set=OurSet(expert_path=expert_path)
    # )

    # _analysis(path=bc_actions_path, prefix='bc_actions', **kwargs)
    # _analysis(path=dqn_no_exp_path, prefix='dqn_no_exp', **kwargs)
    _analysis(path=gail_actions_path, prefix='dqn_actions', **kwargs)
    # _analysis(path=gail_tracks_path, prefix='dqn_tracks', **kwargs)


import itertools
from collections import Counter

from core.common.SemanticSimilarityAlgorithm import SemanticSimilarityAlgorithm

import numpy as np

class SifCosineSimilarity(SemanticSimilarityAlgorithm):

    def get_label(self):
        return "Smooth Inverse Frequency + Cosine Distance"

    def __init__(self, embedding, stemmer):
        super().__init__()
        self.a = 0.001
        self.embedding = embedding
        self.stemmer = stemmer

    def _absolute_score(self, sentence_1, sentence_2):
        splitted_1 = self.stemmer.stem_sentence(sentence_1)
        splitted_2 = self.stemmer.stem_sentence(sentence_2)
        vecs = self._sif_feature_vectors(splitted_1, splitted_2, self.embedding)
        A = vecs[0]
        B = vecs[1]
        return A @ B / (np.linalg.norm(A) * np.linalg.norm(B))

    def _normalize(self, absolute):
        return absolute * 50 + 50

    # Based on
    # medium.com/@Intellica.AI/comparison-of-different-word-embeddings-on-text-similarity-a-use-case-in-nlp-e83e08469c1c
    def _sif_feature_vectors(self, splitted_1, splitted_2, embedding, a=0.001):
        s_1 = [token for token in splitted_1 if token in embedding.model.wv.vocab]
        s_2 = [token for token in splitted_2 if token in embedding.model.wv.vocab]
        word_counts = self._map_word_frequency((s_1 + s_2))
        sentence_set = []
        for sentence in [s_1, s_2]:
            vs = np.zeros(embedding.model_shape)
            sentence_length = len(sentence)
            for word in sentence:
                a_value = a / (a + word_counts[word])  # smooth inverse frequency, SIF
                vs = np.add(vs, np.multiply(a_value, embedding.model.wv[word]))  # vs += sif * word_vector
            vs = np.divide(vs, sentence_length)  # weighted average
            sentence_set.append(vs)
        return sentence_set

    def _map_word_frequency(self, document):
        return Counter(itertools.chain(*document))


import argparse
import sys

import matplotlib.pyplot as plt
import numpy as np
import gurobipy as gb

sys.path.append("../..")
from utils.auxiliary_energy import *
from utils.evolution import out_of_sample_test

parser = argparse.ArgumentParser()
# name of example
parser.add_argument('--name', help='example name', type=str, default='Energy')
# number of qubits
parser.add_argument('--n', help='number of qubits', type=int, default=2)
# number of edges for generating regular graph
parser.add_argument('--num_edges', help='number of edges for generating regular graph', type=int,
                    default=1)
# if generate the graph randomly
parser.add_argument('--rgraph', help='if generate the graph randomly', type=int, default=0)
# number of instances
parser.add_argument('--g_seed', help='graph seed', type=int, default=0)
# seed for uncertainties
parser.add_argument('--r_seed', help='seed for uncertainties', type=int, default=0)
# number of scenarios
parser.add_argument('--num_scenario', help='number of scenarios', type=int, default=100)
# evolution time
parser.add_argument('--evo_time', help='evolution time', type=float, default=2)
# time steps
parser.add_argument('--n_ts', help='time steps', type=int, default=40)
# initial control file
parser.add_argument('--control', help='file name of control', type=str, default=None)
# threshold of CVaR
parser.add_argument('--cvar', help='threshold of CVaR', type=float, default=0.01)
parser.add_argument('--bound', help='bound of uniform sample', type=float, default=0.3)

args = parser.parse_args()

y0 = uniform(args.n)

if args.rgraph == 0:
    Jij, edges = generate_Jij_MC(args.n, args.num_edges, 100)

    C = get_ham(args.n, True, Jij)
    B = get_ham(args.n, False, Jij)

    args.seed = 0

if args.rgraph == 1:
    Jij = generate_Jij(args.n, args.g_seed)
    C = get_ham(args.n, True, Jij)
    B = get_ham(args.n, False, Jij)

if args.num_scenario == 1:
    c_uncertainty = np.zeros((2, 1))
else:
    np.random.seed(args.r_seed)
    c_uncertainty = np.random.uniform(low=-args.bound, high=args.bound, size=(2, args.num_scenario))
prob = np.ones(args.num_scenario) * 1 / args.num_scenario

if not os.path.exists("../../output/OutSampleTest/"):
    os.makedirs("../../output/OutSampleTest/")

output_num = "../../output/OutSampleTest/" + args.control.split('/')[-1].split('.csv')[0] + "_out_test" + \
             "cvar{}_scenario{}_group{}_bound{}".format(args.cvar, args.num_scenario, args.r_seed, args.bound) + ".log"

f = open(output_num, "w+")
control = np.loadtxt(args.control, delimiter=',')
obj = out_of_sample_test(args.num_scenario, c_uncertainty, np.zeros((2 ** args.n, 2 ** args.n), dtype=complex),
                         [B, C], args.n_ts, args.evo_time, control, y0[0:2 ** args.n], None, 'energy')
print("********* Results *****************", file=f)
m = gb.Model()
zeta_var = m.addVar(lb=-np.infty)
cos = m.addVars(args.num_scenario, lb=0)
m.addConstrs(cos[l] >= obj[l] - zeta_var for l in range(args.num_scenario))
m.setObjective(zeta_var + 1 / args.cvar * gb.quicksum(prob[l] * cos[l] for l in range(args.num_scenario)))
m.optimize()
print("Final original objective value", file=f)
for l in range(args.num_scenario):
    print(obj[l], file=f)

print("********* Summary *****************", file=f)
print("zeta", zeta_var.x, file=f)
print("objective value", m.objval, file=f)
print("Final maximum energy {}".format(-min(abs(obj))), file=f)
print("Final average original objective value {}".format(sum(obj) / len(obj)), file=f)
print("Final original objective value exceeding current objective value {}".format(obj[obj > m.objval]), file=f)