# coding: utf-8
#!/usr/bin/ruby
require "numru/gphys"
require 'getoptlong'
#require "./gp_dcpam_methods_v1.0"
include NumRu

def print_usage
  <<~USAGE
  Usage : 
    $ ruby vinteg.rb Ps.nc IN.nc OUT.nc (options)

      Ps.nc  : NetCDF file for surface pressure
      IN.nc  : Input NetCDF file whose vertical level will be converged
      OUT.nc : Output NetCDF file

    options:
      --merge
        Distributed files are used as input. Those files are merged
        and output is one file.
        Note that the name of NetCDF files to be merged is IN_rank??????.nc, 
        if this option is given.
     --grav <gravitational acceleration>
        Gravitational acceleration is given. 
     --varname <variable name>
        [optional] Variable name is given.
  USAGE
end


parser = GetoptLong.new

parser.set_options(
  ['--merge', '-m',              GetoptLong::NO_ARGUMENT],
  ['--grav',                     GetoptLong::REQUIRED_ARGUMENT],
  ['--varname',                  GetoptLong::OPTIONAL_ARGUMENT],
  ['--time_index_s', '--tis',    GetoptLong::REQUIRED_ARGUMENT],
  ['--time_index_e', '--tie',    GetoptLong::REQUIRED_ARGUMENT],
#  ['--plev', '-p',               GetoptLong::REQUIRED_ARGUMENT],
)

$OPT_merge = false
$OPT_grav = 9.81
$OPT_varname = ""
$OPT_time_index_s = 0
$OPT_time_index_e = -1
begin
  parser.each_option do |name, arg|
    eval "$OPT_#{name.sub(/^--/, '').gsub(/-/, '_')} = '#{arg}'"
#    print name, ":", arg, "\n"
    if name == "--merge" then
      $OPT_merge = true
    end
    if name == "--grav" then
      $OPT_grav = $OPT_time_index_s.to_f
    end
    if name == "--time_index_s" then
      $OPT_time_index_s = $OPT_time_index_s.to_i
    end
    if name == "--time_index_e" then
      $OPT_time_index_e = $OPT_time_index_e.to_i
    end
  end
rescue
  exit(1)
end

#print $OPT_merge, "\n"
#print $OPT_grav, "\n"
#print $OPT_varname
#exit

grav = $OPT_grav

if ARGV.size < 3 then
  puts print_usage
  exit
end


# 惑星表面圧力
ncfn_ps = ARGV[0]
vname_ps = "Ps"
# 入力
ncfn_in = ARGV[1]
# 出力
ncfn_out = ARGV[2]


if File.exist?(ncfn_out) then
  print "File, ", ncfn_out, " exists.\n"
  print "Overwrite the file? (yes/no)\n"
  input = $stdin.gets
  if input.chomp != 'yes' then
    print "STOP\n"
    exit
  end
end

ncfn = ncfn_in
outncfn = ncfn_out

if $OPT_varname.size > 1 then
  vname = $OPT_varname
else
  unless ncfn[-3..-1] == '.nc' then
    print "ERROR : Unexpected extention of file name: ", ncfn, "\n"
    exit
  end
  is = ncfn.rindex("/") != nil ? ncfn.rindex("/") : -1
  is += 1
  ie = -4
  vname = ncfn[is..ie]
end


print "   Input (Ps)    : ", ncfn_ps, "\n"
print "   Input         : ", ncfn, "\n"
print "   Variable name : ", vname, "\n"
print "   Output        : ", outncfn, "\n"
print "   Gravity       : ", grav, "\n"

if $OPT_merge then
  url = ncfn_ps[0..-4] + "_rank??????.nc@" + vname_ps
else
  url = ncfn_ps + "@" + vname_ps
end
gp_ps = GPhys::IO.open_gturl( url )
#
na_time = gp_ps.coord('time').val
times = na_time[$OPT_time_index_s]
timee = na_time[$OPT_time_index_e]
gp_ps = gp_ps.cut('time'=>times..timee)
#
if $OPT_merge then
  url = ncfn[0..-4] + "_rank??????.nc@" + vname
else
  url = ncfn + "@" + vname
end
gp = GPhys::IO.open_gturl( url )
#
na_time = gp.coord('time').val
times = na_time[$OPT_time_index_s]
timee = na_time[$OPT_time_index_e]
gp = gp.cut('time'=>times..timee)
#
if $OPT_merge then
  url = ncfn[0..-4] + "_rank000000.nc@" + 'sig_weight'
else
  url = ncfn + "@" + 'sig_weight'
end
gp_sig_weight = GPhys::IO.open_gturl( url )
na_sig_weight = gp_sig_weight.val

itime = 0
itimes = $OPT_time_index_s
if $OPT_time_index_e < 0 then
  itimee = $OPT_time_index_e + na_time.size
else
  itimee = $OPT_time_index_e
end
ntime = (itimee-itimes+1)

outfile = NetCDF.create(outncfn)
GPhys::NetCDF_IO.each_along_dims_write(gp, outfile, -1) do |sub|
  # https://qiita.com/hokaccha/items/3abd55aa23894b57ffd1#comment-00bb23380f8b3b7cc489
  progress = ((itime+1).to_f/ntime.to_f*100).round(3)
  print "working... "+progress.round(3).to_s+"%\r"
  STDOUT.flush
  if (itime+1) == ntime then
    print "\n"
  end
  gp_vinteg = ( sub * na_sig_weight.reshape(1,1,na_sig_weight.size,1) ).sum('sig')
  gp_vinteg *= gp_ps[true,true,itime] / grav
#  gp_onplev.put_att('missing_value',[9.969209968386869e+36])
  itime += 1
  [gp_vinteg]
end
outfile.close