module TensorStream
  module Train
    # convenience methods used for training
    module Utils
      def create_global_step(graph = nil)
        target_graph = graph || TensorStream.get_default_graph
        raise TensorStream::ValueError, '"global_step" already exists.' unless get_global_step(target_graph).nil?

        TensorStream.variable_scope.get_variable(TensorStream::GraphKeys::GLOBAL_STEP, shape: [],
                                                                                       dtype: :int64,
                                                                                       initializer: TensorStream.zeros_initializer,
                                                                                       trainable: false,
                                                                                       collections: [TensorStream::GraphKeys::GLOBAL_VARIABLES,
                                                                                                     TensorStream::GraphKeys::GLOBAL_STEP,])
      end

      def get_global_step(graph = nil)
        target_graph = graph || TensorStream.get_default_graph
        global_step_tensors = target_graph.get_collection(TensorStream::GraphKeys::GLOBAL_STEP)
        global_step_tensor = if global_step_tensors.nil? || global_step_tensors.empty?
          begin
            target_graph.get_tensor_by_name("global_step:0")
          rescue TensorStream::KeyError
            nil
          end
        elsif global_step_tensors.size == 1
          global_step_tensors[0]
        else
          TensorStream.logger.error("Multiple tensors in global_step collection.")
          nil
        end
        global_step_tensor
      end
    end
  end
end
