#! /usr/bin/ruby
# -*- coding: utf-8 -*-

# 2012-09-05 石渡 覚え書き
# 時間範囲はハードコーディングしているので改良したい.
# 図を書くところまではやらないでデータを作るだけにしたい.



# =begin
# = 質量流線関数の計算と描画
# 
# == 概要
# 
# 東西風速と地表面気圧を読み込み, 質量流線関数の画像 dcl.ps を出力する.
# 
# == 計算方法
# 
# 子午面方向の質量流線関数 \bar{\psi} は以下のように定義される:
# %
# \begin{align}
#   \frac{\partial \bar{\psi}}{\partial p}
#   = - \frac{2 \pi a \cos \phi}{g} \bar{v}.
# \end{align}
# %
# ここで $p$ は気圧, $a$ は惑星半径, $\phi$ は緯度[rad], $g$ は重力加速度, $v$ は東西風速である.
# 両辺を $p$ で積分すると
# %
# %
# \begin{align}
#   \bar{\psi} = = - \frac{2 \pi a \cos \phi}{g} \int \bar{v} dp
# \end{align}
# %
# となる.
# 境界条件として $p=0$ で $\bar{\psi} = 0$ をとる.
# 
# 入力データの鉛直座標が $\sigma$ 座標である場合, 上記の式はそのまま使えず, 
# 計算の際には $p = p_s \sigma$ の関係を利用して置換する必要がある.
# 本来は $p$ 一定の場所で時間方向に積分する必要があるが, 
# 本スクリプトでは簡単のために $\sigma$ 一定の格子点での値を用いているため
# 地表面気圧の変動分程度の誤差が入ると考えられるので注意されたい.
# 
# 
# == 参考文献
# 
# * 小倉「気象力学通論」東京大学出版
# 
# 
# == 更新履歴
# 
# * 2013/06/26  石渡 正樹   オプションを入れられるように変更
# * 2009/11/24  納多 哲史   計算式の記入
# * 2009/11/24  納多 哲史   鉛直方向の軸の表記を修正
# * 2009/11/10  納多 哲史   新規作成
# 
# =end

require 'getoptlong'
require "numru/ggraph"
include NumRu

## オプション解析
parser = GetoptLong.new
parser.set_options(
        ###    global option   ###
        ['--lat',                      GetoptLong::REQUIRED_ARGUMENT],
        ['--time_start',                      GetoptLong::REQUIRED_ARGUMENT],
        ['--time_end',                      GetoptLong::REQUIRED_ARGUMENT],
        ['--lon_start',                      GetoptLong::REQUIRED_ARGUMENT],
        ['--lon_end',                      GetoptLong::REQUIRED_ARGUMENT],
        ['--range',                      GetoptLong::REQUIRED_ARGUMENT],
        ['--cint',                      GetoptLong::REQUIRED_ARGUMENT],
        ['--wsn',                      GetoptLong::REQUIRED_ARGUMENT]
                   )
parser.each_option do |name, arg|
  eval "$OPT_#{name.sub(/^--/, '').gsub(/-/, '_')} = '#{arg}'"  # strage option value to $OPT_val
end

wsn = ($OPT_wsn||4)

#time_start = 937
#time_end   = 1096

time_start = ($OPT_time_start||1).to_f
time_end = ($OPT_time_end||50).to_f

#lon_start = 0.0     # [degree]
#lon_end =   360.0   # [degree]

lon_start = ($OPT_lon_start||0).to_f # [degree]
lon_end = ($OPT_lon_end||360).to_f # [degree]

# range for drawing [10^10 kg s-1]
#maxval_draw = 100
#minval_draw = -100
#interval_draw = 5.0

range = ($OPT_range||'-100:100')
minval_draw = range.split(':')[0].to_f
maxval_draw = range.split(':')[1].to_f

interval_draw = ($OPT_cint||5).to_f

p interval_draw

# parameters
PI = 3.14159265358979323846
PlanetRadius = 6.4e+6       # [m]
Grav = 9.8                  # [m s-2]

# show help message
def help()
  puts <<EOM
Usage: $ ruby #{$0} file_V varname_V file_Ps varname_Ps

Example: $ ruby #{$0} V.nc V Ps.nc Ps

  *** command option ***
-h, --help: show this message.
--debug   : run debug mode.
EOM
end


#if ARGV[0] == "-h" or ARGV[0] == "--help" then
#  help()
#  exit 0
#elsif ARGV[0] == "--debug" then
#  datadir = "/home/noda/work/dcpam/ifort_opt_dcpam5-20090405_practice/090923_ape_Omega1.0E"
#  file_v     = File.join(datadir, 'V.nc')
#  varname_v  = 'V'
#  file_ps    = File.join(datadir, 'Ps.nc') 
#  varname_ps = 'Ps'
#elsif ARGV.length != 4 then
#  puts "ERROR: The number of argument is invalid."
#  help()
#  exit 1
#else
#  file_v     = ARGV[0]
#  varname_v  = ARGV[1]
#  file_ps    = ARGV[2]
#  varname_ps = ARGV[3]
#end

if ARGV.length < 4 then
  puts "ERROR: The number of argument is invalid."
  help()
  exit 1
else
  file_v     = ARGV[0]
  varname_v  = ARGV[1]
  file_ps    = ARGV[2]
  varname_ps = ARGV[3]
end


# ファイルが存在するか確認
if !FileTest.file?(file_v) then
  puts "ERROR: #{file_v} does not exist."
  exit 2
elsif !FileTest.file?(file_ps) then
  puts "ERROR: #{file_ps} does not exist."
  exit 3
end



# V (南北風速) [m s-1]
lon = GPhys::IO.open(file_v, "lon").cut('lon'=>lon_start..lon_end).data.val
lat = GPhys::IO.open(file_v, "lat").data.val
sig = GPhys::IO.open(file_v, "sig").data.val
sigm = GPhys::IO.open(file_v, "sigm").data.val
time = GPhys::IO.open(file_v, "time").cut('time'=>time_start..time_end).data.val

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

v = GPhys::IO.open(file_v, varname_v)
v = v.cut('lon'=>lon_start..lon_end).cut('time'=>time_start..time_end).data.val

# Ps (地表面気圧) [Pa]
ps = GPhys::IO.open(file_ps, varname_ps)
ps = ps.cut('lon'=>lon_start..lon_end).cut('time'=>time_start..time_end).data.val


# 質量流線関数の計算
#   台形で積分

msf = NArray.sfloat(nlat, nsigm).fill(0.0)

# 鉛直層数 nsigm が nsig より一つ大きいことに注意
msf_tmp = NArray.sfloat(nlon, nlat, nsigm).fill(0.0)

delta_sigma = NArray.sfloat(nsig).fill(0.0)

for k in 0..nsig-1
  delta_sigma[k] = sigm[k] - sigm[k+1]
end


# 時間方向に回す
for itime in 0..ntime-1


  # 初期化
  msf_tmp.fill!(0.0)

# 大気上端
# 境界条件として, $p=0$ で \psi = 0 とする
  k = nsig
  for i in 0..nlon-1
    for j in 0..nlat-1
      msf_tmp[i,j,k] = 0.0
    end
  end

  k = nsig-1
  for i in 0..nlon-1
    for j in 0..nlat-1
      delta_p = delta_sigma[k] * ps[i,j,itime]
      msf_tmp[i,j,k] = 0.5 * delta_p * v[i,j,k,itime]
    end
  end

# 大気上端以外
  for k in 2..nsig
    for j in 0..nlat-1
      for i in 0..nlon-1
        delta_p = delta_sigma[nsig-k] * ps[i,j,itime]
        msf_tmp[i,j,nsig-k] = msf_tmp[i,j,nsig-k+1]  \
                              + delta_p * v[i,j,nsig-k,itime]
      end
    end
  end

# 帯状平均を取って msf に追加 (時間平均操作込み)
  for k in 0..nsig-1
    for j in 0..nlat-1
      tmp = 0.0
      for i in 0..nlon-1
        tmp += (msf_tmp[i,j,k] - tmp) / (i+1.0)
      end
      msf[j,k] += (tmp - msf[j,k]) / (itime+1.0)
    end
  end

end


# 定数を掛ける

for k in 0..nsig-1
  for j in 0..nlat-1
    cos_phi = cos(2.0 * PI * lat[j]/360.0)
# 1e-10 は, tone の色が飛ばないようにするための工夫
    msf[j,k] *= 1e-10 * 2.0 * PI * PlanetRadius * cos_phi / Grav
#    msf[j,k] *= 2.0 * PI * PlanetRadius * cos_phi / Grav
  end
end


# GPhys オブジェクトの作成

lat_a = VArray.new( lat,
                    {"long_name"=>"latitude","units"=>"degrees_north"},
                    "lat" )
lat_axis = Axis.new.set_pos(lat_a)

sigm_a = GPhys::IO.open(file_v, "sigm").data
sigm_axis = Axis.new.set_pos(sigm_a)

data = VArray.new( msf,
                   {"long_name"=>"mass stream function [1e+10 kg s-1]", "units"=>"1e+10 kg s-1"},
                   "msf" )

msf_gphys = GPhys.new( Grid.new(lat_axis, sigm_axis), data )



# change colormap
#DCL.sgscmn(1)
DCL.sgscmn(14)
#DCL.sgscmn(66)


DCL.gropn(wsn)

GGraph.set_fig('viewport'=>[0.15,0.80,0.15,0.6])
GGraph.set_axes('ytitle'=>'sigma')

DCL.sgpset('lfull', true)     # 全画面表示
DCL.sgpset('lcntl', false)
DCL.uzfact(0.6)

GGraph.set_linear_contour_options('interval'=>interval_draw, 'max'=>maxval_draw, 'min'=>minval_draw)
GGraph.set_linear_tone_options('interval'=>interval_draw, 'max'=>maxval_draw, 'min'=>minval_draw)

GGraph.tone( msf_gphys )
GGraph.contour( msf_gphys, false )
GGraph.color_bar

DCL.grcls