# This file has ben automatically generated by stubgen
# DO NOT EDIT
#
module TensorStream
  module OpStub

    ##
    # Returns x + y element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def add(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:add, input_a, input_b, name: name)
    end

    ##
    # Returns the index with the largest value across axes of a tensor.
    #
    #
    # @param input_a tensor X (of type NUMERIC_TYPES)
    # @param axis Describes which axis of the input tensor to reduce across. For vectors, use axis = 0 (of type INTEGER_TYPES)
    #
    # Options:
    # @option name Optional name
    # @option dimension Same as axis
    # @option output_type Output data type defaults to int32 default (:int32)
    # @return Tensor
    def argmax(input_a, axis = nil, name: nil, dimension: nil, output_type: :int32)
      check_allowed_types(input_a, TensorStream::Ops::NUMERIC_TYPES)
      check_allowed_types(axis, TensorStream::Ops::INTEGER_TYPES)
      _op(:argmax, input_a, axis, name: name, dimension: dimension, output_type: output_type)
    end

    ##
    # Returns the index with the smallest value across axes of a tensor.
    #
    #
    # @param input_a tensor X (of type NUMERIC_TYPES)
    # @param axis Describes which axis of the input tensor to reduce across. For vectors, use axis = 0 (of type INTEGER_TYPES)
    #
    # Options:
    # @option name Optional name
    # @option dimension Same as axis
    # @option output_type Output data type defaults to int32 default (:int32)
    # @return Tensor
    def argmin(input_a, axis = nil, name: nil, dimension: nil, output_type: :int32)
      check_allowed_types(input_a, TensorStream::Ops::NUMERIC_TYPES)
      check_allowed_types(axis, TensorStream::Ops::INTEGER_TYPES)
      _op(:argmin, input_a, axis, name: name, dimension: dimension, output_type: output_type)
    end

    ##
    # Returns element-wise smallest integer in not less than x
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def ceil(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:ceil, input_a, name: name)
    end

    ##
    # Computes cos of input element-wise.
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def cos(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:cos, input_a, name: name)
    end

    ##
    # Returns x / y element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def div(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:div, input_a, input_b, name: name)
    end

    ##
    # Returns the truth value of (x == y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def equal(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:equal, input_a, input_b, name: name)
    end

    ##
    # Inserts a dimension of 1 into a tensor's shape.
    # Given a tensor input, this operation inserts a dimension of 1 at the dimension index axis of input's shape. The
    # dimension index axis starts at zero; if you specify a negative number for axis it is counted backward from the end.
    #
    #
    # @param input A tensor
    # @param axis Specifies the dimension index at which to expand the shape of input. Must be in the range [-rank(input) - 1, rank(input)].
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def expand_dims(input, axis, name: nil)
      _op(:expand_dims, input, axis, name: name)
    end

    ##
    # This operation creates a tensor of shape dims and fills it with value.
    #
    #
    # @param dims tensor shape
    # @param value scalar value to fill with
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def fill(dims, value, name: nil)
      _op(:fill, dims, value, name: name)
    end

    ##
    # Returns element-wise largest integer not greater than x.
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def floor(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:floor, input_a, name: name)
    end

    ##
    # Returns element-wise integer divistion.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def floor_div(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:floor_div, input_a, input_b, name: name)
    end

    ##
    # Returns the truth value of (x > y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def greater(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:greater, input_a, input_b, name: name)
    end

    ##
    # Returns the truth value of (x >= y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def greater_equal(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:greater_equal, input_a, input_b, name: name)
    end

    ##
    # Returns the truth value of (x < y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def less(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:less, input_a, input_b, name: name)
    end

    ##
    # Returns the truth value of (x <= y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def less_equal(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:less_equal, input_a, input_b, name: name)
    end

    ##
    # Computes natural logarithm of x element-wise.
    #
    #
    # @param input tensor X
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def log(input, name: nil)
      _op(:log, input, name: name)
    end

    ##
    # Multiplies matrix a by matrix b, producing a * b. The inputs must, following any transpositions, be tensors of rank 2 .
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option transpose_a Transpose matrix A first default (false)
    # @option transpose_b Transpose matrix B first default (false)
    # @option name Optional name
    # @return Tensor
    def mat_mul(input_a, input_b, transpose_a: false, transpose_b: false, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:mat_mul, input_a, input_b, transpose_a: transpose_a, transpose_b: transpose_b, name: name)
    end

    alias_method :matmul, :mat_mul

    ##
    # Returns the max of x and y (i.e. x > y ? x : y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X (of type NUMERIC_TYPES)
    # @param input_b tensor Y (of type NUMERIC_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def max(input_a, input_b, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::NUMERIC_TYPES)
      check_allowed_types(input_b, TensorStream::Ops::NUMERIC_TYPES)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:max, input_a, input_b, name: name)
    end

    ##
    # Returns the min of x and y (i.e. x < y ? x : y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X (of type NUMERIC_TYPES)
    # @param input_b tensor Y (of type NUMERIC_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def min(input_a, input_b, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::NUMERIC_TYPES)
      check_allowed_types(input_b, TensorStream::Ops::NUMERIC_TYPES)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:min, input_a, input_b, name: name)
    end

    ##
    # Returns element-wise remainder of division.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def mod(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:mod, input_a, input_b, name: name)
    end

    ##
    # Returns x * y element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def mul(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:mul, input_a, input_b, name: name)
    end

    ##
    # Computes numerical negative value element-wise.
    #
    #
    # @param input tensor X
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def negate(input, name: nil)
      _op(:negate, input, name: name)
    end

    ##
    # Returns the truth value of (x != y) element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def not_equal(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:not_equal, input_a, input_b, name: name)
    end

    ##
    # Creates a tensor with all elements set to 1.
    # Given a single tensor (tensor), this operation returns a
    # tensor of the same type and shape as tensor with all elements set to 1.
    # Optionally, you can specify a new type (dtype) for the returned tensor.
    #
    #
    # @param input A tensor
    #
    # Options:
    # @option dtype Optional new data type to cast into
    # @option name Optional name
    # @return Tensor
    def ones_like(input, dtype: nil, name: nil)
      _op(:ones_like, input, data_type: dtype, name: name)
    end

    ##
    # Computes the power of one value to another X^Y element wise
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def pow(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:pow, input_a, input_b, name: name)
    end

    ##
    # Computes the product of elements across dimensions of a tensor.
    # Reduces input_tensor along the dimensions given in axis. Unless keepdims is true, the rank of the
    # tensor is reduced by 1 for each entry in axis. If keepdims is true, the reduced dimensions are
    # retained with length 1.
    # If axis has no entries, all dimensions are reduced, and a tensor with a single element is returned.
    #
    #
    # @param input_a tensor X
    # @param axis tensor X (of type INTEGER_TYPES)
    #
    # Options:
    # @option name Optional name
    # @option keepdims If true, retains reduced dimensions with length 1. default (false)
    # @return Tensor
    def prod(input_a, axis = nil, name: nil, keepdims: false)
      check_allowed_types(axis, TensorStream::Ops::INTEGER_TYPES)
      input_a = TensorStream.convert_to_tensor(input_a)
      return input_a if input_a.shape.scalar?
      axis = cast_axis(input_a, axis)
      _op(:prod, input_a, axis, name: name, keepdims: keepdims)
    end

    alias_method :reduce_prod, :prod

    ##
    # Outputs random values from a uniform distribution.
    #
    #
    # @param shape A 1-D integer Tensor or array. The shape of the output tensor.
    #
    # Options:
    # @option name Optional name
    # @option dtype The type of the output: float16, float32, float64, int32, or int64 default (:float32)
    # @option minval A 0-D Tensor or ruby value of type dtype. The lower bound on the range of random values to generate. Defaults to 0. default (0)
    # @option maxval A 0-D Tensor or ruby value of type dtype. The upper bound on the range of random values to generate. Defaults to 1 if dtype is floating point. default (1)
    # @option seed  A ruby integer. Used to create a random seed for the distribution. See set_random_seed for behavior.
    # @return Tensor
    def random_uniform(shape, name: nil, dtype: :float32, minval: 0, maxval: 1, seed: nil)
      _op(:random_uniform, shape, name: name, dtype: dtype, minval: minval, maxval: maxval, seed: seed)
    end

    ##
    # Creates a sequence of numbers.
    # Creates a sequence of numbers that begins at start and extends by increments of delta up to but not including limit.
    #
    #
    # @param start Acts as first entry in the range if limit is not nil; otherwise, acts as range limit and first entry defaults to 0.
    # @param limit Upper limit of sequence, exclusive. If nil, defaults to the value of start while the first entry of the range defaults to 0.
    # @param delta Number that increments start. Defaults to 1.
    #
    # Options:
    # @option name  A name for the operation. Defaults to "range". default ("range")
    # @option dtype The type of the elements of the resulting tensor.
    # @option output_type Output data type defaults to int32 default (:int32)
    # @return Tensor
    def range(start = 0, limit = 0, delta = 1, name: "range", dtype: nil, output_type: :int32)
      _op(:range, start, limit, delta, name: name, dtype: dtype, output_type: output_type)
    end

    ##
    # Returns the rank of a tensor
    #
    #
    # @param input A tensor
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def rank(input, name: nil)
      input = convert_to_tensor(input)
      return cons(input.shape.ndims) if input.shape.known?
      _op(:rank, input, name: name)
    end

    ##
    # Reshapes a tensor.
    # Given tensor, this operation returns a tensor that has the same values as tensor with shape shape.
    #
    #
    # @param input A tensor
    # @param shape A new tensor shape
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def reshape(input, shape, name: nil)
      _op(:reshape, input, shape, name: name)
    end

    ##
    # Rounds the values of a tensor to the nearest integer, element-wise
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def round(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:round, input_a, name: name)
    end

    ##
    # Computes reciprocal of square root of x element-wise.
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def rsqrt(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:rsqrt, input_a, name: name)
    end

    ##
    # This operation returns a 1-D integer tensor representing the shape of input
    #
    #
    # @param input A tensor
    #
    # Options:
    # @option name Optional name
    # @option out_type Optional output type default (:int32)
    # @return Tensor
    def shape(input, name: nil, out_type: :int32)
      return constant(shape_eval(input, out_type), dtype: out_type, name: "Shape/#{name}") if input.is_a?(Array) && !input[0].is_a?(Tensor)
      return constant(input.shape.shape, dtype: out_type, name: "Shape/#{input.name}_c") if shape_full_specified(input)
      _op(:shape, input, name: name, out_type: out_type)
    end

    ##
    # Computes sigmoid of x element-wise.
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def sigmoid(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:sigmoid, input_a, name: name)
    end

    ##
    # Computes sign of input element-wise.
    # <tt>y = sign(x) = -1 if x < 0; 0 if x == 0 or tf.is_nan(x); 1 if x > 0.</tt>
    # Zero is returned for NaN inputs.
    #
    #
    # @param input_a tensor X
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def sign(input_a, name: nil)
      _op(:sign, input_a, name: name)
    end

    ##
    # Computes sin of input element-wise.
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def sin(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:sin, input_a, name: name)
    end

    ##
    # Returns the size of a tensor.
    # Returns a 0-D Tensor representing the number of elements in input of type out_type. Defaults to :int32.
    #
    #
    # @param input A tensor
    #
    # Options:
    # @option name Optional name
    # @option out_type Optional output type default (:int32)
    # @return Tensor
    def size(input, name: nil, out_type: :int32)
      _op(:size, input, name: name, out_type: out_type)
    end

    ##
    # Extracts a strided slice of a tensor
    # this op extracts a slice of size `(end-begin)/stride`
    #   from the given `input_` tensor. Starting at the location specified by `begin`
    #   the slice continues by adding `stride` to the index until all dimensions are
    #   not less than `end`.
    #   Note that a stride can be negative, which causes a reverse slice.
    #
    #
    # @param input A tensor
    # @param _begin start index
    # @param _end end index
    # @param strides end index
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def strided_slice(input, _begin, _end, strides = nil, name: nil)
      _op(:strided_slice, input, _begin, _end, strides, name: name)
    end

    ##
    # Returns x - y element-wise.
    #
    # This operation supports broadcasting
    #
    # @param input_a tensor X
    # @param input_b tensor Y
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def sub(input_a, input_b, name: nil)
      input_a, input_b = apply_data_type_coercion(input_a, input_b)
      _op(:sub, input_a, input_b, name: name)
    end

    alias_method :subtract, :sub

    ##
    # Computes the sum of elements across dimensions of a tensor.
    # Reduces input_tensor along the dimensions given in axis. Unless keepdims is true, the rank of the
    # tensor is reduced by 1 for each entry in axis. If keepdims is true, the reduced dimensions are
    # retained with length 1.
    # If axis has no entries, all dimensions are reduced, and a tensor with a single element is returned.
    #
    #
    # @param input_a tensor X
    # @param axis_p tensor X (of type INTEGER_TYPES)
    #
    # Options:
    # @option axis axis
    # @option name Optional name
    # @option keepdims If true, retains reduced dimensions with length 1. default (false)
    # @return Tensor
    def sum(input_a, axis_p = nil, axis: nil, name: nil, keepdims: false)
      check_allowed_types(axis_p, TensorStream::Ops::INTEGER_TYPES)
      input_a = TensorStream.convert_to_tensor(input_a)
      return input_a if input_a.shape.scalar?
      axis_p = axis_p || axis
      axis_p = cast_axis(input_a, axis_p)
      _op(:sum, input_a, axis_p, name: name, keepdims: keepdims)
    end

    alias_method :reduce_sum, :sum

    ##
    # Computes tan of input element-wise.
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def tan(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:tan, input_a, name: name)
    end

    ##
    # Computes tanh of input element-wise.
    #
    #
    # @param input_a tensor X (of type FLOATING_POINT_TYPES)
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def tanh(input_a, name: nil)
      check_allowed_types(input_a, TensorStream::Ops::FLOATING_POINT_TYPES)
      _op(:tanh, input_a, name: name)
    end

    ##
    # Constructs a tensor by tiling a given tensor.
    # This operation creates a new tensor by replicating input multiples times.
    # The output tensor's i'th dimension has input.dims(i) * multiples[i] elements,
    # and the values of input are replicated multiples[i] times along the 'i'th dimension. For example, tiling [a b c d] by [2] produces [a b c d a b c d].
    #
    #
    # @param input A tensor
    # @param multiples Must be one of the following types: int32, int64. 1-D. Length must be the same as the number of dimensions in input
    #
    # Options:
    # @option name Optional name
    # @return Tensor
    def tile(input, multiples, name: nil)
      _op(:tile, input, multiples, name: name)
    end

    ##
    # Finds values and indices of the `k` largest entries for the last dimension.
    #
    #
    # @param input 1-D or higher `Tensor` with last dimension at least `k`.
    # @param k 0-D `int32` `Tensor`.  Number of top elements to look for along the last dimension (along each row for matrices)
    #
    # Options:
    # @option sorted If true the resulting `k` elements will be sorted by the values in descending order. default (true)
    # @option name Optional name
    # @return Tensor
    def top_k(input, k = 1, sorted: true, name: nil)
      result = _op(:top_k, input, k, sorted: sorted, name: name)
      [result[0], result[1]]
    end

    ##
    # Creates a tensor with all elements set to zero
    #
    #
    # @param shape A 1-D integer Tensor or ruby array. The shape of the output tensor.
    #
    # Options:
    # @option dtype Optional name default (:float32)
    # @option name Optional name
    # @return Tensor
    def zeros(shape, dtype: :float32, name: nil)
      _op(:zeros, shape, dtype: dtype, name: name)
    end

  end
end


module TensorStream
  class Freezer
    include TensorStream::OpHelper

    ##
    # Utility class to convert variables to constants for production deployment
    #
    def convert(session, checkpoint_folder, output_file)
      model_file = File.join(checkpoint_folder, "model.yaml")
      TensorStream.graph.as_default do |current_graph|
        YamlLoader.new.load_from_string(File.read(model_file))
        saver = TensorStream::Train::Saver.new
        saver.restore(session, checkpoint_folder)

        # collect all assign ops and remove them from the graph
        remove_nodes = Set.new(current_graph.nodes.values.select { |op| op.is_a?(TensorStream::Operation) && op.operation == :assign }.map { |op| op.consumers.to_a }.flatten.uniq)

        output_buffer = TensorStream::Yaml.new.get_string(current_graph) { |graph, node_key|
          node = graph.get_tensor_by_name(node_key)
          case node.operation
          when :variable_v2
            value = node.container
            options = {
              value: value,
              data_type: node.data_type,
              shape: shape_eval(value),
            }
            const_op = TensorStream::Operation.new(current_graph, inputs: [], options: options)
            const_op.name = node.name
            const_op.operation = :const
            const_op.data_type = node.data_type
            const_op.shape = TensorShape.new(shape_eval(value))

            const_op
          when :assign
            nil
          else
            remove_nodes.include?(node.name) ? nil : node
          end
        }
        File.write(output_file, output_buffer)
      end
    end
  end
end


module TensorStream
  # A class that defines a TensorStream graph
  class Graph
    include OpHelper

    attr_accessor :nodes, :collections, :eager_execution, :random_seed, :constants
    attr_reader :node_keys

    def initialize
      @eager_execution = false
      @nodes = {}
      @node_keys = []
      @collections = {
        :"#{GraphKeys::GLOBAL_VARIABLES}" => [],
        :"#{GraphKeys::TRAINABLE_VARIABLES}" => [],
      }
      @constants = {}
    end

    def reset
      @placeholder_counter = 0
      @const_counter = 0
      @var_counter = 0
      @op_counter = 0
      @random_seed = nil
      @nodes = {}
      @node_keys = []
      @collections = {
        :"#{GraphKeys::GLOBAL_VARIABLES}" => [],
        :"#{GraphKeys::TRAINABLE_VARIABLES}" => [],
      }
      @constants = {}
    end

    def as_default
      Thread.current[:tensor_stream_current_graph_queue] ||= []
      Thread.current[:tensor_stream_current_graph_queue] << Graph.get_default_graph

      Thread.current[:tensor_stream_current_graph] = self
      yield(self) if block_given?
      Thread.current[:tensor_stream_current_graph] = Thread.current[:tensor_stream_current_graph_queue].pop
      self
    end

    def name_scope(name = nil)
      Thread.current["ts_graph_#{object_id}"] ||= {}
      Thread.current["ts_graph_#{object_id}"][:current_scope] ||= []
      Thread.current["ts_graph_#{object_id}"][:current_scope] << name

      begin
        yield get_name_scope if block_given?
      ensure
        Thread.current["ts_graph_#{object_id}"][:current_scope].pop
      end
    end

    ##
    # Returns a context manager that specifies the default device to use.
    def device(device_name)
      Thread.current["ts_graph_#{object_id}"] ||= {}
      Thread.current["ts_graph_#{object_id}"][:default_device] ||= []
      Thread.current["ts_graph_#{object_id}"][:default_device] << device_name
      begin
        yield
      ensure
        Thread.current["ts_graph_#{object_id}"][:default_device].pop
      end
    end

    def self.get_default_graph
      Thread.current[:tensor_stream_current_graph] || create_default
    end

    def self.create_default
      Thread.current[:tensor_stream_current_graph] = TensorStream::Graph.new
    end

    def get_collection(name, _options = {})
      @collections[name.to_sym]
    end

    def add_to_collection(collection_name, val)
      @collections[collection_name.to_sym] ||= []
      @collections[collection_name.to_sym] << val
    end

    def add_node(node, name = nil)
      raise "Placeholder cannot be used when eager_execution is enabled" if @eager_execution && node.is_a?(Placeholder)

      if name.nil?
        node.name = if @nodes[node.name]
          uniqunify(node.name)
        else
          node.name
        end
      end

      node.device = get_device_scope
      @node_keys << node.name
      @nodes[node.name] = node
      @constants[node.name] = node if node.is_const

      node.send(:propagate_outputs)
      node.send(:propagate_consumer, node)
    end

    def node_added?(name)
      @nodes.key?(name)
    end

    def get_node(name)
      @nodes[name]
    end

    def get_tensor_by_name(name)
      raise TensorStream::KeyError, "#{name} not found" unless @nodes.key?(name)

      get_node(name)
    end

    def [](name)
      get_node(name)
    end

    def add_node!(name, node)
      @nodes[name] = node
      node
    end

    def add_op(operation, *args)
      options = if args.last.is_a?(Hash)
        args.pop
      else
        {}
      end

      inputs = args.map { |i| TensorStream.convert_to_tensor(i) }.map { |i| i ? i.op : nil }

      new_op = Operation.new(self, inputs: inputs, options: options)
      new_op.source = format_source(caller_locations)
      new_op.operation = operation
      new_op.shape = TensorShape.new(TensorStream::InferShape.infer_shape(new_op))
      new_op.rank = new_op.shape.rank
      new_op.name = options[:internal_name] || [get_name_scope, options[:name] || set_operation_name(new_op)].compact.reject(&:empty?).join("/")
      new_op.internal = options[:internal]

      new_op.data_type = new_op.set_data_type(options[:data_type])
      new_op.is_const = new_op.infer_const

      new_op.given_name = new_op.name

      new_op
    end

    def add_op!(operation, *args)
      add_op(operation, *args).tap { |node| add_node(node) }
    end

    def set_operation_name(op)
      op.operation.to_s
    end

    def add_variable(node, options = {})
      scope = _variable_scope

      raise "duplicate variable detected #{node.name} and reuse=false in current scope" if @nodes[node.name] && !scope.reuse
      return @nodes[node.name] if @nodes[node.name]
      raise "shape is not declared for #{node.name}" if node.shape.nil?

      if !options[:collections].nil? && !options[:collections].empty?
        options[:collections] = [options[:collections]] unless options[:collections].is_a?(Array)
        options[:collections].each { |coll| add_to_collection(coll, node) }
      end

      add_to_collection(GraphKeys::GLOBAL_VARIABLES, node)
      add_to_collection(GraphKeys::TRAINABLE_VARIABLES, node) if node.trainable?

      node
    end

    def add_variable!(node, options = {})
      node = add_variable(node, options)
      op = Graph.get_default_graph.add_op!(:variable_v2, container: node, internal_name: node.name, shape: options[:shape], data_type: options[:data_type])
      node.name = op.name
      op
    end

    def control_dependencies(control_inputs = [])
      Thread.current["ts_graph_#{object_id}"] ||= {}
      Thread.current["ts_graph_#{object_id}"][:control_dependencies] ||= []
      Thread.current["ts_graph_#{object_id}"][:control_dependencies] << Graph.get_default_graph.add_op!(:no_op, *control_inputs)
      begin
        yield
      ensure
        Thread.current["ts_graph_#{object_id}"][:control_dependencies].pop
      end
    end

    def enable_eager_execution
      @eager_execution = true
    end

    def disable_eager_execution
      @eager_execution = false
    end

    def executing_eagerly?
      @eager_execution
    end

    def get_operation_counter
      @op_counter ||= 0

      name = @op_counter.zero? ? "" : "_#{@op_counter}"

      @op_counter += 1

      name
    end

    def get_placeholder_counter
      @placeholder_counter ||= 0
      @placeholder_counter += 1

      return "" if @placeholder_counter == 1

      "_#{@placeholder_counter}"
    end

    def get_var_counter
      @var_counter ||= 0
      @var_counter += 1

      return "" if @var_counter == 1
      "_#{@var_counter}"
    end

    def get_const_counter
      @const_counter ||= 0

      name = @const_counter.zero? ? "" : "_#{@const_counter}"

      @const_counter += 1
      name
    end

    def get_name_scope
      graph_thread_storage = Thread.current["ts_graph_#{object_id}"]
      return nil if graph_thread_storage.nil? || graph_thread_storage[:current_scope].nil?

      graph_thread_storage[:current_scope].join("/")
    end

    def get_dependency_scope
      graph_thread_storage = Thread.current["ts_graph_#{object_id}"]
      return nil if graph_thread_storage.nil? || graph_thread_storage[:control_dependencies].nil?
      graph_thread_storage[:control_dependencies].last
    end

    def get_device_scope
      graph_thread_storage = Thread.current["ts_graph_#{object_id}"]
      return :default if graph_thread_storage.nil? || graph_thread_storage[:default_device].nil?
      graph_thread_storage[:default_device].last
    end

    def as_graph_def
      TensorStream::Pbtext.new.get_string(self)
    end

    def self.parse_from_string(buffer)
      builder = TensorStream::GraphBuilder.new(Graph.new)
      builder.build(buffer)
    end

    def graph_def_versions
      "producer: 26"
    end

    protected

    def _variable_scope
      return VariableScope.new(name: "", reuse: false, initializer: nil) if Thread.current[:tensor_stream_variable_scope].nil? || Thread.current[:tensor_stream_variable_scope].empty?
      scope = Thread.current[:tensor_stream_variable_scope].last
      scope
    end

    def uniqunify(name)
      counter = 0
      new_name = name
      Kernel.loop do
        counter += 1
        new_name = "#{name}_#{counter}"

        break unless @nodes.key?(new_name)
      end
      new_name
    end
  end
end


module TensorStream
  ##
  # Class for deserialization from a YAML file
  class YamlLoader
    def initialize(graph = nil)
      @graph = graph || TensorStream.get_default_graph
    end

    ##
    # Loads a model Yaml file and builds the model from it
    #
    # Args:
    # filename: String - Location of Yaml file
    #
    # Returns: Graph where model is restored to
    def load_from_file(filename)
      load_from_string(File.read(filename))
    end

    ##
    # Loads a model Yaml file and builds the model from it
    #
    # Args:
    # buffer: String - String in Yaml format of the model
    #
    # Returns: Graph where model is restored to
    def load_from_string(buffer)
      serialized_ops = YAML.safe_load(buffer, [Symbol], [], true)
      serialized_ops.each do |op_def|
        inputs = op_def[:inputs].map { |i| @graph.get_tensor_by_name(i) }
        options = {}

        new_var = nil
        if op_def.dig(:attrs, :container)
          new_var = Variable.new(op_def.dig(:attrs, :data_type))
          var_shape = op_def.dig(:attrs, :container, :shape)
          var_options = op_def.dig(:attrs, :container, :options)
          var_options[:name] = op_def[:name]

          new_var.prepare(var_shape.size, var_shape, TensorStream.get_variable_scope, var_options)
          options[:container] = new_var

          @graph.add_variable(new_var, var_options)
        end

        new_op = Operation.new(@graph, inputs: inputs, options: op_def[:attrs].merge(options))
        new_op.operation = op_def[:op].to_sym
        new_op.name = op_def[:name]
        new_op.shape = TensorShape.new(TensorStream::InferShape.infer_shape(new_op))
        new_op.rank = new_op.shape.rank
        new_op.data_type = new_op.set_data_type(op_def.dig(:attrs, :data_type))
        new_op.is_const = new_op.infer_const
        new_op.given_name = new_op.name
        new_var.op = new_op if new_var

        @graph.add_node(new_op)
      end
      @graph
    end
  end
end


module TensorStream
  module Utils
    def float32
      Types.float32
    end

    def graph
      TensorStream::Graph.new
    end

    def get_default_graph
      TensorStream::Graph.get_default_graph
    end

    def reset_default_graph
      TensorStream::Graph.get_default_graph.reset
    end

    def enable_eager_execution
      TensorStream::Graph.get_default_graph.enable_eager_execution
    end

    def disable_eager_execution
      TensorStream::Graph.get_default_graph.disable_eager_execution
    end

    def executing_eagerly?
      TensorStream::Graph.get_default_graph.executing_eagerly?
    end

    ##
    # List available evaluators + devices in the current local environment
    # Returns:
    # - An array containing the names of those devices
    def list_local_devices
      local_name = "job:localhost"
      TensorStream::Evaluator.evaluators.collect { |k, v|
        v[:class].query_supported_devices.collect do |device_str|
          [local_name, "ts:#{k}:#{device_str.name}"].join("/")
        end
      }.flatten
    end

    ##
    # Creates a variable
    # A variable maintains state across sessions
    def variable(value, name: nil, initializer: nil, graph: nil, dtype: nil, trainable: true)
      op = Graph.get_default_graph.add_op(:assign, nil, value)
      common_options = {
        initializer: initializer || op,
        name: name,
        graph: graph,
        dtype: dtype,
        trainable: trainable,
      }
      tensor = if value.is_a?(String)
        i_var(dtype || :string, 0, [], get_variable_scope, common_options)
      elsif value.is_a?(Integer)
        i_var(dtype || :int32, 0, [], get_variable_scope, common_options)
      elsif value.is_a?(Float)
        i_var(dtype || :float32, 0, [], get_variable_scope, common_options)
      else
        i_var(dtype || :float32, 0, nil, get_variable_scope, common_options)
      end
      op.set_input(0, tensor.op)
      Graph.get_default_graph.add_node(op)
      tensor
    end

    ##
    # Defines a variable context manager
    def variable_scope(scope = nil, default_name = nil, reuse: nil, initializer: nil)
      Thread.current[:tensor_stream_variable_scope] ||= [VariableScope.new]

      # uniquenifier
      if scope.nil? && default_name
        same_names = get_variable_scope.used_names.select { |s| s.start_with?(default_name) }
        new_name = default_name
        index = 1
        while same_names.include?(new_name)
          new_name = "#{default_name}_#{index}"
          index += 1
        end
        scope = new_name
      end

      variable_scope = VariableScope.new(name: scope, reuse: reuse, initializer: initializer)
      get_variable_scope.register_name(scope || "")
      Thread.current[:tensor_stream_variable_scope] << variable_scope
      scope_name = __v_scope_name
      if block_given?
        begin
          TensorStream.get_default_graph.name_scope(scope) do
            yield(scope_name)
          end
        ensure
          Thread.current[:tensor_stream_variable_scope].pop
        end
      else
        variable_scope
      end
    end

    def device(device_uri, &block)
      get_default_graph.device(device_uri, &block)
    end

    def name_scope(name, default_name = nil, default: nil, values: nil)
      if values
        graph_count = values.select { |v| v.is_a?(Tensor) }.map(&:graph).map(&:object_id).uniq.size
        raise "values are not on the same graph" if graph_count > 1
      end

      get_default_graph.name_scope(name || default_name || default) do |scope|
        yield scope if block_given?
      end
    end

    def get_variable_scope
      unless Thread.current[:tensor_stream_variable_scope]
        variable_scope = VariableScope.new
        Thread.current[:tensor_stream_variable_scope] = [variable_scope]
        return variable_scope
      end

      Thread.current[:tensor_stream_variable_scope].last
    end

    def __v_scope_name
      Thread.current[:tensor_stream_variable_scope].map(&:name).compact.reject(&:empty?).join("/")
    end

    ##
    # Creates a session context where operations can be executed
    #
    # Args:
    #     evaluator: Specific evaluator to use, otherwise the best evaluator will automatically be determined
    #
    # Options:
    #     thread_pool_class: Class to use to manage thread pooling
    #     log_device_placement: Show assigned device/evalutor for each tensor op
    #     profile_enabled: Log performance metrics for each operation
    def session(evaluator = nil, thread_pool_class: Concurrent::ImmediateExecutor, log_device_placement: false, profile_enabled: false)
      session = TensorStream::Session.new(evaluator, thread_pool_class: thread_pool_class, log_device_placement: log_device_placement, profile_enabled: profile_enabled)
      yield session if block_given?

      session
    end

    def colocate_with(op, ignore_existing: false)
      # noop for now
      yield
    end

    def program
      yield self
    end

    def layers
      TensorStream::Layers
    end

    def constant(value, dtype: nil, shape: nil, internal: false, name: "Const")
      shared_options = {const: true, value: value, name: name, internal: internal}

      if value.is_a?(Float)
        TensorStream::Constant.new(dtype || :float32, 0, shape || [], shared_options)
      elsif value.is_a?(Integer)
        TensorStream::Constant.new(dtype || :int32, 0, shape || [], shared_options)
      elsif value.is_a?(String)
        TensorStream::Constant.new(dtype || :string, 0, shape || [], shared_options)
      elsif !!value == value
        TensorStream::Constant.new(dtype || :boolean, 0, shape || [], shared_options)
      elsif value.is_a?(Array)
        dimension = shape || shape_eval(value)
        rank = dimension.size
        TensorStream.check_if_dense(value)

        cur_dtype = dtype || Tensor.detect_type(value.flatten.last)
        value = Tensor.cast_dtype(value, cur_dtype) unless dtype.nil?

        shared_options[:value] = value
        TensorStream::Constant.new(cur_dtype, rank, dimension, shared_options)
      end
    end

    def group(inputs, name: nil)
      TensorStream::ControlFlow.new(:group, inputs, nil, name: name)
    end

    def dynamic_stitch(indices, data, name: nil)
      TensorStream::DynamicStitch.new(:dynamic_stitch, [indices, data], name: name)
    end

    def get_variable(name, dtype: nil, shape: nil, initializer: nil, trainable: true, collections: nil, validate_shape: false)
      get_variable_scope.get_variable(name, dtype: dtype, shape: shape, initializer: initializer, trainable: trainable, collections: collections)
    end

    def get_collection(name, options = {})
      Graph.get_default_graph.get_collection(name, options)
    end

    def assign(ref, value, name: nil)
      raise "#{ref.name} not a variable" unless ref.is_a?(Variable)
      ref.assign(value, name: name)
    end

    ##
    # Inserts a placeholder for a tensor that will be always fed.
    def placeholder(dtype, shape: nil, name: nil)
      TensorStream::Placeholder.new(dtype, nil, shape, name: name)
    end

    def global_variables_initializer
      TensorStream::Variable.global_variables_initializer
    end

    def train
      TensorStream::Trainer
    end

    def math
      TensorStream::Maths
    end

    def image
      TensorStream::Images
    end

    def trainable_variables
      TensorStream.get_default_graph.get_collection(TensorStream::GraphKeys::TRAINABLE_VARIABLES)
    end

    ##
    # Sets random seed to use for the default graph
    def set_random_seed(seed)
      TensorStream.get_default_graph.random_seed = seed
    end

    def control_dependencies(control_inputs, &block)
      TensorStream.get_default_graph.control_dependencies(control_inputs, &block)
    end

    def convert_to_tensor(value, dtype: nil, name: nil)
      return value if value.is_a?(Tensor)
      return convert_to_tensor(value.call) if value.is_a?(Proc)
      # raise "Invalid tensor value" if value.nil?

      if value.is_a?(Array) && value.detect { |v| v.is_a?(Tensor) }
        return TensorStream.stack(value) if value.size > 1

        return TensorStream.expand_dims(value[0], 0)
      end

      if value.is_a?(TensorShape)
        value = value.shape
      end

      check_if_dense(value)
      i_cons(value, dtype: dtype || Tensor.detect_type(value), name: name)
    end

    ##
    # Check to make sure passed array is dense
    #
    def check_if_dense(value, expected_shape = nil)
      return unless value.is_a?(Array)
      return if value.empty?

      expected_shape ||= shape_eval(value)

      s = expected_shape.shift
      raise TensorStream::ValueError, "Argument must be a dense tensor: #{value}, expected size #{s} got #{value.size}" if value.size != s

      return if expected_shape.empty?

      value.each do |item|
        check_if_dense(item, expected_shape.dup)
      end
    end

    def check_allowed_types(input, types)
      return input unless input.is_a?(Tensor)
      return input if input.data_type.nil?

      raise "#{input.source}: Parameter data type #{input.data_type} passed not in #{types.join(",")}" unless types.include?(input.data_type.to_sym)
    end

    def check_data_types(*args)
      unique_types = args.select { |a| a.is_a?(Tensor) }. map { |a| DataTypeUtils.norm_dtype(a.data_type) }.uniq

      if unique_types.size > 1
        raise TensorStream::ValueError, "Value Error: Tensor conversion requested dtypes are different -> #{unique_types}"
      end

      unique_types.first
    end

    ##
    # Auto cast ruby constant data types to the same
    # tensor types of other operands
    def apply_data_type_coercion(*args)
      coerced_type = check_data_types(*args)
      args.map { |a| a.is_a?(Tensor) ? a : convert_to_tensor(a, dtype: coerced_type) }
    end
  end
end


module PgSlice
  module Helpers
    SQL_FORMAT = {
      day: "YYYYMMDD",
      month: "YYYYMM",
      year: "YYYY"
    }

    protected

    # output

    def log(message = nil)
      error message
    end

    def log_sql(message = nil)
      say message
    end

    def abort(message)
      raise Thor::Error, message
    end

    # database connection

    def connection
      @connection ||= begin
        url = options[:url] || ENV["PGSLICE_URL"]
        abort "Set PGSLICE_URL or use the --url option" unless url

        uri = URI.parse(url)
        params = CGI.parse(uri.query.to_s)
        # remove schema
        @schema = Array(params.delete("schema") || "public")[0]
        uri.query = URI.encode_www_form(params)

        ENV["PGCONNECT_TIMEOUT"] ||= "1"
        PG::Connection.new(uri.to_s)
      end
    rescue PG::ConnectionBad => e
      abort e.message
    rescue URI::InvalidURIError
      abort "Invalid url"
    end

    def schema
      connection # ensure called first
      @schema
    end

    def execute(query, params = [])
      connection.exec_params(query, params).to_a
    end

    def run_queries(queries)
      connection.transaction do
        execute("SET LOCAL client_min_messages TO warning") unless options[:dry_run]
        log_sql "BEGIN;"
        log_sql
        run_queries_without_transaction(queries)
        log_sql "COMMIT;"
      end
    end

    def run_query(query)
      log_sql query
      unless options[:dry_run]
        begin
          execute(query)
        rescue PG::ServerError => e
          abort("#{e.class.name}: #{e.message}")
        end
      end
      log_sql
    end

    def run_queries_without_transaction(queries)
      queries.each do |query|
        run_query(query)
      end
    end

    def server_version_num
      execute("SHOW server_version_num")[0]["server_version_num"].to_i
    end

    # helpers

    def sql_date(time, cast, add_cast = true)
      if cast == "timestamptz"
        fmt = "%Y-%m-%d %H:%M:%S UTC"
      else
        fmt = "%Y-%m-%d"
      end
      str = "'#{time.strftime(fmt)}'"
      add_cast ? "#{str}::#{cast}" : str
    end

    def name_format(period)
      case period.to_sym
      when :day
        "%Y%m%d"
      when :month
        "%Y%m"
      else
        "%Y"
      end
    end

    def partition_date(partition, name_format)
      DateTime.strptime(partition.name.split("_").last, name_format)
    end

    def round_date(date, period)
      date = date.to_date
      case period.to_sym
      when :day
        date
      when :month
        Date.new(date.year, date.month)
      else
        Date.new(date.year)
      end
    end

    def assert_table(table)
      abort "Table not found: #{table}" unless table.exists?
    end

    def assert_no_table(table)
      abort "Table already exists: #{table}" if table.exists?
    end

    def advance_date(date, period, count = 1)
      date = date.to_date
      case period.to_sym
      when :day
        date.next_day(count)
      when :month
        date.next_month(count)
      else
        date.next_year(count)
      end
    end

    def quote_ident(value)
      PG::Connection.quote_ident(value)
    end

    def quote_table(table)
      table.quote_table
    end

    def quote_no_schema(table)
      quote_ident(table.name)
    end

    def create_table(name)
      if name.include?(".")
        schema, name = name.split(".", 2)
      else
        schema = self.schema
      end
      Table.new(schema, name)
    end

    def make_index_def(index_def, table)
      index_def.sub(/ ON \S+ USING /, " ON #{quote_table(table)} USING ").sub(/ INDEX .+ ON /, " INDEX ON ") + ";"
    end

    def make_fk_def(fk_def, table)
      "ALTER TABLE #{quote_table(table)} ADD #{fk_def};"
    end
  end
end


module PgSlice
  class Table
    attr_reader :schema, :name

    def initialize(schema, name)
      @schema = schema
      @name = name
    end

    def to_s
      [schema, name].join(".")
    end

    def exists?
      execute("SELECT COUNT(*) FROM pg_catalog.pg_tables WHERE schemaname = $1 AND tablename = $2", [schema, name]).first["count"].to_i > 0
    end

    def columns
      execute("SELECT column_name FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2", [schema, name]).map{ |r| r["column_name"] }
    end

    # http://www.dbforums.com/showthread.php?1667561-How-to-list-sequences-and-the-columns-by-SQL
    def sequences
      query = <<-SQL
        SELECT
          a.attname as related_column,
          s.relname as sequence_name
        FROM pg_class s
          JOIN pg_depend d ON d.objid = s.oid
          JOIN pg_class t ON d.objid = s.oid AND d.refobjid = t.oid
          JOIN pg_attribute a ON (d.refobjid, d.refobjsubid) = (a.attrelid, a.attnum)
          JOIN pg_namespace n ON n.oid = s.relnamespace
        WHERE s.relkind = 'S'
          AND n.nspname = $1
          AND t.relname = $2
      SQL
      execute(query, [schema, name])
    end

    def foreign_keys
      execute("SELECT pg_get_constraintdef(oid) FROM pg_constraint WHERE conrelid = #{regclass} AND contype ='f'").map { |r| r["pg_get_constraintdef"] }
    end

    # http://stackoverflow.com/a/20537829
    def primary_key
      query = <<-SQL
        SELECT
          pg_attribute.attname,
          format_type(pg_attribute.atttypid, pg_attribute.atttypmod)
        FROM
          pg_index, pg_class, pg_attribute, pg_namespace
        WHERE
          nspname = $1 AND
          relname = $2 AND
          indrelid = pg_class.oid AND
          pg_class.relnamespace = pg_namespace.oid AND
          pg_attribute.attrelid = pg_class.oid AND
          pg_attribute.attnum = any(pg_index.indkey) AND
          indisprimary
      SQL
      execute(query, [schema, name]).map { |r| r["attname"] }
    end

    def index_defs
      execute("SELECT pg_get_indexdef(indexrelid) FROM pg_index WHERE indrelid = #{regclass} AND indisprimary = 'f'").map { |r| r["pg_get_indexdef"] }
    end

    def quote_table
      [quote_ident(schema), quote_ident(name)].join(".")
    end

    def intermediate_table
      self.class.new(schema, "#{name}_intermediate")
    end

    def retired_table
      self.class.new(schema, "#{name}_retired")
    end

    def trigger_name
      "#{name}_insert_trigger"
    end

    def column_cast(column)
      data_type = execute("SELECT data_type FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 AND column_name = $3", [schema, name, column])[0]["data_type"]
      data_type == "timestamp with time zone" ? "timestamptz" : "date"
    end

    def max_id(primary_key, below: nil, where: nil)
      query = "SELECT MAX(#{quote_ident(primary_key)}) FROM #{quote_table}"
      conditions = []
      conditions << "#{quote_ident(primary_key)} <= #{below}" if below
      conditions << where if where
      query << " WHERE #{conditions.join(" AND ")}" if conditions.any?
      execute(query)[0]["max"].to_i
    end

    def min_id(primary_key, column, cast, starting_time, where)
      query = "SELECT MIN(#{quote_ident(primary_key)}) FROM #{quote_table}"
      conditions = []
      conditions << "#{quote_ident(column)} >= #{sql_date(starting_time, cast)}" if starting_time
      conditions << where if where
      query << " WHERE #{conditions.join(" AND ")}" if conditions.any?
      (execute(query)[0]["min"] || 1).to_i
    end

    def partitions
      query = <<-SQL
        SELECT
          nmsp_child.nspname  AS schema,
          child.relname       AS name
        FROM pg_inherits
          JOIN pg_class parent            ON pg_inherits.inhparent = parent.oid
          JOIN pg_class child             ON pg_inherits.inhrelid   = child.oid
          JOIN pg_namespace nmsp_parent   ON nmsp_parent.oid  = parent.relnamespace
          JOIN pg_namespace nmsp_child    ON nmsp_child.oid   = child.relnamespace
        WHERE
          nmsp_parent.nspname = $1 AND
          parent.relname = $2
      SQL
      execute(query, [schema, name]).map { |r| Table.new(r["schema"], r["name"]) }
    end

    def fetch_comment
      execute("SELECT obj_description(#{regclass}) AS comment")[0]
    end

    def fetch_trigger(trigger_name)
      execute("SELECT obj_description(oid, 'pg_trigger') AS comment FROM pg_trigger WHERE tgname = $1 AND tgrelid = #{regclass}", [trigger_name])[0]
    end

    # legacy
    def fetch_settings(trigger_name)
      needs_comment = false
      trigger_comment = fetch_trigger(trigger_name)
      comment = trigger_comment || fetch_comment
      if comment
        field, period, cast, version = comment["comment"].split(",").map { |v| v.split(":").last } rescue []
        version = version.to_i if version
      end

      unless period
        needs_comment = true
        function_def = execute("SELECT pg_get_functiondef(oid) FROM pg_proc WHERE proname = $1", [trigger_name])[0]
        return [] unless function_def
        function_def = function_def["pg_get_functiondef"]
        sql_format = SQL_FORMAT.find { |_, f| function_def.include?("'#{f}'") }
        return [] unless sql_format
        period = sql_format[0]
        field = /to_char\(NEW\.(\w+),/.match(function_def)[1]
      end

      # backwards compatibility with 0.2.3 and earlier (pre-timestamptz support)
      unless cast
        cast = "date"
        # update comment to explicitly define cast
        needs_comment = true
      end

      version ||= trigger_comment ? 1 : 2
      declarative = version > 1

      [period, field, cast, needs_comment, declarative, version]
    end

    protected

    def execute(*args)
      PgSlice::CLI.instance.send(:execute, *args)
    end

    def quote_ident(value)
      PG::Connection.quote_ident(value)
    end

    def regclass
      "'#{quote_table}'::regclass"
    end

    def sql_date(time, cast, add_cast = true)
      if cast == "timestamptz"
        fmt = "%Y-%m-%d %H:%M:%S UTC"
      else
        fmt = "%Y-%m-%d"
      end
      str = "'#{time.strftime(fmt)}'"
      add_cast ? "#{str}::#{cast}" : str
    end
  end
end