# coding: utf-8
#---------------------------------------------------------------
# 2 層モデルの計算結果から Barotropic/Baroclinic 渦度を取り出し,
# リスタートファイルを出力するスクリプト
# 主に, 次元を上げて再計算することを想定している
#---------------------------------------------------------------
#
# 履歴
#   2016/02/12     岡崎 正悟       作成

require "numru/gphys"
require "numru/netcdf"
require "narray"
include NumRu

# 入出力ファイル名
ifile = ARGV[0].to_s                  # データを取り出すファイル名
ofile = NetCDF.create("restart.nc")      # 出力ファイル名

# 変数の次元設定
wx, wy = 341, 341   # 渦度の波数(出力)
kx, ly = wx*2+1, wy*2+1

# 変数の定義
#----- 次元の変数 ----- 
ofile.def_dim("k",kx)
ofile.def_dim("l",ly)
k = ofile.def_var("k","sfloat",["k"])
l = ofile.def_var("l","sfloat",["l"])
k.put_att("longname","X-Wavenumber")     # 属性の設定
k.put_att("units","1")
l.put_att("longname","Y-Wavenumber")     # 属性の設定
l.put_att("units","1")
#----- 渦度の変数 -----
ee_tro_vor = ofile.def_var("ee_tro_vor","sfloat",["k","l"])
ee_tro_vor.put_att("longname","Barotropic Vorticity Spectra")     # 属性の設定
ee_tro_vor.put_att("units","1")

ee_cli_vor = ofile.def_var("ee_cli_vor","sfloat",["k","l"])
ee_cli_vor.put_att("longname","Baroclinic Vorticity Spectra")     # 属性の設定
ee_cli_vor.put_att("units","1")

# defineモード終わり
ofile.enddef

# 波数の情報を入れる
k.put( NArray.float(kx).indgen!(-wx) )
l.put( NArray.float(ly).indgen!(-wy) )

#----- データの取り出し ---------
file = NetCDF.open(ifile)

# 最終時刻の渦度の値のみを出力
lasttime   = (GPhys::IO.open(file, 't')).val.size
tro_vor_in = ((GPhys::IO.open(file, 'tro_vorsp')).cut({'t'=>lasttime})).val
cli_vor_in = ((GPhys::IO.open(file, 'cli_vorsp')).cut({'t'=>lasttime})).val

# 取り出す渦度の波数を取得
k_in = (((GPhys::IO.open(file, 'k')).val.size) - 1 ) / 2
l_in = (((GPhys::IO.open(file, 'l')).val.size) - 1 ) / 2

#----- 渦度の情報を代入 -----
# もとの結果よりも高端数の領域にはゼロを代入する

# 一時的に渦度を代入する変数を定義
ee_tro_vor_tmp = NArray.sfloat(ly,kx)
ee_cli_vor_tmp = NArray.sfloat(ly,kx)

for k in 0..kx-1 do
  for l in 0..ly-1 do
    if (171 <= k && k <= 511) && (171 <= l && l <= 511) then
    # if (k_in <= k && k <= (kx - k_in - 1 )) && (l_in <= l && l <= (ly - l_in -1)) then
      ee_tro_vor_tmp[l,k] = tro_vor_in[l-171,k-171]
      # ee_tro_vor_tmp[l,k] = tro_vor_in[l-l_in,k-k_in]
    else
      ee_tro_vor_tmp[l,k] = 0.0
    end
  end
end

for k in 0..kx-1 do
  for l in 0..ly-1 do
  if (171 <= k && k <= 511) && (171 <= l && l <= 511) then
  # if (k_in <= k && k <= (kx - k_in - 1)) && (l_in <= l && l <= (ly - l_in - 1)) then
     ee_cli_vor_tmp[l,k] = cli_vor_in[l-171,k-171]
     # ee_cli_vor_tmp[l,k] = cli_vor_in[l-l_in,k-k_in]
    else
      ee_cli_vor_tmp[l,k] = 0.0
    end
  end
end

# GPhys オブジェクトに値を入れる
ee_tro_vor.put(ee_tro_vor_tmp)
ee_cli_vor.put(ee_cli_vor_tmp)

ofile.close