#!/usr/bin/ruby

$: << File::dirname(__FILE__)

require 'test/unit'
require 'build_SWIG/FenrirMath.so'

require 'complex'
require 'matrix'

if RUBY_VERSION >= '1.9.0'
class Complex
  def Complex.new(x, y)
    return Complex(x, y)
  end
end
end

unless Float.method_defined?(:image) then
class Float
  def image; return self.imaginary; end
end
end
unless Complex.method_defined?(:image) then
class Complex
  alias :image :imag
end
end

$: << "../../ruby"
require 'FFT'

begin
  require 'gsl'
  class GSL::Matrix
    def row_size; self.size1; end
    def column_size; self.size2; end
  end
  class GSL::Vector
    def row_size; self.size; end
    def column_size; 1; end
  end
  class GSL::Vector::Complex
    def row_size; self.size; end
    def column_size; 1; end
  end
  class GSL::Complex
    def image; self.imag; end
  end
rescue LoadError
end

def print_and_eval(exp)
  puts "#{exp}: #{eval(exp, TOPLEVEL_BINDING)}"
end

class TestComplex < Test::Unit::TestCase

  def setup
    @a = FenrirMath::DComplex::new(1,3)
    @b = FenrirMath::DComplex::new(2,3)
    
    @a_ruby = Complex::new(1,3)
    @b_ruby = Complex::new(2,3)
  end
  
  def teardown
  end
  
  def test_to_s
    assert_equal(@a.to_s, "1 + 3i")
  end
  
  def test_real
    assert_equal(@a.real, @a_ruby.real)
  end
  
  def test_real_eq
    assert_equal((@a.real=3), 3.0)
  end
  
  def test_imaginary
    assert_equal(@a.imaginary, @a_ruby.image)
    assert_equal(@a.image, @a_ruby.image)
  end
  
  def test_imaginary_eq
    assert_equal((@a.imaginary=2), 2.0)
    assert_equal((@a.image=2), 2.0)
  end
  
  [:abs2, :abs, :arg].each{|item|
    module_eval <<-__TEXT__
      def test_#{item}
        assert_equal(@a.#{item}, @a_ruby.#{item})
      end
    __TEXT__
  }
  
  def compare(v, v_ruby)
    assert_equal(v.real, v_ruby.real)
    assert_equal(v.imaginary, v_ruby.image)
  end
  
  def compare2(v, v_ruby, delta)
    assert_in_delta(v.real, v_ruby.real, delta)
    assert_in_delta(v.imaginary, v_ruby.image, delta)
  end
  
  def test_power
    c_ruby = Complex::new(2,0)
    compare2(@a.power(2), (@a_ruby ** c_ruby), 1E-8)
  end
  
  def test_conjugate
    compare(@a.conjugate, @a_ruby.conjugate)
  end
  
  def test_exp
    c_ruby = Complex::new(Math::cos(@a_ruby.image), Math::sin(@a_ruby.image))
    c_ruby *= Complex::new(Math::exp(@a_ruby.real), 0)
    compare(FenrirMath::DComplex::exp(@a), c_ruby)
  end
  
  def test_eq
    assert_equal((@a == @b), false)
    assert_equal((@a != @b), true)
  end
  
  def test_substitute
    @a = @b
    assert_equal((@a == @b), true)
    assert_equal((@a != @b), false)
  end
  
  def test_op_scalar
    c_ruby = @a_ruby
    c_ruby = Complex::new(c_ruby.real + 1, c_ruby.image)
    compare(@a + 1, c_ruby)
    compare((@a += 1), c_ruby)
    c_ruby = Complex::new(c_ruby.real - 1, c_ruby.image)
    compare(@a - 1, c_ruby)
    compare(@a -= 1, c_ruby)
    c_ruby = Complex::new(c_ruby.real * 2, c_ruby.image * 2)
    compare(@a * 2, c_ruby)
    compare(@a *= 2, c_ruby)
    c_ruby = Complex::new(c_ruby.real / 2, c_ruby.image / 2)
    compare(@a / 2, c_ruby)
    compare(@a /= 2, c_ruby)
  end
  
  def test_op_complex
    compare(@a + @b, @a_ruby + @b_ruby)
    compare((@a += @b), @a_ruby += @b_ruby)
    compare(@a - @b, @a_ruby - @b_ruby)
    compare(@a -= @b, @a_ruby -= @b_ruby)
    compare(@a * @b, @a_ruby * @b_ruby)
    compare(@a *= @b, @a_ruby *= @b_ruby)
    compare(@a / @b, @a_ruby / @b_ruby)
    compare(@a /= @b, @a_ruby /= @b_ruby)
  end
end

class Matrix
  def []=(i,j,x)
    @rows[i][j]=x
  end
end


class TestMatrix < Test::Unit::TestCase

  def setup
    @a = FenrirMath::DMatrix::zero(8)
    @b = FenrirMath::DMatrix::zero(8)
    
    @a_ruby = Matrix::zero(8)
    @b_ruby = Matrix::zero(8)
  end
  
  def teardown
  end
  
  def assign(v)
    v.row_size.times{|i|
      v.column_size.times{|j|
        v[i, j] = yield(i, j)
      }
    }
  end
  
  def elm_assign(v1, v2)
    v1.row_size.times{|i|
      v1.column_size.times{|j|
        v2[i, j] = v1[i, j]
      }
    }
  end
  
  def compare_drv(v1, v2)
    assert_equal(v1.row_size, v2.row_size)
    assert_equal(v1.column_size, v2.column_size)
    v1.row_size.times{|i|
      v1.column_size.times{|j|
        yield(v1[i, j], v2[i, j])
      }
    }
  end
  
  def compare(v1, v2)
    compare_drv(v1, v2){|a, b|
      a, b = yield(a, b) if block_given?
      assert_equal(a, b)
    }
  end
  
  def compare_complex(v1, v2)
    compare_drv(v1, v2){|a, b|
      a, b = yield(a, b) if block_given?
      assert_equal(a.real, b.real)
      assert_equal(a.image, b.image)
    }
  end
  
  def compare2(v1, v2, delta)
    compare_drv(v1, v2){|a, b|
      a, b = yield(a, b) if block_given?
      assert_in_delta(a, b, delta)
    }
  end
  
  def compare2_complex(v1, v2, delta)
    compare_drv(v1, v2){|a, b|
      a, b = yield(a, b) if block_given?
      assert_in_delta(a.real, b.real, delta)
      assert_in_delta(a.image, b.image, delta)
    }
  end
  
  def test_get
    assert_equal(@a[0, 0], 0)
  end
  
  def test_set
    assert_equal((@a[0, 0] = 1), 1)
  end
  
  def test_copy
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    a_backup = @a.copy
    elm_assign(a_backup, @a_ruby)
    assign(@a){|i, j| elm = rand}
    elm_assign(a_backup, @a_ruby)
  end
  
  def test_rows
    assert_equal(@a.rows, 8)
    assert_equal(@a.row_size, 8)
  end
  
  def test_columns
    assert_equal(@a.columns, 8)
    assert_equal(@a.column_size, 8)
  end
  
  def test_unit
    compare(FenrirMath::DMatrix::unit(8), Matrix::unit(8))
    compare(FenrirMath::DMatrix::identity(8), Matrix::identity(8))
    compare(FenrirMath::DMatrix::I(8), Matrix::I(8))
  end
  
  def test_scalar
    compare(FenrirMath::DMatrix::scalar(8, 2), Matrix::scalar(8, 2))
  end
  
  def test_zero
    compare(@a, @a_ruby)
  end
  
  def test_clear
    assign(@a){|i, j| elm = rand}
    @a.clear
    compare(@a, @b)
  end
  
  def test_transpose
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    compare(@a.transpose, @a_ruby.transpose)
    compare(@a.t, @a_ruby.t)
  end
  
  def test_partial
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    compare(@a.partial(4, 4, 2, 2), @a_ruby.minor(2, 4, 2, 4))
    compare(@a.minor(2, 4, 2, 4), @a_ruby.minor(2, 4, 2, 4))
  end
  
  def test_row_vector
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    @a.row_size.times{|i|
      compare(@a.row_vector(i), Matrix::rows([@a_ruby.row_vectors[i].to_a]))
      compare(@a.rowVector(i), Matrix::rows([@a_ruby.row_vectors[i].to_a]))
      compare(@a.row_vectors[i], Matrix::rows([@a_ruby.row_vectors[i].to_a]))
    }
  end
  
  def test_column_vector
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    @a.column_size.times{|j|
      compare(@a.column_vector(j), Matrix::columns([@a_ruby.column_vectors[j].to_a]))
      compare(@a.columnVector(j), Matrix::columns([@a_ruby.column_vectors[j].to_a]))
      compare(@a.column_vectors[j], Matrix::columns([@a_ruby.column_vectors[j].to_a]))
    }
  end
  
  def test_ex_rows
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    i_1 = 2
    i_2 = 3
    @a_ruby.column_size.times{|j|
      temp = @a_ruby[i_1, j]
      @a_ruby[i_1, j] = @a_ruby[i_2, j]
      @a_ruby[i_2, j] = temp
    }
    a_backup = @a.copy
    compare(@a.exchangeRows(i_1, i_2), @a_ruby)
    @a = a_backup
    compare(@a.exchange_rows!(i_1, i_2), @a_ruby)
  end
  
  def test_ex_columns
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    j_1 = 2
    j_2 = 3
    @a_ruby.row_size.times{|i|
      temp = @a_ruby[i, j_1]
      @a_ruby[i, j_1] = @a_ruby[i, j_2]
      @a_ruby[i, j_2] = temp
    }
    a_backup = @a.copy
    compare(@a.exchangeColumns(j_1, j_2), @a_ruby)
    @a = a_backup
    compare(@a.exchange_columns!(j_1, j_2), @a_ruby)
  end
  
  def test_check
    assign(@a){|i, j| elm = (i == j ? rand : 0)}
    elm_assign(@a, @a_ruby)
    assert_equal(@a.isSquare, @a_ruby.square?)
    assert_equal(@a.square?, @a_ruby.square?)
    assert_equal(@a.isDiagonal, true)
    assert_equal(@a.diagonal?, true)
    assert_equal(@a.isSymmetric, true)
    assert_equal(@a.symmetric?, true)
    assert_equal(@a.isDifferentSize(@b), false)
    assert_equal(@a.different_size?(@b), false)
  end
  
  def test_trace
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    assert_equal(@a.trace, @a_ruby.trace)
    assert_equal(@a.tr, @a_ruby.tr)
  end
  
  def test_comatrix
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    compare(@a.coMatrix(0, 0), 
        @a_ruby.minor(1, @a_ruby.row_size - 1, 1, @a_ruby.column_size - 1))
  end
  
  def test_LU
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    lu = @a.decomposeLU
    lu.L.row_size{|i|
      lu.L.column_size{|j|
        if i > j then
          assert_equal(lu.U[i, j], 0)
        elsif i < j then
          assert_equal(lu.L[i, j], 0)
        end
      }
    }
    compare2(@a, lu.L * lu.U, 1E-8)
    
    # GSLの定義とあわせる
    lu.L.row_size.times{|i|
      factor = lu.L[i, i]
      lu.L.column_size.times{|j|
        lu.L[j, i] /= factor
        lu.U[i, j] *= factor
      }
    }
    
    begin
      require 'gsl'
      @a_ruby = GSL::Matrix[*@a_ruby.to_a]
      lu_ruby, perm, sign = GSL::Linalg::LU.decomp(@a_ruby)
      l_ruby = lu_ruby.clone.lower
      l_ruby.size1.times{|i| l_ruby[i, i] = 1.0}
      u_ruby = lu_ruby.clone.upper
      #compare2(lu.L * lu.U, l_ruby * u_ruby, 1E-8)
    rescue LoadError
    end
  end
  
  def test_UD
    assign(@a){|i, j| elm = (i <= j ? rand : @a[j, i])}
    elm_assign(@a, @a_ruby)
    ud = @a.decomposeUD
    ud.U.row_size{|i|
      ud.D.column_size{|j|
        if i > j then
          assert_equal(ud.U[i, j], 0)
          assert_equal(ud.D[i, j], 0)
        elsif i < j then
          assert_equal(ud.D[i, j], 0)
        end
      }
    }
    compare2(@a, ud.U * ud.D * ud.U.t, 1E-8)
    
    begin
      require 'gsl'
      
    rescue LoadError
    end
  end

  def test_det
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    assert_in_delta(@a.det, @a_ruby.det, 1E-8)
    assert_in_delta(@a.determinant, @a_ruby.determinant, 1E-8)
  end
  
  def test_op
    assign(@a){|i, j| elm = rand}
    elm_assign(@a, @a_ruby)
    assign(@b){|i, j| elm = rand}
    elm_assign(@b, @b_ruby)
    
    compare(@a *= 2, @a_ruby *= 2)
    compare(@a * 2, @a_ruby * 2)
    compare(@a /= 2, @a_ruby /= 2)
    compare(@a / 2, @a_ruby / 2)
    compare(@a += @b, @a_ruby += @b_ruby)
    compare(@a + @b, @a_ruby + @b_ruby)
    compare(@a -= @b, @a_ruby -= @b_ruby)
    compare(@a - @b, @a_ruby - @b_ruby)
    compare2(@a *= @b, @a_ruby *= @b_ruby, 1E-8)
    compare2(@a * @b, @a_ruby * @b_ruby, 1E-8)
    compare2(-@a, @a_ruby * -1, 1E-8)
    compare2(@a.inverse, @a_ruby.inverse, 1E-8)
    compare2(@a.inv, @a_ruby.inv, 1E-8)
    compare2(@a / @b, @a_ruby / @b_ruby, 1E-8)
    compare2(@a /= @b, @a_ruby /= @b_ruby, 1E-8)
    compare2(@a.sqrt * @a.sqrt, @a_ruby, 1E-6){|a, b| [a.real, b]}
  end
  
  def test_eigen
    assign(@a){|i, j| elm = rand}
    a_complex = FenrirMath::CDMatrix::new(@a.rows, @a.columns)
    assign(a_complex){|i, j| FenrirMath::DComplex::new(@a[i, j], 0)}
    compare2_complex(a_complex, @a, 1E-8){|a, b| [a.real, b]}
    a_eigen = @a.eigen rescue nil
    a_eigen_vals = []
    a_eigen_vecs = []
    
    if a_eigen then
      a_eigen_val_vecs = []
      @a.rows.times{|j|
        a_eigen_val_vecs << \
            [a_eigen[j, @a.rows], a_eigen.partial(@a.rows, 1, 0, j)]
      }
      # 固有値が大きい順に並べる
      a_eigen_val_vecs.sort!{|a, b|
        (b[0] == a[0].conjugate) ? \
            (b[0].image <=> a[0].image) : (b[0].abs <=> a[0].abs) 
      }.each{|item|
        a_eigen_vals << item[0]
        a_eigen_vecs << item[1]
      }
      
      @a.rows.times{|j|
        lhs = a_complex * a_eigen_vecs[j]
        rhs = a_eigen_vecs[j] * a_eigen_vals[j]
        #compare2_complex(lhs, rhs, 1E-8)
      }
    end
    
    begin
      require 'gsl'
      elm_assign(@a, @a_ruby)
      @a_ruby = GSL::Matrix[*@a_ruby.to_a]
      a_ruby_eigen_vals, a_ruby_eigen_vecs = GSL::Eigen::nonsymmv(@a_ruby)
      a_ruby_eigen_val_vecs = []
      a_ruby_eigen_vals.size.times{|i|
        a_ruby_eigen_val_vecs << [a_ruby_eigen_vals[i], a_ruby_eigen_vecs[i]]
      }
      # 固有値が大きい順に並べかえて、比較
      a_ruby_eigen_val_vecs.sort!{|a, b|
        (b[0] == a[0].conjugate) ? \
            (b[0].image <=> a[0].image) : (b[0].abs <=> a[0].abs)
      }.each_with_index{|item, i|
        # 固有ベクトル
        def (item[1]).[](i, j); self.get(i); end
        
        if a_eigen then # 比較
          assert_in_delta(item[0].real, a_eigen_vals[i].real, 1E-8) # 固有値 実部
          assert_in_delta(item[0].image, a_eigen_vals[i].image, 1E-8) # 固有値 虚部
          #compare2_complex(item[1], a_eigen_vecs[i], 1E-8)
        end
      }
    rescue LoadError
    end
  end
end

class TestFFT < Test::Unit::TestCase
  def setup
    @a_ruby = []
    @b_ruby = []
    
    256.times{|i|
      @a_ruby << rand
      @b_ruby << Complex::new(rand,rand)
    }
  
    @a = FenrirMath::DVector::new
    @a_ruby.each{|item| @a.push item}
    
    @b = FenrirMath::DComplexVector::new
    @b_ruby.each{|item| @b.push FenrirMath::DComplex::new(item.real, item.image)}
  end
  
  def teardown
  end
  
  def compare(converted1, converted2)
    assert_equal(converted1.length, converted2.length)
    converted1.length.times{|i|
      #p converted1[i], converted2[i]
      assert_in_delta(converted1[i].real, converted2[i].real, 1E-4)
      assert_in_delta(converted1[i].image, converted2[i].image, 1E-4)
    }
  end

  def test_ft
    res_a = []
    @a.each_with_index{|v, i|
      res_a << FenrirMath::FFT::D_ft(@a, i)
    }
    res_a_ruby = []
    @a_ruby.each_with_index{|v, i|
      res_a_ruby << FFT::ft(@a_ruby, i)
    }
    compare(res_a, res_a_ruby)
  end
  
  def test_fft
    res_b = FenrirMath::FFT::D_fft(@b)
    def res_b.length; return self.size; end
    res_b_ruby = FFT::fft(@b_ruby)
    compare(res_b, res_b_ruby)
  end
end