require "narray"

module NumRu

  # == Overview
  # 
  # Miscellaneous functions and classes to facilitate programming.
  # 
  # == Index
  # 
  # CLASSES
  # 
  # * class KeywordOpt : supports keyword arguments with default values.
  # * class NArray (http://masa16.github.io/narray/index.ja.html) : used in EMath
  # 
  # MODULES
  # 
  # * module MD_Iterators : A Mixin for classes with
  #   multi-dimension indexing support (such as NArray).
  # * module EMath :
  #   To be included instead of the Math predefined module (or NMath in NArray).
  #   Unlike Math and NMath, EMath handles unknown classes by calling its
  #   native instance method (assuming the same name).
  # 
  # 

   module Misc
      module_function

      #  Check the consistency of array shapes (multi-dim such as NArray).
      #  Exception is raised if inconsistent.
      # 
      #  ARGUMENTS
      #  * cshapes (String) : description of the shapes of the args.
      #    Delimited by one-or-more spaces between arrays,
      #    and the shape of each array is delimited by a comma. The lengths are 
      #    expressed with string names as identifiers (in that case, length 
      #    values are unquestioned) or specified as positive integers.
      #    Use '..' or '...' for repetition of the last shape.
      #    See EXAMPLES below.
      # 
      #  * args (multi-dim arrays such as NArray): arrays to be checked
      # 
      #  RETURN VALUE
      #  * nil
      # 
      #  POSSIBLE EXCEPTIONS
      #  * exception is raised if cshapes and args are inconsistent:
      #    
      #    * RuntimeError, if the arrays do not have shapes specified by cshapes.
      # 
      #    * ArgeumentError, if the number of args are inconsistent with cshapes. 
      #      This is likely a coding error of the user.
      # 
      #  EXAMPLES
      # 
      #  *  to check whether three arrays u, v, and w are shaped as
      #     u[nx], v[ny], and w[nx,ny], where nx and ny are any integer:
      # 
      #        NumRu::Misc.check_shape_consistency('nx ny nx,ny',u,v,w)
      # 
      #     Or equivalently,
      # 
      #        NumRu::Misc.check_shape_consistency('m  n  m,n',u,v,w)
      # 
      #     because actual strings does not matter.
      # 
      #  * To specify fixed lengths, use integers instead of names:
      # 
      #        NumRu::Misc.check_shape_consistency('4  n  4,n',u,v,w)
      # 
      #    In this case, u,v,w must have shapes [4], [ny], and [4,ny],
      #    where ny is any length.
      # 
      #  * Use '..' or '...' to repeat the same shape:
      # 
      #        NumRu::Misc.check_shape_consistency('nx,ny ...',u,v,w)
      # 
      #    This ensures that u, v, and w are 2D arrays with the same shape.
      #    Note: '..' and '...' are the same, so you can use whichever you like.
      def check_shape_consistency(cshapes, *args)
	 ranks = Array.new
	 elm2idx = Hash.new
	 spl = cshapes.split(' ')
	 if spl.length >= 2 && /^\.\.\.?$/ =~ spl[-1]  # '..' or '...'
            ((spl.length-1)...args.length).each{|i|
	       spl[i]=spl[i-1]
	    }	       
	 end
	 if spl.length != args.length
	    raise ArgumentError,"# of the argument (#{args.length}) is inconsistent with the 1st arg '#{cshapes}'"
	 end
         spl.each_with_index{|csh,i| 
	    sh = csh.split(',')
	    ranks.push( sh.length )
	    sh.each_with_index{|tag,j|
	       elm2idx[tag] = Array.new if !elm2idx[tag]
	       elm2idx[tag].push([i,j])
	    }
	 }
	 ranks.each_with_index{|len,i|
	    if args[i].rank != len
	       raise "(#{i+1}th arg) unexepected rank #{args[i].rank} for #{len}"
	    end
	 }
	 elm2idx.each{|tag,ary|
	    if tag.to_i > 0   # numeric (positive integer)
	       size = tag.to_i
	       start = 0
	    else              # alphabet
	       size = args[ary[0][0]].shape[ary[0][1]]
	       start = 1
	    end
	    (start...ary.length).each{|i|
	       if args[ary[i][0]].shape[ary[i][1]] != size
		  if start == 0
		     raise "length of dim #{ary[i][1]} of #{ary[i][0]+1}th "+
                           "arg is unexpected " +
                           "(#{args[ary[i][0]].shape[ary[i][1]]} for #{size})"
		  else
		     raise "Dimension lengths inconsistent between "+
			"dim #{ary[0][1]} of #{ary[0][0]+1}th arg and " +
		        "dim #{ary[i][1]} of #{ary[i][0]+1}th arg"
		  end
	       end
	    }
	 }
	 nil
      end

   end
end

if __FILE__ == $0
   include NumRu

   puts '* test A *'
   u = NArray.float(3)
   v = NArray.float(5)
   w = NArray.float(3,5)
   Misc.check_shape_consistency('nx ny nx,ny',u,v,w)
   puts ' OK'
   Misc.check_shape_consistency('3 ny 3,ny',u,v,w)
   puts ' OK'
   begin
      Misc.check_shape_consistency('6 ny 6,ny',u,v,w)
   rescue
      puts "  exception raised as expected\n#{$!}"
      puts ' OK'
   end

   puts '* test B *'
   Misc.check_shape_consistency('nx,ny ...',w,w,w,w)
   puts ' OK'

   puts '* test C *'
   begin
      u = NArray.float(4)
      v = NArray.float(5)
      w = NArray.float(3,5)
      Misc.check_shape_consistency('nx ny nx,ny',u,v,w)
   rescue
      puts "  exception raised as expected\n#{$!}"
      puts ' OK'
   end

   puts '* test D *'
   begin
      u = NArray.float(3,5)
      v = NArray.float(5)
      w = NArray.float(3,5)
      Misc.check_shape_consistency('nx ny nx,ny',u,v,w)
   rescue
      puts "  exception raised as expected\n#{$!}"
      puts ' OK'
   end

   puts '* test E *'
   u = NArray.float(3,5,7)
   v = NArray.float(5)
   w = NArray.float(3,5)
   1000.times{ Misc.check_shape_consistency('nx,ny,nz ny nx,ny',u,v,w) }
   puts ' OK'

end