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

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

      IN.nc  : Input NetCDF file
      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.
      --varname <variable name>
        [optional] Variable name is given.
  USAGE
end

parser = GetoptLong.new

parser.set_options(
  ['--merge', '-m',              GetoptLong::NO_ARGUMENT],
  ['--planet',                   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_planet = "Earth"
$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 == "--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_planet, "\n"
#print $OPT_varname
#print a_plev, "\n"
#exit


if ARGV.size < 2 then
  puts print_usage
  exit
end


# 惑星表面圧力
# 入力
ncfn_in = ARGV[0]
# 出力
ncfn_out = ARGV[1]


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         : ", ncfn, "\n"
print "   Variable name : ", vname, "\n"
print "   Output        : ", outncfn, "\n"


# lat_weight
if $OPT_merge then
  a_w = []
  path = ncfn_in[0..-4] + "_rank??????.nc"
  a_fn = Dir.glob( path ).sort
  a_fn.each { |fn|
    url = fn + "@" + "lat_weight"
    gp_w = GPhys::IO.open_gturl( url )
    a_w_s = gp_w.val[0..(gp_w.val.size/2-1)].to_a
    a_w_n = gp_w.val[(gp_w.val.size/2)..(gp_w.val.size-1)].to_a
    a_w = a_w_s + a_w + a_w_n
  }
  na_w = NArray.to_na(a_w)
  na_w = na_w / na_w.sum(0)
else
  url = ncfn + "@" + 'lat_weight'
  gp_w = GPhys::IO.open_gturl( url )
  na_w = gp_w.val
  na_w = na_w / na_w.sum(0)
end

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)


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

#  p
#  p sub.rank
#  exit

  case sub.rank
  when 3 then
    na_w_reshaped = na_w.reshape!(1,na_w.size)
  when 4 then
    na_w_reshaped = na_w.reshape!(1,na_w.size,1)
  else
    print 'ERROR : Unexpected rank of data', "\n"
    exit
  end
  gp_gm = ( sub * na_w_reshaped ).sum(1).mean(0)

  itime += 1
  [gp_gm]
end
outfile.close