$LOAD_PATH.unshift( File.expand_path("../lib") )
require "narray"
require "numru/ganalysis/fitting"
require "test/unit"
include NumRu

class FittingTest < Test::Unit::TestCase
  def setup
    nx = 7
    ny = 5
    @x = NArray.float(nx).indgen! - nx/2
    @y = NArray.float(ny).indgen! - ny/2

    x = @x.newdim(-1)
    y = @y.newdim(0)
    @data = x + 2*y + 10     # then the coefficients will be [1, 2, 10]
    dist = (x*x)*(y*y)*1e-1
    dist -= dist.mean
    @data += dist  # disturbance that does not change the fitting result

    @data2 = NArray.float(nx,2,ny)
    @data2[true,0,true] = @data - 1
    @data2[true,1,true] = @data + 1

    @data3 = 2*NMath.cos(x) + 3*NMath.sin(y)

    @f_cosx = proc {|x,y| NMath.cos(x).newdim!(-1)}
    @f_sinx = proc {|x,y| NMath.sin(x).newdim!(-1)}
    @f_cosy = proc {|x,y| NMath.cos(y).newdim!(0)}
    @f_siny = proc {|x,y| NMath.sin(y).newdim!(0)}

    d = NArray.float(nx,2,ny)
    d[true,0,true] = x + 2*y + 10 - 1
    d[true,1,true] = x + 2*y + 10 + 1
    mask = NArray.byte( *d.shape )
    mask[0..3,true,0..3] = mask[-3..-1,true,true] = 1
    mask[0,0,0] = 0           # this is the source of a small descrepancy
    @data2m = NArrayMiss.to_nam( d, mask )

    @data4 = NArray.float(nx,2,ny,2)
    @data4[false,0] = d - 0.1
    @data4[false,1] = d + 0.1
  end

  def test_least_square_fit_1
    f_x = GAnalysis::Fitting::X
    f_y = GAnalysis::Fitting::Y

    c, bf, dif = GAnalysis::Fitting.least_square_fit(@data, [@x,@y], [f_x, f_y])
    err = (c - NArray[1.0, 2.0, 10.0]).abs.max
    assert( err < 1e-5 )

    c, bf, dif = GAnalysis::Fitting.least_square_fit(@data2, [@x,@y], 
                                                     [f_x,f_y], [1])
    err = (c - NArray[1.0, 2.0, 10.0]).abs.max
    assert( err < 1e-5 )
  end

  def test_least_square_fit_2
    c, bf, dif = GAnalysis::Fitting.least_square_fit(@data3, [@x,@y], 
                                      [@f_cosx, @f_sinx, @f_cosy, @f_siny])
    err = (c - NArray[2.0, 0, 0, 3, 0]).abs.max
    assert( err < 1e-10 )
    assert( dif < 1e-10 )
  end

  def test_least_square_fit_3
    f_x = GAnalysis::Fitting::X
    f_y = GAnalysis::Fitting::Y

    c, bf, dif = GAnalysis::Fitting.least_square_fit(@data2m, [@x,@y], 
                                                     [f_x,f_y],[1])
    err = (c - NArray[1.0, 2.0, 10.0]).abs.max
    assert( err < 3e-2 )
  end

  def test_least_square_fit_4
    f_x = GAnalysis::Fitting::X
    f_y = GAnalysis::Fitting::Y

    c, bf, dif = GAnalysis::Fitting.least_square_fit(@data4, [@x,@y], [f_x,f_y],
                                                nil, [1,3])
    assert_equal( c.rank, 3 )  
    assert( dif < 1e-10 )   # should be a perfect fit (except round-off)


    c, bf, dif = GAnalysis::Fitting.least_square_fit(@data4, [@x,@y], [f_x,f_y],
                                                [3], [1])
    assert_equal( c.rank, 2 )
    assert_equal( bf.rank, 4 )
    assert_equal( bf.shape[3], 1 )
  end

  def test_least_square_fit_5
    va = VArray.new(@data4,{"long_name"=>"data4","units"=>""},"data4")
    xax = Axis.new().set_pos( VArray.new(@x,{"units"=>""},"x") )
    yax = Axis.new().set_pos( VArray.new(@y,{"units"=>""},"y") )
    v1 = NArray.float(@data4.shape[1]).indgen!
    ax1 = Axis.new().set_pos( VArray.new(v1,nil,"d1") )
    v3 = NArray.float(@data4.shape[3]).indgen!
    ax3 = Axis.new().set_pos( VArray.new(v3,nil,"d3") )
    grid = Grid.new(xax, ax1, yax, ax3)
    gphys = GPhys.new(grid, va)

    f_x = GAnalysis::Fitting::X
    f_y = GAnalysis::Fitting::Y

    c, bf, dif = gphys.least_square_fit([f_x,f_y], [3], [1])
    assert_equal( c.rank, 2 )
    assert_equal( bf.class, NumRu::GPhys )
    assert_equal( bf.rank, 3 )   # rank reduced <-- ensemble_dims are deleted
    assert_equal( bf.shape, [7,2,5] )   # [nx, len of indep_dim, ny ]
  end

end