# -*- coding: euc-jp -*-
# 温度 (Temp) からジオポテンシャルを計算し, netCDF ファイルに出力する
# 

# 計算式を載せること

# 参考文献


# 更新履歴
# 
# 2012/06/26 納多 哲史   新規作成
# 

require "numru/ggraph"
include NumRu

prefix="."

# debug
#prefix="/home/noda/work/dcpam/dcpam5-20100224/practice/100225_nonsr_Omega1.0E_T21L16_S1380"
#prefix="/home/noda/data/101220_diff-DryAir-Vapor_T21L16_sr_Omega1.0_dt20m"

file_var = "Temp.nc"
name_var = "Temp"

TITLE_DEFAULT = "Geopotential"
OUTFN_PHI = "Phi.nc"

title = nil

# 物理定数
GasRUniv = 8.314 #  $ R^{*} $ [J K-1 mol-1]. 普遍気体定数.
MolWtDry = 28.964e-3  # $ M $ [kg mol-1]. 乾燥大気の平均分子量.
CpDry = 1004.6    # $ C_p $ [J kg-1 K-1]. 乾燥大気の定圧比熱.
GasRDry = GasRUniv / MolWtDry   # $ R $ [J kg-1 K-1].  乾燥大気の気体定数
Kappa = GasRDry / CpDry

starttime =    0
endtime   = 1095

startlon =   0.0   # degree
endlon =   359.0   # degree




def help()
  print <<EOM

  Usage:
 
    ruby #{$0} --var file@var [options]

  Options:

    --title str            set title of a picture
    --range val:val        range of value to draw [default -5e+15:5e+15]
    --range_val val:val    same as "--range"
    --range_lon  val:val   range of longitude [default 0:360]
    --range_time val:val   range of time [default 1000:2000]
    --prefix path          set prefix of files [default "."]
    -h, --help             show this help

EOM
end

def get_file_var(str)
  p str_a = str.split(/\@/)
  if str_a.size == 2
    filename = str_a[0]
    varname  = str_a[1]
    return filename, varname
  else
    puts "ERROR: illegal format!\n"
    help
    exit 10
  end
end

def get_range(str)
  p str_a = str.split(/\:/)
  if str_a.size == 2
    range_start = eval(str_a[0])
    range_end   = eval(str_a[1])
    return range_start, range_end
  else
    puts "ERROR: illegal format!\n"
    help
    exit 10
  end
end



# コマンドライン引数処理

while ARGV.size > 0
  case ARGV[0]
  when "--var"
    ARGV.shift
    file_var, name_var = get_file_var(ARGV[0])
  when "--title"
    ARGV.shift
    title = ARGV[0]
  when "--range_time"
    ARGV.shift
    starttime, endtime = get_range(ARGV[0])
  when "--range_lon"
    ARGV.shift
    startlon, endlon = get_range(ARGV[0])
  when "--range", "--range_val"
    ARGV.shift
    val_min, val_max = get_range(ARGV[0])
  when "--prefix"
    ARGV.shift
    prefix = ARGV[0]
  when "--wsn"
    ARGV.shift
    n = ARGV[0].to_i
    if n >=1 and n <= 4
      wsn = n
    else
      puts "ERROR: Invalid number of draw number."
      help
      exit 70
    end
  when "-v", "--verbose"
    verbose = true
  when "-h", "--help"
    help
    exit 0
  else
    puts "ERROR: unknown argument!\n"
    help
    exit 20
  end
  ARGV.shift
end

# タイトル未定時の自動生成.
title = TITLE_DEFAULT

# debug info
if verbose
  print <<EOM
*** info ***

prefix    = #{prefix}
file_var  = #{file_var}
name_var  = #{name_var}
title     = "#{title}"

time : #{starttime}:#{endtime}
lon  : #{startlon}:#{endlon}

CpDry : #{CpDry}

EOM
end

# ファイルの存在確認


file = File.join(prefix, file_var)
temp = GPhys::IO.open(file, name_var)
temp = temp.cut('lon'=>startlon..endlon).cut('time'=>starttime..endtime)
temp = temp.data.val


# axis etc.
lon        = GPhys::IO.open(file, "lon").cut('lon'=>startlon..endlon).data.val
lat        = GPhys::IO.open(file, "lat").data.val
gp_lat_weight = GPhys::IO.open(file, "lat_weight")
lat_weight = gp_lat_weight.data.val * 0.5   # 和を 1 に規格化
sig        = GPhys::IO.open(file, "sig").data.val
gp_sig_weight = GPhys::IO.open(file, "sig_weight")
sig_weight = gp_sig_weight.data.val
sigm       = GPhys::IO.open(file, "sigm").data.val
time       = GPhys::IO.open(file, "time").cut('time'=>starttime..endtime).data.val

nlon = lon.size
nlat = lat.size
nsig = sig.size
nsigm = sigm.size
ntime = time.size

p nlon

# 重み (HydroAlpha, HydroBeta) の計算
# z_HydroAlpha(k) = ( r_Sigma(k-1) / z_Sigma(k) ) ** Kappa - 1.
# z_HydroBeta(k) = 1. - ( r_Sigma(k) / z_Sigma(k) ) ** Kappa

alpha = NArray.sfloat(nsig).fill(0.0)
beta  = NArray.sfloat(nsig).fill(0.0)

alpha[0..(nsigm-2)] = (sigm[0..(nsigm-2)] / sig) ** Kappa - 1.0
beta[0..(nsigm-2)]  = 1.0 - (sigm[1..(nsigm-1)] / sig) ** Kappa

# ジオポテンシャルの計算
# xyz_Phi(:,:,1) = CpDry * z_HydroAlpha(1) * xyz_Temp(:,:,1)
# do k = 2, kmax
#  xyz_Phi(:,:,k) =   xyz_Phi(:,:,k-1) &
#   &              + CpDry * z_HydroAlpha(k  ) * xyz_Temp(:,:,k  )   &
#   &              + CpDry * z_HydroBeta (k-1) * xyz_Temp(:,:,k-1)
# enddo

phi = NArray.sfloat(nlon,nlat,nsig,ntime).fill(0.0)

phi[true, true, 0, true] = CpDry * alpha[0] * temp[true, true, 0, true]

for k in 1..(nsig-1)
  phi[true, true, k, true] = \
    phi[true, true, k-1, true] + \
#    CpDry * alpha[k].reshape(1,1,nsig-1)  * temp[true, true, k, true] + \
#    CpDry * beta[k-1].reshape(1,1,nsig-1) * temp[true, true, k-1, true]
    CpDry * alpha[k]  * temp[true, true, k, true] + \
    CpDry * beta[k-1] * temp[true, true, k-1, true]
end





# 次から時間平均をとる




# time mean
# phi = phi.mean(3)


# GPhys オブジェクトの作成
va_lon = VArray.new( lon,
                    {"long_name"=>"longitude", "units"=>"degrees_east"},
                    "lon" )
ax_lon = Axis.new.set_pos(va_lon)
 
va_lat = VArray.new( lat,
                    {"long_name"=>"latitude","units"=>"degrees_north"},
                    "lat" )
ax_lat = Axis.new.set_pos(va_lat)
 
va_sig = VArray.new( sig,
                    {"long_name"=>"sigma at layer midpoints","units"=>"1"},
                    "sig" )
ax_sig = Axis.new.set_pos(va_sig)

# (2014-5-8) 追加
va_time = VArray.new( time,
                    {"long_name"=>"time","units"=>"day"},
                    "time" )
ax_time = Axis.new.set_pos(va_time)


 
data = VArray.new( phi,
                   {"long_name"=>"Geopotential", "units"=>"J kg-1"},
                   "Phi" )
gphys = GPhys.new( Grid.new(ax_lon,ax_lat,ax_sig,ax_time), data )


# netCDF ファイルを出力
outfile = NetCDF.create(OUTFN_PHI)
GPhys::NetCDF_IO.write( outfile, gp_lat_weight )
GPhys::NetCDF_IO.write( outfile, gp_sig_weight )
GPhys::NetCDF_IO.write( outfile, gphys )
outfile.close