# coding: cp932
require "numru/gphys"
require "numru/gphys/derivative"

module NumRu
  class GPhys
    # Wave activity flux by Takata&Nakamura(2001) (Cartesian; Stationary; Horizontal component only)
    # 
    # Simplified version neglecting the C_U*M term, which is zero for stationary
    # waves. (if the waves are not stationary, the result is the wave activity
    # flux relative to their phase propagation)
    # 
    # HISTORY
    # * T.Horinouchi 2012-12-07
    # 
    # CAUTION: The above neglect is equivalent to assume that the
    # ground-based phase speed is much smaller than the backgound wind
    # speed, WHICH MAY NOT BE VALID IF THE BACKGROUND WIND IS WEAK!!
    # 
    # self must be geopotential height anomaly (usu. deviation from 
    # climatology) on isobaric surface(s)
    # with longitude and latitude as the 1st & 2nd coordinates
    # 
    # ARGUMENTS:
    # * ubk (GPhys/nil) : background zonal wind (w/ same grid as self).
    #   If nil (special case), bg wind is not used (equiv to assume vbk==0)
    # * vbk (GPhys/nil) : background merid wind (w/ same grid as self).
    #   If nil (special case), bg wind is not used (equiv to assume vbk==0)
    # * prs (nil or UNumeric; default: nil) : pressure level (good for single 
    #   pressure-level data); if nil, the 2nd dim is assumed to be pressue
    # * reflat (Numeric; default: 45degN) : referential latitude (in degN)
    # 
    # RETURN VALUE:
    # * [wx ,wy] : where wx&wy are the x&y components of the 
    #   wave activity flux (the vertical component is not derived) 
    # 
    def tn2001fluxCSH(ubk, vbk, prs=nil, reflat=45.0)
      #< assumptions on dimensions : you can modify it to configure it! >
      xd = 0     # longitudinal dimension
      yd = 1     # latitudinal dimension
      pd = 2     # [used only if prs is nil] pressure dimension

      #< constants >
      g = UNumeric[9.8, "m.s-2"]  # gravity
      a = UNumeric[6.37e6, "m"]   # Earth's radius
      trot = 86400.0              # Earth's rotation in sec (to calc f0 below)
      y0 = Math::PI * reflat / 180.0   # reference latitude (radian)
      f0 = UNumeric[4*Math::PI / trot * Math::sin(y0), "s-1"]  # Coriolis para

      #< preparation >
      za = self      # geoptential height anomaly
      psi = za * (g/f0)            # quasi-geostrophic stream function
      x = psi.axis(xd).pos.convert_units("rad") * (a*Math::cos(y0))
      y = psi.axis(yd).pos.convert_units("rad") * a
      x.units = "m"   # rad.m -> m
      y.units = "m"   # rad.m -> m
      psi.axis(xd).set_pos(x)      # replace lon ax for differentiation
      psi.axis(yd).set_pos(y)      # replace lat ax for differentiation
      prs = psi.axis(pd).to_gphys if !prs

      #< calculation >
      psi_x = psi.threepoint_O2nd_deriv(xd)
      psi_xx = psi.deriv2nd(xd)
      psi_y = psi.threepoint_O2nd_deriv(yd)
      psi_yy = psi.deriv2nd(yd)
      psi_xy = psi_x.threepoint_O2nd_deriv(yd)
      if ubk && vbk
        #Case: TN2001 Eq.32 except the CuM term
        c = prs / (2*UNumeric[1000.0,"hPa"]) / (ubk**2 + vbk**2).sqrt #rho0/2U
        ub = ubk*c
        vb = vbk*c
        wx = ub*(psi_x**2 - psi*psi_xx) + vb*(psi_x*psi_y - psi*psi_xy) 
        wy = ub*(psi_x*psi_y - psi*psi_xy) + vb*(psi_y**2 - psi*psi_yy) 
        wx.long_name = "WAF-x (stationary)"
        wy.long_name = "WAF-y (stationary)"
      else
        #Case: Further simplification assuming V=0 (or the zonal momentum flux)
        c = prs / (2*UNumeric[1000.0,"hPa"])
        wx = (psi_x**2 - psi*psi_xx) * c
        wy = (psi_x*psi_y - psi*psi_xy) * c
        wx.long_name = "WAF-x (stationary,V=0)"
        wy.long_name = "WAF-y (stationary,V=0)"
      end

      #< post processing >
      grid = za.grid
      wx = GPhys.new(grid,wx.data)     # use the original grid
      wy = GPhys.new(grid,wy.data)     # use the original grid
      wx.name = "wx"
      wy.name = "wy"
      [wx,wy]
    end
  end
end

######################################################################
## Main program (test part)

if __FILE__ == $0
  require "numru/ggraph"
  include NumRu

  def prep_grph
    iws= 1
    DCL.sgscmn(4)                         # カラーマップ番号
    DCL.swpset("ldump",true) if iws==4    # 画像ファイルダンプ
    DCL.swpset('iwidth',1000)
    DCL.swpset('iheight',500)
    DCL.gropn(iws)                        # iws : DCL出力デバイス．1:画面, 2:file
    DCL.sldiv("y",2,1)
    #DCL.sgpset("lfull",true)
    DCL.sgpset('isub', 96)                # 下付き添字制御文字を '_' から '`' に
    DCL.glpset('lmiss',true)              # DCLの欠損値処理を on に
    #GGraph.set_fig "itr"=>10,"viewport"=>[0.05,0.85,0.2,0.5]  # 10:経度緯度座標
    GGraph.set_fig "itr"=>32,"viewport"=>[0.06,0.86,0.1,0.9] # 32:正距方位図法
    GGraph.set_map "coast_world"=>true
  end

  plev = 300
  latran = 20..75
  z = GPhys::IO.open('hgt.2012-01.nc', 'hgt').cut('level'=>plev,"lat"=>latran)
  clifl = 'ncep2.Jan.clim.1981-2010.nc'
  zcli = GPhys::IO.open(clifl,'hgt').cut('level'=>plev,"lat"=>latran)
  ubk = GPhys::IO.open(clifl,'uwnd').cut('level'=>plev,"lat"=>latran)
  vbk = GPhys::IO.open(clifl,'vwnd').cut('level'=>plev,"lat"=>latran)

  prep_grph
  (8..18).step(5) do |d0|                  # 2012年1月の日付でループ
    t0 = DateTime.new(2012,1,d0)          # 2012-01-d0 00UTC (時刻指定省略のため)
    t1 = t0 + 9.9                  # 5日間の終わり(データは毎正時なので +4 でも可)

    zp = z.cut("time"=>t0..t1).mean("time") # 半旬(5日)平均
    za = zp - zcli
    [ [ubk, vbk], [nil,nil ] ].each do |ub,vb|
      wx,wy = za.tn2001fluxCSH(ub, vb, UNumeric[plev,"hPa"], 45.0 )
      su = !ub.nil? ? "with UV" : "w/o UV" 
      GGraph.tone za,true,"max"=>500,"min"=>-500,
         "annot"=>false,  # 右脇のアノテーション非表示
         "title"=>"Z#{plev} & WAF #{su} #{t0.strftime('%Y/%m/%d')}-#{t1.strftime('%d')}"
      GGraph.contour zp,false,"int"=>100
      GGraph.vector wx,wy,false,"xint"=>3,"yint"=>2,"unit"=>true,"fact"=>1.1
      GGraph.color_bar "vcent"=>0.4,"vlength"=>0.3
    end

  end

  GGraph.tone ubk,true,"max"=>75,"min"=>-75,"int"=>15,"title"=>"clim U & (U,V)"
  GGraph.vector ubk,vbk,false,"xint"=>3,"yint"=>2,"unit"=>true,"fact"=>2
  GGraph.color_bar "vcent"=>0.4,"vlength"=>0.3

  uu = (ubk**2 + vbk**2).sqrt
  GGraph.tone uu,true,"max"=>70,"min"=>0,"int"=>5,"title"=>"clim WindSpeed & (U,V)"
  GGraph.vector ubk,vbk,false,"xint"=>3,"yint"=>2,"unit"=>true,"fact"=>2
  GGraph.color_bar "vcent"=>0.4,"vlength"=>0.3

  DCL.grcls
end