=begin 
= 自作 GPhys 拡張モジュール
ここには NumRu::GPhys クラスに対して追加, 変更したメソッドを記述する.
本来はモジュール関数を定義すべきなのだろうが, こちらの方が簡単なので
とりあえずそうする. マナーに反する 等の問題があった場合は, とりやめる.
=end
require "numru/netcdf"
require "numru/ggraph"

include NumRu

class NArray # NArray???????????
  def cshift(n) # Fortran90?cshift????
    #if(n > x.size) ...
    y = self.dup
    y[n..-1] = self[0..-1-n]
    y[0..n-1] = self[-n..-1]
    return y
  end
end


class GPhys
  # gtstrm メソッド GPhys に追加
  def enhanced_strm_function(vname)
    grid = self.grid
    
    lat = grid.axis(0).pos.val      # [90.0, 87.5, ... , -87.5, -90.0]
    ps  = grid.axis(1).pos.val*100  # [100, 200, ... , 100000] milibar => N/m^2
    time = grid.axis(2).pos.val      # [90.0, 87.5, ... , -87.5, -90.0]
    date = grid.axis(3).pos.val   # [100, 200, ... , 100000] milibar => N/m^2
    data = self.data.val        

    pi = NMath::PI
    g  = 9.8
    r  = 6.370e6
    cos = NMath::cos(lat*pi*2.0/360.0)

    keisuu = 2*pi*r*cos/g

    axs = grid.shape[0]
    ays = grid.shape[1]
    ats = grid.shape[2]
    ads = grid.shape[3]


    # 流線関数 psi の初期化
    psi = NArray.float(axs, ays, ats, ads).fill!(0.0)

    # 台形公式による積分
    0.upto(ads-1) {|l|
      0.upto(ats-1) {|k|
	0.upto(axs-1) {|i|
	  0.upto(ays-1){|j|
	    v = data[i, j, k, l]
	    if j == 0 then
	      psi[i, j, k, l] = 0.5*(v+0)*ps[0]*keisuu[i]
	    else
	      dp = (ps[j] - ps[j-1])
	      psi[i, j, k, l] = 0.5*(v+data[i, j-1, k, l])*dp*keisuu[i] + psi[i, j - 1, k, l]
	    end
	  }
	}
      }
    }

    psi = GPhys.new(grid, VArray.new(psi).rename(vname))
    return psi
  end



  # gtstrm メソッド GPhys に追加
  def strm_function(vname)
    grid = self.grid
    
    lat = grid.axis(0).pos.val      # [90.0, 87.5, ... , -87.5, -90.0]
    ps  = grid.axis(1).pos.val*100  # [100, 200, ... , 100000] milibar => N/m^2
    data = self.data.val        

    axs = grid.shape[0]
    ays = grid.shape[1]

    pi = NMath::PI
    g  = 9.8
    r  = 6.370e6
    cos = NMath::cos(lat*pi*2.0/360.0)

    keisuu = 2*pi*r*cos/g

    # 流線関数 psi の初期化
    psi = NArray.float(axs, ays).fill!(0.0)

    # 台形公式による積分
    0.upto(axs-1) {|i|
	0.upto(ays-1){|j|
           v = data[i, j]
           if j == 0 then
              psi[i, j] = 0.5*(v+0)*ps[0]*keisuu[i]
           else
              dp = (ps[j] - ps[j-1])
              psi[i, j] = 0.5*(v+data[i, j-1])*dp*keisuu[i] + psi[i, j - 1]
           end
	  }
      }

    psi = GPhys.new(grid, VArray.new(psi).rename(vname))
    return psi
  end

  # gttheta メソッド GPhys に追加
  def theta(vname)
    grid = self.grid
    
    lon = grid.axis(0).pos.val  # [0.0, 2.5, ... , -359.5, -360.0]
    lat = grid.axis(1).pos.val  # [90.0, 87.5, ... , -87.5, -90.0]
    ps  = grid.axis(2).pos.val  # [10, 20, ... , 1000]
    time= grid.axis(3).pos.val  # [0, 600, 1200, 1800]
    date= grid.axis(4).pos.val  # []
    data = self.data.val        

    axs = grid.shape[0]
    ays = grid.shape[1]
    azs = grid.shape[2]
    ats = grid.shape[3]
    ads = grid.shape[4]

    # 温位 theta の初期化
#    theta = NArray.float(axs, ays, azs).fill!(0.0)
    theta = NArray.float(axs, ays, azs, ats, ads).fill!(0.0)

=begin
    # theta
    0.upto(axs-1) {|i|
      0.upto(ays-1){|j|
	0.upto(azs-1){|k|
	  t = data[i, j, k]
	  theta[i, j, k] = t*(1000.0/ps[k])**0.288
	}
      }
    }
=end

    # theta
    0.upto(azs-1){|k|
      t = data[true, true, k, true, true]
      theta[true, true, k, true, true] = t*(1000.0/ps[k])**0.288
    }

    gttheta = GPhys.new(grid, VArray.new(theta).rename(vname))
    return gttheta
  end

  # p_mean メソッド GPhys に追加
  def p_bar(vname)
    theta_bar = self.mean(0, 1).data.val
    grid = self.grid
    
    lon = grid.axis(0).pos.val  # [0.0, 2.5, ... , -359.5, -360.0]
    lat = grid.axis(1).pos.val  # [90.0, 87.5, ... , -87.5, -90.0]
    ps  = grid.axis(2).pos.val  # [10, 20, ... , 1000]
    data = self.data.val        

    axs = grid.shape[0]
    ays = grid.shape[1]
    azs = grid.shape[2]

    # p_bar の初期化
    p_bar = NArray.float(axs, ays, azs).fill!(0.0)

    # p_bar
    0.upto(axs-1) {|i|
      0.upto(ays-1){|j|
	theta = data[i, j, true]
	0.upto(azs-1){|k|
	  if k == 0 then
	    p_bar[i, j, k] = 1.0 - ((ps[k] - (theta[k] - theta_bar[k])/(fd_dp(theta, k, ps[k+1]-ps[k])))/ps[k])
	  elsif k == (azs-1) then
	    p_bar[i, j, k] = 1.0 - ((ps[k] - (theta[k] - theta_bar[k])/(bd_dp(theta, k, ps[k]-ps[k-1])))/ps[k])
	  else
	    p_bar[i, j, k] = 1.0 - ((ps[k] - (theta[k] - theta_bar[k])/(cd_dp(theta, k, ps[k+1]-ps[k-1])))/ps[k])
	  end
	}
      }
    }

    gtpbar = GPhys.new(grid, VArray.new(p_bar).rename(vname))
    return gtpbar
  end

  def cd_dp(v, k, dp) # center 
    (v[k+1] - v[k-1])/(dp)
  end

  def fd_dp(v, k, dp) # forward
    (v[k+1] - v[k])/(dp) 
  end

  def bd_dp(v, k, dp) # backward
    (v[k] - v[k-1])/(dp) 
  end

  def ape(vname)
    grid = self.grid
    
    kappa = 0.288

    lon = grid.axis(0).pos.val  # [0.0, 2.5, ... , -359.5, -360.0]
    lat = grid.axis(1).pos.val  # [90.0, 87.5, ... , -87.5, -90.0]
    ps  = grid.axis(2).pos.val  # [10, 20, ... , 1000]
    date= grid.axis(3).pos.val  # [10, 20, ... , 1000]
    time  = grid.axis(4).pos.val  # [10, 20, ... , 1000]
    theta = self.data.val        

    axs = grid.shape[0]
    ays = grid.shape[1]
    azs = grid.shape[2]

    # theta の平均値を 3 次元に拡張
    theta_bar1 = self.mean(0, 1).data.val
    theta_bar = NArray.float(axs, ays, azs).fill!(0.0)
    for i in 0..axs-1
      for j in 0..ays-1
	theta_bar[i, j, true] = theta_bar1
      end
    end

    # ps を 3 次元に拡張
    psn = NArray.sint(axs, ays, azs).fill!(0)
    for i in 0..axs-1
      for j in 0..ays-1
	psn[i, j, true] = ps
      end
    end

    # ape
    ape = theta_bar**2*psn**(kappa-1)*((theta-theta_bar)/theta_bar)**2

    gtape = GPhys.new(grid, VArray.new(ape).rename(vname))
    ngtape = gtape.integrate(0)
    ngtape = ngtape.integrate(0)
    ngtape = ngtape.integrate(0)
    return ngtape
  end

#------------------------------------------------------------------------------#

  def stspct_fft(vname=nil)
    #!!注意 !!#
    # 上記の点に関して修正すべし.
    if !vname
      vname = self.data.name
    end

    # グリッド情報取得   
    grid = self.grid    
    space = grid.axis(0).pos.val  # 空間グリッド
    time  = grid.axis(1).pos.val  # 時間グリッド
    axs = grid.shape[0] # 空間方向のデータ数
    ays = grid.shape[1] # 時間方向のデータ数

    # 複素空間フーリエ係数 f_k(t) を計算
    w = self.data.val                         # 元データ w(x,t)
    f_k = NArray.complex(axs, ays).fill!(0.0)  # f_k の箱を用意
    for i in 0..(axs-1)
      f_k[i,true] = FFTW.fftw(w[i, true], 1)/(ays/2)
    end

    # f_k(t) の実部 c_k, 虚部 s_k を取得
    c_k = f_k.real
    s_k = f_k.imag

    # c_k,s_k のパワースペクトル pc, ps を求める.
    pwc = NArray.complex(axs, ays).fill!(0.0)  # P_w(c) の箱を用意
    pws = NArray.complex(axs, ays).fill!(0.0)  # P_w(s) の箱を用意
    for i in 0..(ays-1)
      pwc[true,i] = FFTW.fftw(c_k[true, i], 1)/(axs/2) # forward FT
      pws[true,i] = FFTW.fftw(s_k[true, i], 1)/(axs/2) # forward FT
    end

    pc = (pwc).abs**2/2.0
    ps = (pws).abs**2/2.0


    # クアドラチャースペクトルを求める.
    qcs = (pwc.conj*pws).imag/2.0

    # パワースペクトル    
    p_ww = (pc + ps - 2*qcs)/4.0     # 西進波のパワースペクトル
    p_we = (pc + ps + 2*qcs)/4.0     # 東進波のパワースペクトル
    p_ww = p_ww[-1..0, true]         # 符号逆向き


    # 以下, 空間軸が偶数の場合と奇数の場合で分けなければならない.

    if ays/2 == 0
      nays = (ays+1)/2
    else
      nays = ays/2
    end
    if axs/2 == 0
      naxs = (axs+1)/2
    else
      naxs = axs/2
    end

    axs = 2*naxs - 1

    stspct = NArray.float(axs, nays)
    stspct[0..naxs-2, true] = p_ww[(naxs)..-2, 0..(nays-1)]
    stspct[naxs-1, true] = p_ww[-1, 0..(nays-1)] + p_we[0, 0..(nays-1)] # 波数 0
    stspct[naxs..-1, true] = p_we[1..naxs-1, 0..(nays-1)]
    stspct[true, 0] = 0 # 波数 0, 周波数 0
    stspct[0, true] = 0 # 波数 0, 周波数 0
    data = VArray.new(stspct).rename(vname)

    # 軸の生成
    aks = VArray.new(NArray.sfloat(axs).indgen!(0)-(naxs-1)).rename("wvn")
    aws = VArray.new(NArray.sfloat(nays).indgen!(0)/ays).rename("freq")

    xax = Axis.new().set_pos(aks)
    yax = Axis.new(true).set_cell_guess_bounds(aws).set_pos_to_center

    grid = Grid.new(xax, yax)

    # GPhys オブジェクトを生成
    gtspw = GPhys.new(grid, data)

    gtspw.coord(0).set_att('units', '1')
    gtspw.coord(1).set_att('units', '1/DAY')

    return gtspw
  end

#------------------------------------------------------------------------------#

#------------------------------------------------------------------------------#

  def stspct_coh(vname)
    # グリッド情報取得   
    grid = self.grid    
    space = grid.axis(0).pos.val  # 空間グリッド
    time  = grid.axis(1).pos.val  # 時間グリッド
    axs = grid.shape[0] # 空間方向のデータ数
    ays = grid.shape[1] # 時間方向のデータ数

    # 複素空間フーリエ係数 f_k(t) を計算
    w = self.data.val                         # 元データ w(x,t)
    f_k = NArray.complex(axs, ays).fill!(0.0)  # f_k の箱を用意
    for i in 0..(axs-1)
      f_k[i,true] = FFTW.fftw(w[i, true], 1)/(ays/2)
    end

    # f_k(t) の実部 c_k, 虚部 s_k を取得
    c_k = f_k.real
    s_k = f_k.imag

    # c_k,s_k のパワースペクトル pc, ps を求める.
    pwc = NArray.complex(axs, ays).fill!(0.0)  # P_w(c) の箱を用意
    pws = NArray.complex(axs, ays).fill!(0.0)  # P_w(s) の箱を用意
    for i in 0..(ays-1)
      pwc[true,i] = FFTW.fftw(c_k[true, i], 1)/(axs/2) # forward FT
      pws[true,i] = FFTW.fftw(s_k[true, i], 1)/(axs/2) # forward FT
    end

    pc = (pwc).abs**2/2.0            # C_k のパワースペクトル
    ps = (pws).abs**2/2.0            # S_k のパワースペクトル

    qcs = (pwc.conj*pws).imag/2.0    # クアドラチャースペクトル
    kcs = (pwc.conj*pws).real/2.0    # コスペクトル

    # パワースペクトル    
    p_ww = (pc + ps - 2*qcs)/4.0     # 西進波のパワースペクトル
    p_we = (pc + ps + 2*qcs)/4.0     # 東進波のパワースペクトル
    p_ww = p_ww[-1..0, true]         # 符号逆向き

    # P-R コヒーレンス
    coh_cs = (kcs**2 + qcs**2)/pc/ps # C_k, C_s のコヒーレンス
    coh_pr = NMath::sqrt(1-(1-coh_cs**2)/4*pc*ps/p_ww/p_we)
                                     # P-R コヒーレンス


    # 以下, 空間軸が偶数の場合と奇数の場合で分けなければならない.

    if ays/2 == 0
      nays = (ays+1)/2
    else
      nays = ays/2
    end
    if axs/2 == 0
      naxs = (axs+1)/2
    else
      naxs = axs/2
    end

    axs = 2*naxs - 1

    stspct = NArray.float(axs, nays)
    stspct[0..naxs-2, true] = p_ww[(naxs)..-2, 0..(nays-1)]
    stspct[naxs-1, true] = p_ww[-1, 0..(nays-1)] + p_we[0, 0..(nays-1)] # 波数 0
    stspct[naxs..-1, true] = p_we[1..naxs-1, 0..(nays-1)]
    stspct[true, 0] = 0 # 波数 0, 周波数 0
    stspct[0, true] = 0 # 波数 0, 周波数 0
    data = VArray.new(stspct).rename(vname)

    # 軸の生成
    aks = VArray.new(NArray.sfloat(axs).indgen!(0)-(naxs-1)).rename("wvn")
    aws = VArray.new(NArray.sfloat(nays).indgen!(0)/ays).rename("freq")

    xax = Axis.new().set_pos(aks)
    yax = Axis.new(true).set_cell_guess_bounds(aws).set_pos_to_center

    grid = Grid.new(xax, yax)

    # GPhys オブジェクトを生成
    gtspw = GPhys.new(grid, data)
    p gtspw.max
    return gtspw
  end

#------------------------------------------------------------------------------#

  def stspct_st_tr(vname)
    #!!注意 !!#
    # 上記の点に関して修正すべし.

    # グリッド情報取得   
    grid = self.grid    
    space = grid.axis(0).pos.val  # 空間グリッド
    time  = grid.axis(1).pos.val  # 時間グリッド
    axs = grid.shape[0] # 空間方向のデータ数
    ays = grid.shape[1] # 時間方向のデータ数

    # 複素空間フーリエ係数 f_k(t) を計算
    w = self.data.val                         # 元データ w(x,t)
    f_k = NArray.complex(axs, ays).fill!(0.0)  # f_k の箱を用意
    for i in 0..(axs-1)
      f_k[i,true] = FFTW.fftw(w[i, true], 1)/(ays/2)
    end

    # f_k(t) の実部 c_k, 虚部 s_k を取得
    c_k = f_k.real
    s_k = f_k.imag

p c_k.abs.max
p s_k.abs.max
   
    # c_k,s_k のパワースペクトル pc, ps を求める.
    pwc = NArray.complex(axs, ays).fill!(0.0)  # P_w(c) の箱を用意
    pws = NArray.complex(axs, ays).fill!(0.0)  # P_w(s) の箱を用意
    for i in 0..(ays-1)
      pwc[true,i] = FFTW.fftw(c_k[true, i], 1)/(axs/2) # forward FT
      pws[true,i] = FFTW.fftw(s_k[true, i], 1)/(axs/2) # forward FT
    end

    pc = (pwc).abs**2/2.0
    ps = (pws).abs**2/2.0

    qcs = (pwc.conj*pws).imag/2.0    # クアドラチャースペクトル
    kcs = (pwc.conj*pws).real/2.0    # コスペクトル

    # パワースペクトル    
    p_ww = (pc + ps - 2*qcs)/4.0     # 西進波のパワースペクトル
    p_we = (pc + ps + 2*qcs)/4.0     # 東進波のパワースペクトル

    # 定在波, 移動波のパワースペクトル
    var = ((pc-ps)**2)/4 + kcs**2
    p_st = NMath::sqrt(var)          # 定在波のパワースペクトル
    p_tr = (p_ww + p_we) - p_st      # 移動波のパワースペクトル
    p_tr = p_tr[-1..0, true]

    # 以下, 空間軸が偶数の場合と奇数の場合で分けなければならない.

    if ays/2 == 0
      nays = (ays+1)/2
    else
      nays = ays/2
    end
    if axs/2 == 0
      naxs = (axs+1)/2
    else
      naxs = axs/2
    end

    axs = 2*naxs - 1

    stspct = NArray.float(axs, nays)
    stspct[0..naxs-2, true] = p_tr[(naxs)..-2, 0..(nays-1)]
    stspct[naxs-1, true] = p_tr[-1, 0..(nays-1)] + p_st[0, 0..(nays-1)] # 波数 0
    stspct[naxs..-1, true] = p_st[1..naxs-1, 0..(nays-1)]
    stspct[true, 0] = 0 # 波数 0, 周波数 0
    stspct[0, true] = 0 # 波数 0, 周波数 0
    data = VArray.new(stspct).rename(vname)

    # 軸の生成
    aks = VArray.new(NArray.sfloat(axs).indgen!(0)-(naxs-1)).rename("wvn")
    aws = VArray.new(NArray.sfloat(nays).indgen!(0)/ays).rename("freq")

    xax = Axis.new().set_pos(aks)
    yax = Axis.new(true).set_cell_guess_bounds(aws).set_pos_to_center

    grid = Grid.new(xax, yax)

    # GPhys オブジェクトを生成
    gtspw = GPhys.new(grid, data)
    p gtspw.max
    return gtspw
  end


  # gphys オブジェクトを nc ファイルに保存するメソッド
  def save(output, global_attr=nil, vname=nil, var_attr=nil, history=$history, new_vname=nil)
    if !vname
      vname = self.data.name
    end

    nc = NetCDF.create(output)
    GPhys::NetCDF_IO.write(nc, self)

    # 大域属性付加

    global_attr.to_a.each {|attr| 
      nc.put_att(attr[0], attr[1])
    } if global_attr
    # 変数名変更
    if new_vname && vname
      nc.var(vname).name=new_vname 
      vname = new_vname
    end
    # 変数属性付加
    var_attr.to_a.each {|attr|
      nc.var(vname).put_att(attr[0], attr[1])
    }     if var_attr
    # history 付加
    now_hist = global_attr["history"]
    if now_hist then
      hist = now_hist + "\n" + history
      nc.put_att("history", hist)
    else
      nc.put_att("history", history)
    end

    nc.close
    return "saved as #{output}"

  end
end

def global_attr(ncfile)
  global_attr = Hash.new

  nc = NetCDF.open(ncfile, "r")
  nc.att_names.each do |attr| 
    global_attr[attr] = nc.att(attr).get
  end
  return global_attr
end


# 同じ次元構造をもつ複数の nc ファイルを指定した際に
# 平均したオブジェクトを返すメソッド
def mean_gphys(file, var)

  gphys_file = Array.new()

  file.each {|f|
  
    temp = GPhys::NetCDF_IO.open(f, var)
    gphys_file << temp

  }

  sum = gphys_file[0]

  1.upto(gphys_file.size-1) {|f|
    sum += gphys_file[f] 
  }

  p gphys = sum/file.size.to_f
  return gphys
end

def ask_overwrite(output)
  
  if File.exist?(output) then
    
    print "\n"
    print %Q{the file name "#{output}" is already exists current directory.\n}
    print "Do you sure over write? (y/n)    "

    loop{
      flag = STDIN.gets.chomp

      if flag == "y" then
	break
      elsif flag == "n" then
	p "Bye."
	exit
      else
	print "please answer yes or no. (y/n)     "
      end

    }
    
  end
end