$LOAD_PATH.unshift( File.expand_path("../lib") )
require "narray"
require "../gphys_ext"   # extention library is currently at the top directory
require "numru/gphys"
require "numru/ganalysis/met_z"
require "numru/ganalysis/sigma_coord"
require "test/unit"
include NumRu


class MetZTest < Test::Unit::TestCase
  def setup
    nx = 6
    x = NArray.float(nx).indgen!*60
    lam = ( x * (Math::PI/180) ).newdim(-1)
    ny = 5
    y = NArray.float(ny).indgen!*30 - 60.0
    phi = ( y * (Math::PI/180) ).newdim(0)
    ns = 5
    s = 1.0 - NArray.float(ns).indgen / ns
    vx = VArray.new(x,{"units"=>"degrees_east", "long_name"=>"longitude"},"lon")
    vy = VArray.new(y,{"units"=>"degrees_north","long_name"=>"latitude"},"lat")
    vs = VArray.new(s,{"units"=>"1","long_name"=>"sigma"},"sig")
    ax = Axis.new().set_pos(vx)
    ay = Axis.new().set_pos(vy)
    as = Axis.new().set_pos(vs)
    grid_xy = Grid.new(ax, ay)
    grid_xys = Grid.new(ax, ay, as)

    p = s * 1000.0
    vp = VArray.new(p,{"units"=>"hPa","long_name"=>"pressure"},"p")
    ap = Axis.new().set_pos(vp)
    grid_xyp = Grid.new(ax, ay, ap)

    ps = NArray.float(nx,ny).fill!(1000) + NMath.cos(lam)*NMath.cos(phi)*50
    ps[1,1] = 750 
    ps[2,1] = 500 
    ps[3,1] = 700 

    #< for test_mass_strm_p_* and test_mass_strm_* >

    @ps = GPhys.new( grid_xy, VArray.new(ps, {"units"=>"hPa"}, "ps") )

    #v = NArray.float(nx,ny,ns) + (NMath.cos(lam)*NMath.sin(2*phi)*10).newdim(-1)
    v = NArray.float(nx,ny,ns).fill!(1.0)

    @v_on_sig = GPhys.new( grid_xys, VArray.new(v, {"units"=>"m.s-1"}, "v") )
    @v_on_p   = GPhys.new( grid_xyp, VArray.new(v, {"units"=>"m.s-1"}, "v") )

    #v2 = v.clone
    #(0...ns).each{|k| v2[false,k]=ns-k}
    #v2[2,0,false]=300
    #@v_on_p2 = GPhys.new( grid_xyp, VArray.new(v2, {"units"=>"m.s-1"}, "v") )
    
    @ans1 = NArray[ 5.41667e+07, 6.5e+09, 1.3e+10, 1.95e+10, 2.57292e+10 ]

    #< additionally for test_mass_strm_* >

    @p_on_sig = GAnalysis::SigmaCoord.sig_ps2p(@ps, s, 2)

    @pcoord = vs * 1000
    @pcoord.units = "hPa"
    @pcoord.name = "p"
    @pcoord.long_name = "pressure"
    @pcoord.put_att("standard_name","air_pressure")
    @pcoord.put_att("positive","down")

  end

  def test_mass_strm_p_1
    mstr = GAnalysis::MetZ.mass_strm_p(@v_on_sig, @ps)
    v = mstr[0,true].val / (2 * Math::PI)
    eps = (v-@ans1).abs.max / v.abs.max
    assert( eps < 1e-5 )
  end

  #def test_mass_strm_p_2
  #  mstr = GAnalysis::MetZ.mass_strm_p(@v_on_p2, @ps)
  #end

  def test_mass_strm_p_3
    mstr = GAnalysis::MetZ.mass_strm_p(@v_on_p, @ps)
    v = mstr[0,true].val / (2 * Math::PI)
    ans = NArray[ 0, 6.5e+09, 1.3e+10, 1.95e+10, 2.57292e+10 ]
    eps = (v-ans).abs.max / v.abs.max
    assert( eps < 1e-5 )
  end

  def test_mass_strm_any_1
    mstr = GAnalysis::MetZ.mass_strm_any(@v_on_sig, @ps, @p_on_sig/10, @pcoord/10)
    v = mstr[0,true].val / (2 * Math::PI)
    eps = (v - @ans1).abs.max / v.abs.max
    assert( eps < 1e-5 )
  end

end