=begin

=end

module NumRu
  module Misc
    # A Mixin.
    # To be included in a class with multi-dimension indexing support
    # (such as NArray).
    module MD_Iterators

      # Iterator for each sub-array (not each element) specified by
      # dimensions.  
      # 
      # ARGUMENT
      # 
      # <tt>*dims</tt> (integers) : specifies subsets
      # at dimensions specified here with the beginning-to-end selection.
      # For example, [0, 1] to specify the first 2 dimensions (subsets
      # will be 2D then), and [2] to specify the 3rd dimension (subsets
      # will be 1D). Duplication has no effect, so [0,0] and [0] are the
      # same. Also, its order has no effect.  See EXAMPLE below for more.
      # 
      # RETURN VALUE
      # * self
      # 
      # POSSIBLE EXCEPTIONS
      # * exception is raised if ( dims.min<0 || dims.max>=self.rank ).
      # 
      # EXAMPLE
      # 
      # * Suppose that you want to do something with 2D sub-arrays in a
      #   multi-dimension NArray. First, you include this module as follows:
      # 
      #     require "narray"
      #     class NArray
      #        include NumRu::Misc::MD_Iterators
      #     end
      # 
      #   And prepare the array if you have not (here, it is 4D):
      # 
      #     na = NArray.int(10,2,5,2).indgen!
      # 
      #   Then you do the job like this:
      # 
      #     na.each_subary_at_dims(0,2){ |sub|
      #       ...  # do whatever with sub
      #     }
      # 
      #   This is equivalent to the following:
      # 
      #     (0...na.shape[3]).each{|j|
      #       (0...na.shape[1]).each{|i|
      #         sub = na[0..-1, i, 0..-1, j]
      #         ...  # do whatever with sub
      #       }
      #     }
      # 
      #   Note that the loop must be nested 3 times when <tt>na>/tt> is a 5D array,
      #   if the latter approach is used. On the other hand, it will still 
      #   require the same single loop with the former.      
      def each_subary_at_dims( *dims )
	if dims.min<0 || dims.max>=rank
	  raise ArguemntError,"Invalid dims #{dims.inspect} for #{rank}D array"
	end

	loopdims = Array.new
	sh = Array.new
	len = 1
	(0...rank).each{|i| 
	  if !dims.include?(i)
	    loopdims.push(i) 
	    sh.push(shape[i])
	    len *= shape[i]
	  end
	}
	if loopdims.length == 0
	  yield(self)
	  return self
	end
	cs = [1]
	(1...sh.length).each{|i| cs[i] = sh[i-1]*cs[i-1]}
	idx = Array.new
	all = 0..-1
	for i in 0...len do
	  loopdims.each_with_index{|d,j| idx[d] = ( (i/cs[j])%sh[j] )}
	  dims.each{|d| idx[d] = all}
	  sub = self[ *idx ]
	  yield(sub)
	end
	self
      end

    # Like #each_subary_at_dims but the block takes two arguments:
    # subset and the subset specifier (index).
    # 
    # EXAMPLE
    # * Suppose the example above in #each_subary_at_dims (EXAMPLE).
    #   And suppose that you want to overwrite <tt>na</tt> with the result
    #   you get. You can do it like this:
    # 
    #     na.each_subary_at_dims_with_index(0,2){ |sub,idx|
    #       result = (sub + 10) / 2
    #       na[*idx] = result
    #     }
    # 
    #   Here, <tt>idx</tt> is an Array to be fed in the []= or [] methods
    #   with asterisk (ungrouping).
    #
      def each_subary_at_dims_with_index( *dims )
	if dims.min<0 || dims.max>=rank
	  raise ArguemntError,"Invalid dims #{dims.inspect} for #{rank}D array"
	end

	loopdims = Array.new
	sh = Array.new
	len = 1
	(0...rank).each{|i| 
	  if !dims.include?(i)
	    loopdims.push(i) 
	    sh.push(shape[i])
	    len *= shape[i]
	  end
	}
	if loopdims.length == 0
	  yield(self, false)
	  return self
	end
	cs = [1]
	(1...sh.length).each{|i| cs[i] = sh[i-1]*cs[i-1]}
	idx = Array.new
	all = 0..-1
	for i in 0...len do
	  loopdims.each_with_index{|d,j| idx[d] = ( (i/cs[j])%sh[j] )}
	  dims.each{|d| idx[d] = all}
	  sub = self[ *idx ]
	  yield(sub, idx)
	end
	self
      end

    end
  end
end

##################################

if __FILE__ == $0
  require "narray"
  class NArray   # :nodoc:
    include NumRu::Misc::MD_Iterators
  end
  na = NArray.int(10,2,2,2).indgen!
  puts "** test A **"
  na.each_subary_at_dims(0,1){ |sub|
    p sub
  }
  puts "** test B **"
  na.each_subary_at_dims(0,3){ |sub|
    p sub
  }
  puts "** test C **"
  na.each_subary_at_dims(2,1,0){ |sub|     # same as (0,1,2)
    p sub
  }
  puts "** test C **"
  na.each_subary_at_dims(0,1,2,3){ |sub|
    p sub
  }
end
