require 'getopts'
require 'numru/gphys'

def cal_global_mean(gp)

  lat     = gp.coord('lat')
  phi     = lat*2*PI/360.0
  grid    = Grid.new(gp.axis('lat'))
  cos_lat = GPhys.new(grid, cos(phi))

  gp_global_mean = ( gp.average('lon').average('level') * cos_lat ).average('lat')

  return gp_global_mean

end

################################################################################
include NumRu
include Misc::EMath

unless getopts("v:", "hgt:", "temp:", "shum:", "mask:", "output:epflx.nc")
  print "#{$0}:illegal options.\n"
  exit 1
end

Cp  = UNumeric.new(1004.0,   'J.K-1.kg-1')
G   = UNumeric.new(9.81,     "m.s-2") 

nc_vwnd, var_vwnd =   ($OPT_v).split(/\s*@\s*/)
nc_hgt,  var_hgt  =   ($OPT_hgt).split(/\s*@\s*/)
nc_temp, var_temp =   ($OPT_temp).split(/\s*@\s*/)
nc_mask, var_mask =   ($OPT_mask).split(/\s*@\s*/)

gp_v =     GPhys::IO.open( nc_vwnd,  var_vwnd )
gp_hgt =   GPhys::IO.open( nc_hgt,   var_hgt  )
gp_t =     GPhys::IO.open( nc_temp,  var_temp )
gp_mask =  GPhys::IO.open( nc_mask,  var_mask )

ofile = NetCDF.create($OPT_output)

  nt = gp_v.shape[-1]
  i = 0
  GPhys::IO.each_along_dims_write([gp_v, gp_hgt, gp_t, gp_mask], ofile, -1){
    |v, hgt, temp, mask|
    i += 1
    print "processing #{i} / #{nt} ..\n" if (i % (nt/20+1))==1

    v         = v    + mask
    hgt       = hgt  + mask
    temp      = temp + mask

    temp_bar  = temp.mean('lon')
    temp_dash = temp - temp_bar
    hgt_bar   = hgt.mean('lon')
    hgt_dash  = hgt  - hgt_bar
    v_bar     = v.mean('lon')
    v_dash    = v - v_bar

    temp_gl = cal_global_mean(temp)
    hgt_gl  = cal_global_mean(hgt)

    vt    = v * ( temp - temp_gl ) * Cp
    vh    = v * ( hgt  - hgt_gl  ) * G
    dsefl = ( vt + vh ).mean('lon')
    dsefl_bar  = ( ( temp_bar - temp_gl ) * Cp    + ( hgt_bar - hgt_gl )  * G ) * v_bar
    dsefl_dash = ( ( ( temp_dash - temp_gl ) * Cp + ( hgt_dash - hgt_gl ) * G ) * v_dash ).mean('lon')

    dsefl.data.rename!("dsefl")
    dsefl.data.set_att('long_name', 'dry staitic energy transport (total)')
    dsefl.data.set_att('units', 'W')
    dsefl_bar.data.rename!("dsefl_bar")
    dsefl_bar.data.set_att('long_name', 'dry staitic energy transport (bar)')
    dsefl_bar.data.set_att('units', 'W')
    dsefl_dash.data.rename!("dsefl_dash")
    dsefl_dash.data.set_att('long_name', 'dry staitic energy transport (dash)')
    dsefl_dash.data.set_att('units', 'W')
    
  [ dsefl, dsefl_bar, dsefl_dash ].each{|gp|       #  This part will not
      gp.data.att_names.each{|nm|                  #  be needed in future.
	gp.data.del_att(nm) if /^valid_/ =~ nm     #  (Even now, it is not
	gp.data.del_att(nm) if /^title/ =~ nm      #  needed if the valid  
      }                                            #  range is wide enough) 
    }                                                
    

    [ dsefl, dsefl_bar, dsefl_dash ]
  }
  
  ofile.close