require 'complex'

class Sequence

  @@types = []
  # @@conversions = Hash.new( proc { |x| x } )
  @@coercions = Hash.new

  Type = Struct.new( :name, :type, :size, :default, :pack, :unpack )
  class Type
    def inspect
      name
    end
  end

  def Sequence.register_type( symbol, type, size, default, pack, unpack )
    type = Type.new( symbol.to_s, type, size, default, pack, unpack )
    eval "#{symbol.to_s} = type"
    @@coercions[ [ type, type ] ] = type
  end

  register_type( :BYTE, Fixnum, 1, 0, proc { |o| [o].pack("c") },
                proc { |s| s.unpack("c")[0] } )
  register_type( :UBYTE, Fixnum, 1, 0, proc { |o| [o].pack("C") },
                proc { |s| s.unpack("C")[0] } )
  register_type( :USINT, Fixnum, 2, 0, proc { |o| [o].pack("S") },
                proc { |s| s.unpack("S")[0] } )
  register_type( :ULINT, Fixnum, 4, 0, proc { |o| [o].pack("L") },
                proc { |s| s.unpack("L")[0] } )
  register_type( :SINT, Fixnum, 2, 0, proc { |o| [o].pack("s") },
                proc { |s| s.unpack("s")[0] } )
  register_type( :LINT, Fixnum, 4, 0, proc { |o| [o].pack("l") },
                proc { |s| s.unpack("l")[0] } )
  register_type( :SFLOAT, Float, 4, 0.0, proc { |o| [o].pack("f") },
                proc { |s| s.unpack("f")[0] } )
  register_type( :DFLOAT, Float, 8, 0.0, proc { |o| [o].pack("d") },
                proc { |s| s.unpack("d")[0] } )
  register_type( :SCOMPLEX, Complex, 8, Complex( 0.0, 0.0 ),
                proc { |o| [ o.real, o.imag ].pack( "ff" ) },
                proc { |s| Complex( *s.unpack( "ff" ) ) } )
  register_type( :DCOMPLEX, Complex, 16, Complex( 0.0, 0.0 ),
                proc { |o| [ o.real, o.imag ].pack( "dd" ) },
                proc { |s| Complex( *s.unpack( "dd" ) ) } )
  register_type( :OBJECT, Object, 1, nil, proc { |o| [o] },
                proc { |s| s[0] } )

  # def Sequence.register_conversion( from, to, operation )
  #   @@conversions[ [ from, to ] ] = operation
  # end

  def Sequence.register_coercion( result, type1, *type2 )
    type2.each { |t2|
      @@coercions[ [ type1, t2 ] ] = result
      @@coercions[ [ t2, type1 ] ] = result
    }
  end

  register_coercion( USINT, USINT, BYTE, UBYTE )
  register_coercion( ULINT, ULINT, USINT, BYTE, UBYTE )
  register_coercion( SINT, SINT, USINT, BYTE, UBYTE )
  register_coercion( LINT, LINT, ULINT, SINT, USINT, BYTE, UBYTE )
  register_coercion( SFLOAT, SFLOAT, LINT, ULINT, SINT, USINT, BYTE, UBYTE )
  register_coercion( DFLOAT, DFLOAT, SFLOAT, LINT, ULINT, SINT, USINT, BYTE,
                    UBYTE )
  register_coercion( SCOMPLEX, SCOMPLEX, SFLOAT, LINT, ULINT, SINT, USINT,
                    BYTE, UBYTE )
  register_coercion( DCOMPLEX, SCOMPLEX, DFLOAT )
  register_coercion( DCOMPLEX, DCOMPLEX, SCOMPLEX, DFLOAT, SFLOAT, LINT,
                    ULINT, SINT, USINT, BYTE, UBYTE )
  register_coercion( OBJECT, OBJECT, DCOMPLEX, SCOMPLEX, DFLOAT, SFLOAT, LINT,
                    ULINT, SINT, USINT, BYTE, UBYTE )

  def Sequence.import( type, data )
    retval = Sequence.new( type )
    retval.import( data )
  end

  def initialize( type = OBJECT, n = 0, value = nil )
    @type = type
    @data = @type.pack.call( value == nil ? @type.default : value ) * n
  end

  def import( data )
    @data = data
    self
  end

  def inspect
    retval = "["
    i = 0
    while i < size
      v = self[i].inspect
      i += 1
      v += ", " if i < size
      if retval.size + v.size >= 74 - "...]".size
        retval += "..."
        break
      else
        retval += v
      end
    end
    retval = "Sequence." + @type.inspect.downcase + "(" + size.inspect +
      "):\n" + retval + "]"
  end

  def to_s
    @data.to_s
  end

  def dup
    retval = Sequence.new( @type )
    retval.data = @data.dup
    retval
  end

  def coerce( other )
    if other.kind_of?( Sequence )
      target = @@coercions[ [ other.typecode, typecode ] ]
      if target == nil
        raise "Don't now how to coerce Sequence of type " +
          typecode.inspect + " and type " + other.typecode.inspect
      else
        return other.to_type( target ), to_type( target )
      end
    elsif other.kind_of?( Array )
      return other, to_a
    else
      raise "Don't now how to coerce Sequence and " + other.class.inspect
    end
  end

  def each
    for i in 0...size
      yield self[i]
    end
    self
  end

  def collect( type = OBJECT )
    retval = Sequence.new( type, size )
    for i in 0...size
      retval[i] = yield self[i]
    end
    retval
  end

  def to_a
    retval = [ nil ] * size
    for i in 0...size
      retval[i] = self[i]
    end
    retval
  end

  def zip( *args )
    arrs = [ self, *args ]
    n = ( arrs.collect { |arr| arr.size } ).max
    retval = Sequence.new( OBJECT, n )
    ( 0...n ).each { |i| retval[i] = arrs.collect { |arr| arr[i] } }
    retval
  end

  def +( other )
    if other.kind_of?( Array ) or other.typecode != typecode
      y, x = coerce( other )
      retval = x + y
    else
      retval = Sequence.new( @type )
      retval.data = @data + other.data
    end
    retval
  end

  def *( n )
    if n.kind_of?( String )
      join( n )
    else
      retval = Sequence.new( @type )
      retval.data = @data * n
      retval
    end
  end

  def join( sep )
    to_a.join( sep )
  end

  def ==( other )
    if other.kind_of?( Sequence )
      if other.typecode == typecode
        other.data == data
      else
        false
      end
    else
      false
    end
  end

  def ===( other )
    self == other
  end

  def eql?( other )
    if other.kind_of?( Sequence )
      other.data.eql?( @data )
    else
      false
    end
  end

  def hash
    @data.hash
  end

  def clear
    @data = @type.default * 0
  end

  def to_type( type )
    if type == @type
      self
    else
      retval = Sequence.new( type, size )
      # conversion = @@conversions[ [ @type.type, type.type ] ]
      for i in 0...size
        # retval[i] = conversion.call( self[i] )
        retval[i] =self[i]
      end
      retval
    end
  end

  def element_size
    @type.size
  end

  def size
    @data.size / element_size
  end

  def typecode
    @type
  end

  def fill( value )
    @data = @type.pack.call( value ) * size
    self
  end

  def <<( *arr )
    push( *arr )
  end

  def push( *arr )
    arr.each { |o| @data += @type.pack.call( o ) }
    self
  end

  def at( i )
    pos = i * @type.size
    @type.unpack.call( @data[ pos...( pos + @type.size ) ] )
  end

  def slice( *i )
    if i.size > 1
      i.collect { |idx| at( idx ) }
    else
      if i[0].kind_of?( Range )
        i[0].collect { |idx| at( idx ) }
      else
        at( i[0] )
      end
    end
  end

  def []( *i )
    slice( *i )
  end

  def []=( i, o )
    pos = i * @type.size
    @data[ pos...( pos + @type.size ) ] = @type.pack.call( o )
    o
  end

  protected

  def data=( s )
    @data = s
  end

  def data
    @data
  end

end



class MultiArray

  BYTE      = Sequence::BYTE
  UBYTE     = Sequence::UBYTE
  USINT     = Sequence::USINT
  ULINT     = Sequence::ULINT
  SINT      = Sequence::SINT
  LINT      = Sequence::LINT
  SFLOAT    = Sequence::SFLOAT
  DFLOAT    = Sequence::DFLOAT
  SCOMPLEX  = Sequence::SCOMPLEX
  DCOMPLEX  = Sequence::DCOMPLEX
  OBJECT    = Sequence::OBJECT

  def MultiArray.import( type, data, *shape )
    retval = MultiArray.new( type )
    retval.import( data, *shape )
  end

  def initialize( type = OBJECT, *shape )
    @shape = shape
    @strides, size = compute_strides
    @data = Sequence.new( type, size )
  end

  def compute_strides
    strides = []
    size = @shape.inject( 1 ) { |stride,dim|
      if dim <= 0
        raise "All dimensions must be greater than zero (was #{dim})"
      end
      strides.push( stride ); stride * dim
    }
    return strides, size
  end

  def import( data, *shape )
    @shape = *shape
    @strides, size = compute_strides
    @data.import( data )
    self
  end

  def inspect
    "MultiArray." + @data.typecode.inspect.downcase + "(" +
      @shape.join(",") + "):\n" + print( @shape.size - 1, 0, 10 )[1]
  end

  def dup
    retval = MultiArray.new
    retval.shape = @shape.dup
    retval.strides = @strides.dup
    retval.data = @data.dup
    retval
  end

  def to_s
    @data.to_s
  end

  def +@
    self
  end

  def -@
    collect( typecode ) { |x| -x }
  end

  def +( other )
    if other.kind_of?( MultiArray )
      if other.typecode != typecode
        y, x = coerce( other )
      else
        y, x = other, self
      end
      retval = MultiArray.new( typecode, *shape )
      ( 0...size ).each { |i| retval.data[i] = data[i] + other.data[i] }
    else
      retval = collect( typecode ) { |x| x + other }
    end
    retval
  end

  def -( other )
    if other.kind_of?( MultiArray )
      if other.typecode != typecode
        y, x = coerce( other )
      else
        y, x = other, self
      end
      retval = MultiArray.new( typecode, *shape )
      ( 0...size ).each { |i| retval.data[i] = data[i] - other.data[i] }
    else
      retval = collect( typecode ) { |x| x - other }
    end
    retval
  end

  def /( other )
    if other.kind_of?( MultiArray )
      if other.typecode != typecode
        y, x = coerce( other )
      else
        y, x = other, self
      end
      retval = MultiArray.new( typecode, *shape )
      ( 0...size ).each { |i| retval.data[i] = data[i] / other.data[i] }
    else
      retval = collect( typecode ) { |x| x / other }
    end
    retval
  end

  def *( other )
    if other.kind_of?( MultiArray )
      if other.typecode != typecode
        y, x = coerce( other )
      else
        y, x = other, self
      end
      retval = MultiArray.new( typecode, *shape )
      ( 0...size ).each { |i| retval.data[i] = data[i] * other.data[i] }
    else
      retval = collect( typecode ) { |x| x * other }
    end
    retval
  end

  def []=( *parameters )
    @data[ pos( *(parameters[0...-1] ) ) ] = parameters.last
  end

  def []( *indices )
    @data[ pos( *indices ) ]
  end

  def collect( type = OBJECT, &action )
    retval = MultiArray.new
    retval.shape = @shape.dup
    retval.strides = @strides.dup
    retval.data = @data.collect( type, &action )
    retval
  end

  def zip( *args )
    arrs = [ self, *args ]
    # resizing !!!
    retval = MultiArray.new( OBJECT, *arrs[0].shape )
    ( 0...size ).each { |i|
      retval.data[i] = arrs.collect { |arr| arr.data[i] }
    }
    retval
  end

  def to_type( type )
    retval = MultiArray.new
    retval.shape = @shape.dup
    retval.strides = @strides.dup
    retval.data = @data.to_type( type )
    retval
  end

  def fill( value )
    @data.fill( value )
    self
  end

  def rank
    @shape.size
  end

  def size
    @data.size
  end

  def shape
    @shape
  end

  def typecode
    @data.typecode
  end

  protected

  def shape=( shape )
    @shape = shape
  end

  def strides
    @strides
  end

  def strides=( strides )
    @strides = strides
  end

  def data
    @data
  end

  def data=( data )
    @data = data
  end

  def pos( *indices )
    if indices.size != @shape.size
      raise "MultiArray has #{dimensions.size} indices (not #{indices.size})"
    end
    stride = @data.size
    retval = 0
    for i in 0...indices.size
      if not (0...@shape[i]).member?( indices[i] )
        raise "Index #{i} must be in #{0...@shape[i]} (was #{indices[i]})" 
      end
      retval += indices[i] * @strides[i]
    end
    retval
  end

  def print( level, skip, lines )
    if @shape.size == 0
      s = @data[0].inspect
    else
      stride = @strides[ level ]
      if level == 0
        s = "[ "
        for i in 0...@shape[ level ]
          s += ", " if i > 0
          v = @data[ skip + i ].inspect
          if s.size + v.size >= 74 - "...".size - ", ".size - "[  ]".size * ( @shape.size - level - 1 )
            s += "..."
            break
          else
            s += v
          end
        end
        lines -= 1
        s += " ]"
      else
        s = "[ "
        for i in 0...@shape[ level ]
          if lines > 0
            p = print( level - 1, skip + i * stride, lines )
            s += ",\n  " + "  " * ( @shape.size - level - 1 ) if i > 0
            s += p[1]
            lines = p[0]
          else
            if lines == 0
              if i > 0
                s += ",\n..."
              else
                s += "\n..."
              end
              lines = -1
            end
          end
        end
        s += " ]" if lines >= 0
      end
      [ lines, s ]
    end
  end

end