text: # Copyright 2018/2019 The RLgraph authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  from __future__ import absolute_import, division, print_function
  
  import numpy as np
  from rlgraph import get_backend
  from rlgraph.components.layers.preprocessing.preprocess_layer import PreprocessLayer
  from rlgraph.utils.decorators import rlgraph_api
  from rlgraph.utils.util import SMALL_NUMBER
  
  if get_backend() == "tf":
      import tensorflow as tf
  elif get_backend() == "pytorch":
      import torch
  
  
  class MovingStandardize(PreprocessLayer):
      """
      Standardizes inputs using a moving estimate of mean and std.
      """
      def __init__(self, batch_size=1, scope="moving-standardize", **kwargs):
          """
          Args:
              batch_size (int): Number of samples processed per step.
          """
          super(MovingStandardize, self).__init__(scope=scope, **kwargs)
          self.batch_size = batch_size
          self.sample_count = None
  
          # Current estimate of state mean.
          self.mean_est = None
  
          # Current estimate of sum of stds.
          self.std_sum_est = None
          #self.output_spaces = None
          self.in_shape = None
  
      def create_variables(self, input_spaces, action_space=None):
          in_space = input_spaces["inputs"]
          #self.output_spaces = in_space
          self.in_shape = (self.batch_size, ) + in_space.shape
  
          if self.backend == "python" or get_backend() == "python" or get_backend() == "pytorch":
              self.sample_count = 0.0
              self.mean_est = np.zeros(self.in_shape, dtype=np.float32)
              self.std_sum_est = np.zeros(self.in_shape, dtype=np.float32)
          elif get_backend() == "tf":
              self.sample_count = self.get_variable(name="sample-count", dtype="float", initializer=0.0, trainable=False)
              self.mean_est = self.get_variable(
                  name="mean-est",
                  shape=self.in_shape,
                  dtype=tf.float32,
                  trainable=False,
                  initializer=tf.zeros_initializer()
              )
              self.std_sum_est = self.get_variable(
                  name="std-sum-est",
                  shape= self.in_shape,
                  dtype=tf.float32,
                  trainable=False,
                  initializer=tf.zeros_initializer()
              )
  
      @rlgraph_api
      def _graph_fn_reset(self):
          if self.backend == "python" or get_backend() == "python" or get_backend() == "pytorch":
              self.sample_count = 0.0
              self.mean_est = np.zeros(self.in_shape)
              self.std_sum_est = np.zeros(self.in_shape)
          elif get_backend() == "tf":
              return tf.variables_initializer([self.sample_count, self.mean_est, self.std_sum_est])
  
      @rlgraph_api
      def _graph_fn_call(self, inputs):
          if self.backend == "python" or get_backend() == "python" or get_backend() == "pytorch":
              # https://www.johndcook.com/blog/standard_deviation/
              inputs = np.asarray(inputs, dtype=np.float32)
              self.sample_count += 1.0
              if self.sample_count == 1.0:
                  self.mean_est[...] = inputs
              else:
                  update = inputs - self.mean_est
                  self.mean_est[...] += update / self.sample_count
                  self.std_sum_est[...] += update * update * (self.sample_count - 1.0) / self.sample_count
  
              # Subtract mean.
              result = inputs - self.mean_est
  
              # Estimate variance via sum of variance.
              if self.sample_count > 1.0:
                  var_estimate = self.std_sum_est / (self.sample_count - 1.0)
              else:
                  var_estimate = np.square(self.mean_est)
              std = np.sqrt(var_estimate) + SMALL_NUMBER
  
              standardized = result / std
              if get_backend() == "pytorch":
                  standardized = torch.Tensor(standardized)
              return standardized
  
          elif get_backend() == "tf":
              assignments = [tf.assign_add(ref=self.sample_count, value=1.0)]
              with tf.control_dependencies(assignments):
                  # 1. Update vars
                  assignments = []
                  update = inputs - self.mean_est
                  mean_update = tf.cond(
                      pred=self.sample_count > 1.0,
                      false_fn=lambda: self.mean_est,
                      true_fn=lambda: update
                  )
                  var_update = update * update * (self.sample_count - 1) / self.sample_count
                  assignments.append(tf.assign_add(ref=self.mean_est, value=mean_update))
                  assignments.append(tf.assign_add(ref=self.std_sum_est, value=var_update))
  
              with tf.control_dependencies(assignments):
                  # 2. Compute var estimate after update.
                  var_estimate = tf.cond(
                      pred=self.sample_count > 1,
                      false_fn=lambda: tf.square(x=self.mean_est),
                      true_fn=lambda: self.std_sum_est / (self.sample_count - 1)
                  )
                  result = inputs - self.mean_est
                  std = tf.sqrt(x=var_estimate) + SMALL_NUMBER
  
                  return result / std
  