module TensorStream
  # TensorStream class that defines an operation
  class Operation < Tensor
    attr_accessor :name, :operation, :inputs, :rank, :options
    attr_reader :outputs

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

      inputs = args

      setup_initial_state(options)

      @operation = operation
      @rank = options[:rank] || 0
      @name = [@graph.get_name_scope, options[:name] || set_name].compact.reject(&:empty?).join('/')
      @internal = options[:internal]
      @given_name = @name

      @options = options

      @inputs = inputs.map { |i| options[:preserve_params_type] ? i : TensorStream.convert_to_tensor(i) }
      @data_type = set_data_type(options[:data_type])
      @is_const = infer_const
      @shape = TensorShape.new(infer_shape)
      @graph.add_node(self)
    end

    def to_s
      @name
    end

    def to_h
      {
        op: operation,
        name: name,
        operands: hashify_tensor(inputs)
      }
    end

    def infer_const
      return false if breakpoint
      case operation
      when :random_standard_normal, :random_uniform, :glorot_uniform, :print
        false
      else
        non_const = @inputs.compact.find { |input| !input.is_const }
        non_const ? false : true
      end
    end

    def set_data_type(passed_data_type)
      case operation
      when :fill
        @inputs[1].data_type
      when :greater, :less, :equal, :not_equal, :greater_equal, :less_equal, :logical_and
        :boolean
      when :shape, :rank
        options[:out_type] || :int32
      when :random_standard_normal, :random_uniform, :glorot_uniform
        passed_data_type || :float32
      when :index
        if @inputs[0].is_a?(ControlFlow)

          if @inputs[1].is_const
            @inputs[0].inputs[@inputs[1].value].data_type
          else
            :unknown
          end
        else
          @inputs[0].data_type
        end
      else
        return passed_data_type if passed_data_type

        if @inputs[0]
          @inputs[0].data_type
        elsif @inputs[1]
          @inputs[1].data_type
        else
          :unknown
        end
      end
    end

    def to_math(name_only = false, max_depth = 99, _cur_depth = 0)
      return @name if max_depth.zero?

      sub_input = auto_math(inputs[0], name_only, max_depth - 1, _cur_depth + 1)
      sub_input2 = auto_math(inputs[1], name_only, max_depth - 1, _cur_depth + 1) if inputs[1]

      out = case operation
      when :argmax
        "argmax(#{sub_input},#{options[:axis]})"
      when :negate
        "-#{sub_input}"
      when :index
        "#{sub_input}[#{sub_input2}]"
      when :slice
        "#{sub_input}[#{sub_input2}]"
      when :assign_sub
        "(#{inputs[0] ? inputs[0].name : 'self'} -= #{auto_math(inputs[1], name_only, 1)})"
      when :assign_add
        "(#{inputs[0] ? inputs[0].name : 'self'} += #{auto_math(inputs[1], name_only, 1)})"
      when :assign
        "(#{inputs[0] ? inputs[0].name : 'self'} = #{auto_math(inputs[1], name_only, 1)})"
      when :sin, :cos, :tanh
        "#{operation}(#{sub_input})"
      when :add
        "(#{sub_input} + #{sub_input2})"
      when :sub
        "(#{sub_input} - #{sub_input2})"
      when :pow
        "(#{sub_input}^#{sub_input2})"
      when :div
        "(#{sub_input} / #{sub_input2})"
      when :mul
        if auto_math(inputs[0]) == 1
          sub_input2
        elsif auto_math(inputs[1]) == 1
          sub_input
        else
          "(#{sub_input} * #{sub_input2})"
        end
      when :sum
        "sum(|#{sub_input}|,  axis=#{sub_input2})"
      when :mean
        "mean(|#{sub_input}|, axis=#{sub_input2})"
      when :prod
        "prod(|#{sub_input}|,  axis=#{sub_input2})"
      when :gradients
        "gradient(#{sub_input})"
      when :stop_gradient
        sub_input
      when :mat_mul
        "#{sub_input}.matmul(#{sub_input2})"
      when :eye
        "eye(#{sub_input})"
      when :transpose
        "transpose(#{sub_input})"
      when :shape
        "#{sub_input}.shape"
      when :exp
        "e^#{sub_input})"
      when :ones
        "ones(#{sub_input})"
      when :ones_like
        "ones_like(#{sub_input})"
      when :flow_group
        "flow_group(#{inputs.collect { |i| auto_math(i, name_only, max_depth - 1, _cur_depth) }.join(',')})"
      when :zeros
        "zeros(#{sub_input})"
      when :reshape
        "reshape(#{sub_input},#{sub_input2})"
      when :rank
        "#{sub_input}.rank"
      when :cond
        "(#{auto_math(options[:pred], name_only, max_depth - 1, _cur_depth)} ? #{sub_input} : #{sub_input2})"
      when :less
        "#{sub_input} < #{sub_input2}"
      when :less_equal
        "#{sub_input} <= #{sub_input2}"
      when :greater
        "#{sub_input} > #{sub_input2}"
      when :greater_equal
        "#{sub_input} >= #{sub_input2}"
      when :square
        "#{sub_input}\u00B2"
      when :log
        "log(#{sub_input})"
      when :identity
        "identity(#{sub_input})"
      when :print
        "print(#{sub_input})"
      when :pad
        "pad(#{sub_input},#{auto_math(options[:paddings])})"
      when :equal
        "#{sub_input} == #{sub_input2}"
      when :not_equal
        "#{sub_input} != #{sub_input2}"
      when :logical_and
        "#{sub_input} && #{sub_input2}"
      when :sqrt
        "sqrt(#{sub_input})"
      when :log1p
        "log1p(#{sub_input})"
      when :zeros_like
        "zeros_like(#{sub_input})"
      when :where
        "where(#{auto_math(options[:pred], name_only, max_depth - 1, _cur_depth)}, #{sub_input}, #{sub_input2})"
      when :max
        "max(#{sub_input},#{sub_input2})"
      when :cast
        "cast(#{sub_input}, #{data_type})"
      when :broadcast_transform
        "broadcast_transform(#{sub_input},#{sub_input2})"
      when :broadcast_gradient_args
        "broadcast_transform(#{sub_input},#{sub_input2})"
      else
        "#{operation}(#{sub_input})" if sub_input
        "#{operation}(#{sub_input}, #{sub_input2})" if sub_input && sub_input2
      end
      ["\n",(_cur_depth + 1).times.collect { ' ' }, out].flatten.join
    end

    def run
      eval
    end

    def op
      self
    end

    private

    def infer_shape
      case operation
      when :index
        input_shape = inputs[0].shape.shape
        return nil if input_shape.nil?
        return input_shape[1, input_shape.size]
      when :mean, :prod, :sum
        return [] if inputs[1].nil?
        return nil if inputs[0].nil?
        input_shape = inputs[0].shape.shape
        return nil if input_shape.nil?
        return nil if inputs[1].is_a?(Tensor) && inputs[1].value.nil?

        axis = inputs[1].is_a?(Tensor) ? inputs[1].value : inputs[1]

        axis = [ axis ] unless axis.is_a?(Array)
        return input_shape.each_with_index.map do |s, index|
          next nil if axis.include?(index)
          s
        end.compact
      when :reshape
        new_shape = inputs[1] && inputs[1].value ? inputs[1].value : nil
        return nil if new_shape.nil?

        input_shape = inputs[0].shape.shape
        return new_shape if input_shape.nil?

        return TensorShape.fix_inferred_elements(new_shape, input_shape.reduce(:*))
      when :flow_group
        return []
      when :zeros, :ones, :fill
        return inputs[0] ? inputs[0].value : options[:shape]
      when :zeros_like, :ones_like
        inputs[0].shape.shape
      when :shape
        return inputs[0].shape.shape ? [inputs[0].shape.shape.size] : nil
      when :mat_mul
        shape1 = inputs[0].shape.shape.nil? ? nil : inputs[0].shape.shape[0]
        shape2 = inputs[1].shape.shape.nil? ? nil : inputs[1].shape.shape[1]
        return [shape1, shape2]
      else
        return inputs[0].shape.shape if inputs.size == 1
        if inputs.size == 2 && inputs[0] && inputs[1]
          return TensorShape.infer_shape(inputs[0].shape.shape, inputs[1].shape.shape)
        end
      end

      nil
    end

    def propagate_consumer(consumer)
      super
      @inputs.compact.each do |input|
        if input.is_a?(Array)
          input.flatten.compact.select { |t| t.is_a?(Tensor) }.each do |t|
            next if t.consumers.include?(consumer.name)
            t.send(:propagate_consumer, consumer)
          end
        elsif input.name != name && !input.consumers.include?(consumer.name)
          input.send(:propagate_consumer, consumer)
        end
      end
    end

    def propagate_outputs
      @inputs.compact.each do |input|
        if input.is_a?(Array)
          input.flatten.compact.each do |t|
            t.send(:setup_output, self) if t.is_a?(Tensor)
          end
        else
          input.send(:setup_output, self) if input.is_a?(Tensor) && (input.name != self.name)
        end
      end
    end

    def set_name
      "#{@operation}#{graph.get_operation_counter}:#{@rank}"
    end
  end
end