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

require "pp"
require 'optparse'
require "numru/ggraph"
include NumRu

require "/GFD_Dennou_Work3/momoko/SyncRotEarthRad/script/load_config_ensemble"

# omegas は ../load_config から読み込む
#p omegas # = [0.0, 0.15, 0.5, 1.0]

# default
#TITLE = "Water vapor transport"
TITLE = ""

Latent = 2500000.0
wsn = 4

time_start = 1000
time_end   = 2000

position = 'net'

# (2013-8-23 石渡) 変更
val_max =  1.5
val_min = -1.5

type = 1  # 1: all, 2: disturbance

component = 1   # 1:x, 2:y, 3:r = sqrt(x^2+y^2)

clrmap = 1

lon = nil
interval = nil

def cut_x(gp, x)
  cut_x = gp.cut('lon'=>x)
end

def gp2axis(gp)
  data = gp.data
  va = VArray.new( gp.val,
                   {"long_name"=>data.long_name,
                     "units"=>data.units.to_s },
                   data.name )
  gp2axis = Axis.new.set_pos(va)
end

def gpdata_reuse_va(nary, gp)
  gpdata_reuse_va = VArray.new( nary,
                                {
                                  "long_name"=>gp.data.long_name,
                                  "units"=>gp.data.units.to_s},
                                gp.data.name )
end



# commandline option
opt = OptionParser.new

opt.on('--lon VAL') {|v| lon = v.to_f }
opt.on('--wsn VAL') {|v| wsn = v.to_i }
opt.on('--type VAL') {|v| type = v.to_i }
opt.on('--position VAL') {|v| position = v }
opt.on('--int VAL') {|v| interval = v.to_f }
opt.on('--clrmap VAL') {|v| clrmap = v.to_i }
opt.on('--component VAL') {|v|
  component = v.to_i
  if component < 1 and component > 3 then
    raise "ERROR: component #{component} is invalid."
  end
}
opt.on('--range VAL') {|v|
  ary = v.split(":")
  if ary.size == 2 and ary[0] < ary[1] then
    val_min = ary[0].to_f
    val_max = ary[1].to_f
  else
    raise "ERROR: #{v} is invalid range."
  end
  if interval == nil then
    interval = (val_max - val_min) / 40.0
  end
}

opt.parse!(ARGV)

if ARGV.size > 0
  raise "ERROR: invalid option: #{ARGV[0]}"
end

# 切り出す lon が未定義ならエラー吐いて死ぬ
raise("ERROR: --lon is not set.") if lon == nil

if type == 1 then
  puts "all"
elsif type == 2 then
  puts "disturbance"
else
  raise "Type #{type} is not supported."
end

var = 'U'
#path = File.join(exp_dirs[omegas[0]], var+'.nc')
path = File.join(exp_dirs[solcons[0]], var+'.nc')

p path

lat = GPhys::IO.open(path, 'lat')
nlat = lat.val.size
na_sig_weight = GPhys::IO.open(path, 'sig_weight').val
nsig = na_sig_weight.size

#na_flux_x_lon0 = NArray.sfloat(nomega, nlat)
#na_flux_y_lon0 = NArray.sfloat(nomega, nlat)
#na_flux_r_lon0 = NArray.sfloat(nomega, nlat)

na_flux_x_lon0 = NArray.sfloat(nsolcon, nlat)
na_flux_y_lon0 = NArray.sfloat(nsolcon, nlat)
na_flux_r_lon0 = NArray.sfloat(nsolcon, nlat)

#na_flux_x_lon180 = NArray.sfloat(nomega, nlat)
#na_flux_y_lon180 = NArray.sfloat(nomega, nlat)
#na_flux_r_lon180 = NArray.sfloat(nomega, nlat)

na_flux_x_lon180 = NArray.sfloat(nsolcon, nlat)
na_flux_y_lon180 = NArray.sfloat(nsolcon, nlat)
na_flux_r_lon180 = NArray.sfloat(nsolcon, nlat)


for solcon in solcons
#for omega in omegas
#p omega
p solcon
#  dir = exp_dirs[omega]
  dir = exp_dirs[solcon]
#  iomega = omegas.index(omega)
  isolcon = solcons.index(solcon)

  # input
  var = 'U'
  path = File.join(dir, var+'.nc')
  gp_u_lon0 = GPhys::IO.open(path, var).cut('time'=>time_start..time_end)
  gp_u_lon0 = cut_x(gp_u_lon0, 0)

  gp_u_lon180 = GPhys::IO.open(path, var).cut('time'=>time_start..time_end)
  gp_u_lon180 = cut_x(gp_u_lon180, 180)

  var = 'V'
  path = File.join(dir, var+'.nc')
  gp_v_lon0 = GPhys::IO.open(path, var).cut('time'=>time_start..time_end)
  gp_v_lon0 = cut_x(gp_v_lon0, 0)

  gp_v_lon180 = GPhys::IO.open(path, var).cut('time'=>time_start..time_end)
  gp_v_lon180 = cut_x(gp_v_lon180, 180)

  var = 'QVap'
  path = File.join(dir, var+'.nc')
  gp_q_lon0 = GPhys::IO.open(path, var).cut('time'=>time_start..time_end)
  gp_q_lon0 = cut_x(gp_q_lon0, 0)

  gp_q_lon180 = GPhys::IO.open(path, var).cut('time'=>time_start..time_end)
  gp_q_lon180 = cut_x(gp_q_lon180, 180)

  # 以降はライブラリ化したほうがいい

  gp_u_q_lon0 = gp_u_lon0 * gp_q_lon0
  gp_v_q_lon0 = gp_v_lon0 * gp_q_lon0

  gp_u_q_lon180 = gp_u_lon180 * gp_q_lon180
  gp_v_q_lon180 = gp_v_lon180 * gp_q_lon180

  # 時間平均
  gp_u_q_lon0 = gp_u_q_lon0.mean('time')
  gp_v_q_lon0 = gp_v_q_lon0.mean('time')
  if type == 2 then
    gp_u_lon0 = gp_u_lon0.mean('time')
    gp_v_lon0 = gp_v_lon0.mean('time')
    gp_q_lon0 = gp_q_lon0.mean('time')
  end

  gp_u_q_lon180 = gp_u_q_lon180.mean('time')
  gp_v_q_lon180 = gp_v_q_lon180.mean('time')
  if type == 2 then
    gp_u_lon180 = gp_u_lon180.mean('time')
    gp_v_lon180 = gp_v_lon180.mean('time')
    gp_q_lon180 = gp_q_lon180.mean('time')
  end

  # 擾乱成分
  if type == 2 then
    gp_ud_qd_lon0 = gp_u_q_lon0 - gp_u_lon0 * gp_q_lon0
    gp_vd_qd_lon0 = gp_v_q_lon0 - gp_v_lon0 * gp_q_lon0
  end

  if type == 2 then
    gp_ud_qd_lon180 = gp_u_q_lon180 - gp_u_lon180 * gp_q_lon180
    gp_vd_qd_lon180 = gp_v_q_lon180 - gp_v_lon180 * gp_q_lon180
  end

  # 鉛直積分
  if type == 1 then
    gp_flux_x_lon0 = (gp_u_q_lon0 * na_sig_weight.reshape(1,nsig)).sum('sig')
    gp_flux_y_lon0 = (gp_v_q_lon0 * na_sig_weight.reshape(1,nsig)).sum('sig')
  elsif type == 2 then
    gp_flux_x_lon0 = (gp_ud_qd_lon0 * na_sig_weight.reshape(1,nsig)).sum('sig')
    gp_flux_y_lon0 = (gp_vd_qd_lon0 * na_sig_weight.reshape(1,nsig)).sum('sig')
  end

  if type == 1 then
    gp_flux_x_lon180 = (gp_u_q_lon180 * na_sig_weight.reshape(1,nsig)).sum('sig')
    gp_flux_y_lon180 = (gp_v_q_lon180 * na_sig_weight.reshape(1,nsig)).sum('sig')
  elsif type == 2 then
    gp_flux_x_lon180 = (gp_ud_qd_lon180 * na_sig_weight.reshape(1,nsig)).sum('sig')
    gp_flux_y_lon180 = (gp_vd_qd_lon180 * na_sig_weight.reshape(1,nsig)).sum('sig')
  end

  na_flux_x_lon0[isolcon, true] = gp_flux_x_lon0.val
  na_flux_y_lon0[isolcon, true] = gp_flux_y_lon0.val

  na_flux_x_lon180[isolcon, true] = gp_flux_x_lon180.val
  na_flux_y_lon180[isolcon, true] = gp_flux_y_lon180.val

  if component == 3 then
    gp_flux_r_lon0 = gp_flux_x_lon0
    gp_flux_r_lon0.val = NMath.sqrt( na_flux_x_lon0[isolcon, true]**2 \
                                + na_flux_y_lon0[isolcon, true]**2 )
    na_flux_r_lon0[isolcon, true] = gp_flux_r_lon0.val
  end

  if component == 3 then
    gp_flux_r_lon180 = gp_flux_x_lon180
    gp_flux_r_lon180.val = NMath.sqrt( na_flux_x_lon180[isolcon, true]**2 \
                                + na_flux_y_lon180[isolcon, true]**2 )
    na_flux_r_lon180[isolcon, true] = gp_flux_r_lon180.val
  end

end

axis_lat = gp2axis(lat)

# うまくいかない
#omega_long_name = DCL::csgi(151)+'|*"'
#omega_long_name = "Omega*"
solcon_long_name = "Solar Constant"
solcon_units = "W/m**2"
va_solcon = VArray.new( NArray.to_na(solcons),
                         {
                           "long_name"=>solcon_long_name,
                           "units"=>solcon_units
                         },
                         "omega"
                         )
#axis_omega = Axis.new.set_pos(va_omega)
axis_solcon = Axis.new.set_pos(va_solcon)

#grid = Grid.new(axis_omega, axis_lat)
grid = Grid.new(axis_solcon, axis_lat)

if component == 1 then
  gp_flux_lon0 = GPhys.new( grid, 
                       gpdata_reuse_va(na_flux_x_lon0, gp_flux_x_lon0) )
  gp_flux_lon180 = GPhys.new( grid, 
                       gpdata_reuse_va(na_flux_x_lon180, gp_flux_x_lon180) )
elsif component == 2 then
  gp_flux_lon0 = GPhys.new( grid,
                       gpdata_reuse_va(na_flux_y_lon0, gp_flux_y_lon0) )
  gp_flux_lon180 = GPhys.new( grid,
                       gpdata_reuse_va(na_flux_y_lon180, gp_flux_y_lon180) )
elsif component == 3 then
  gp_flux_lon0 = GPhys.new( grid,
                        gpdata_reuse_va(na_flux_r_lon0, gp_flux_r_lon0) )
  gp_flux_lon180 = GPhys.new( grid,
                        gpdata_reuse_va(na_flux_r_lon180, gp_flux_r_lon180) )
else
  raise "ERROR: component #{component} is invalid."
end

if position == 'x0'
  gp_netflux = - gp_flux_lon0
elsif position == 'x180'
  gp_netflux = gp_flux_lon180
elsif position == 'net'
  gp_netflux = gp_flux_lon180 - gp_flux_lon0
else
  raise "ERROR: postion #{position} is invalid."
end


DCL.sgscmn(clrmap)  # colormap

DCL.gropn(wsn)
GGraph.set_fig('viewport'=>[0.15,0.85,0.25,0.6])  # same as gpview

DCL.sgpset('lfull', true)     # 全画面表示
DCL.sgpset('lcntl', false)
# うまくいかない
#DCL.sgpset('lcntl', true)

opt_draw = {
  'max'=>val_max, 'min'=>val_min,
  'interval'=>interval,
  'title'=>TITLE
#  'legend'=>false, 'annotate'=>false,
}


GGraph.tone(gp_netflux*Latent, true, opt_draw)
GGraph.contour(gp_netflux*Latent, false, opt_draw)
# (2013-8-23 石渡) 変更
#GGraph.color_bar('landscape'=>true)
GGraph.color_bar('landscape'=>true, 
                 'labelintv'=>10)

branch_x =[0.7, 0.7]
branch_y =[-85.0, 85.0]

DCL.uulinz(branch_x, branch_y,1,5)

DCL.grcls




__END__






DCL.gropn(wsn)

#DCL.sgpset('lcntl', false) ; DCL.uzfact(0.7)
DCL.sgpset('lcntl', true) ; DCL.uzfact(1.0)

# 球の時
#GGraph.set_fig 'itr'=>ITR, 'viewport'=>[0.15,0.95,0.2,0.8], 'map_axis'=>[LON_CNT, 0, 0]
GGraph.set_fig 'viewport'=>[0.15,0.95,0.2,0.8]

DCL.ugpset('LNRMAL', false)  # scale vector automatically

DCL.ugpset('XFACT1', xfact1)
DCL.ugpset('YFACT1', yfact1)

DCL.ugpset('UXUNIT', uxunit) # size of unit vector by dimentional value
DCL.ugpset('UYUNIT', uyunit)

DCL::ugrset('VXUNIT', 0.1)  # size of unit vector
DCL::ugrset('VYUNIT', 0.1)

DCL::uxmttl('T', '', 0.0)  # title (large character)

#DCL.ugpset('LMSG', false)   # no message
DCL.ugpset('LUMSG', false)   # no message

if type == 1 then
  # mean
  GGraph.vector( gp_u_q, gp_v_q, true, 'xintv'=>xintv, 'yintv'=>yintv,
                 'flow_vect'=>false, 'annotate'=>false, 'unit_vect'=>true
                 )
elsif type == 2 then
  # disturbance
  GGraph.vector( gp_ud_qd, gp_vd_qd, true, 'xintv'=>xintv, 'yintv'=>yintv,
                 'flow_vect'=>false, 'annotate'=>false, 'unit_vect'=>true
                 )
else
  raise "Type #{type} is not supported."
end


DCL.grcls