# require 'pry-byebug'
module TensorStream
  # various monkey patches to FixNum types
  module MonkeyPatch
    def self.included(klass)
      ops = if klass == Array
              {:+ => 'add', :- => 'sub', :* => 'mul'}
            else
              {:+ => 'add', :- => 'sub', :/ => 'div', :% => 'mod', :* => 'mul', :** => 'pow' }
            end

      ops.each do |m, name|
        klass.send(:alias_method, :"_tensor_stream_#{name}_orig", m)
        klass.send(:remove_method, m)
      end
    end

    def shape
      TensorStream.shape_eval(self)
    end

    def t(name = nil, dtype: nil)
      TensorStream.convert_to_tensor(self, name: name, dtype: dtype)
    end

    def +(other)
      if other.is_a?(TensorStream::Tensor)
        TensorStream.convert_to_tensor(self, dtype: other.data_type) + other
      else
        _tensor_stream_add_orig(other)
      end
    end

    def -(other)
      if other.is_a?(TensorStream::Tensor)
        TensorStream.convert_to_tensor(self, dtype: other.data_type) - other
      else
        _tensor_stream_sub_orig(other)
      end
    end

    def *(other)
      if other.is_a?(TensorStream::Tensor)
        TensorStream.convert_to_tensor(self, dtype: other.data_type) * other
      else
        _tensor_stream_mul_orig(other)
      end
    end

    def /(other)
      if other.is_a?(TensorStream::Tensor)
        TensorStream.convert_to_tensor(self, dtype: other.data_type) * other
      else
        _tensor_stream_div_orig(other)
      end
    end

    def %(other)
      if other.is_a?(TensorStream::Tensor)
        TensorStream.convert_to_tensor(self, dtype: other.data_type) % other
      else
        _tensor_stream_mod_orig(other)
      end
    end

    def **(other)
      if other.is_a?(TensorStream::Tensor)
        TensorStream.convert_to_tensor(self, dtype: other.data_type)**other
      else
        _tensor_stream_pow_orig(other)
      end
    end
  end
end