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

# Usage
#
#   $ ruby s2p.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.
#     --planet <planet name>
#        Planet name is given. It is used to identify pressure levels
#        used in interpolation.


def set_plev( planetname )

  case ( planetname )
  when "Earth"
    a_plev = [
      1000e2, 925e2, 850e2, 700e2, 600e2, 500e2, 400e2, 300e2, 250e2, 200e2, 150e2,
      100e2, 70e2, 50e2, 30e2, 20e2, 10e2
    ]
  when "Venus"
    a_plev = [
      9210000, 8645000, 8109000, 7601000, 7120000, 6665000, 6235000, 5828000, 5444000, 5081000, 4739000, 4416000, 4112000, 3826000, 3557000, 3304000, 3066000, 2843000, 2633000, 2436000, 2252000, 2079000, 1917000, 1766000, 1625000, 1493000, 1370000, 1256000, 1149000, 1050000, 958100, 872900, 794000, 721100, 653700, 591700, 534600, 482200, 434200, 390300, 350100, 313500, 280200, 249900, 222600, 197900, 175600, 155600, 137500, 121300, 106600, 93470, 81670, 71090, 61600, 53140, 45590, 38910, 33060, 27960, 23570, 16590, 11560, 7970, 5447, 3690, 2476, 1645, 1081, 701.099975585938, 447.600006103516, 280.799987792969, 173.300003051758, 105.300003051758, 63.1199989318848, 37.3600006103516, 21.9099998474121, 12.8100004196167, 7.51900005340576, 4.44999980926514, 2.66000008583069
    ]
  when "Mars"
    a_plev = [
      783.26,
      691.22, 610, 538.32, 475.07, 419.25, 369.98, 326.51, 288.14, 254.29,
      224.41, 198.04, 174.77, 154.23, 136.11, 120.12, 106, 93.547, 82.555,
      72.854, 64.294, 56.739, 50.072, 44.188, 38.996, 34.414, 30.37, 26.802,
      23.652, 20.873, 18.42, 16.256, 14.346, 12.66, 11.173, 9.8597, 8.7012,
      7.6788, 6.7765, 5.9802, 5.2775, 4.6574, 4.1101, 3.6272, 3.201, 2.8249,
      2.4929, 2.2, 1.9415, 1.7134, 1.512, 1.3344, 1.1776, 1.0392, 0.9171,
      0.80934, 0.71424, 0.63031, 0.55625, 0.49089, 0.43321, 0.3823, 0.33738,
      0.29774, 0.26275, 0.23188, 0.20463, 0.18059, 0.15937, 0.14064, 0.12412,
      0.10953, 0.096661, 0.085303, 0.07528, 0.066434, 0.058628, 0.051739,
      0.04566, 0.040294, 0.03556, 0.031381, 0.027694, 0.02444, 0.021568
    ]
  else
    print 'Unexpected name of planet: ', planetname, "\n"
    exit
  end

  return a_plev

end

###############################################################################
# interpolate_on_p
#   Values in GPhys object is interpolated on pressure levels.
#
# arguments:
# plev (array)           : pressure to be interpolated, unit is Pa
# gp_ps_part (GPhys obj) : surface pressure
# gp_part (GPhys obj)    : values to be interpolated
#
# return value:
# (GPhys obj) : interpolated values
#------------------------------------------------------------------------------

def interpolate_on_p( plev, gp_ps_part, gp_part )

  na_plev = NArray.to_na(plev)
  va_plev = VArray.new( na_plev, {"units"=>"Pa"}, "pressure")

  imax   = gp_part.coord('lon').val.size
  jmax   = gp_part.coord('lat').val.size
  kmax   = gp_part.coord('sig').val.size
  tmax   = gp_part.coord('time').val.size
  na_sig = gp_part.coord('sig').val

  # calculate pressure at each grid
  na_p =
    gp_ps_part.val.reshape(imax,jmax,1,tmax) *
    na_sig.reshape(1,1,kmax,1)
  va_p = VArray.new( na_p,
                     {"long_name"=>"air_pressure", "units"=>"Pa"},
                     "pressure" )
  # time axis of object is extracted
  ax_lon     = gp_part.axis('lon')
  ax_lat     = gp_part.axis('lat')
  ax_sig     = gp_part.axis('sig')
  ax_time_sl = gp_part.axis('time')
  # make GPhys object of pressure
  gp_p = GPhys.new( Grid.new(ax_lon,ax_lat,ax_sig,ax_time_sl), va_p )

  # set pressure as an assocated coordinate
  gp_part.set_assoc_coords([gp_p])

  # interpolate values on pressure levels
  gp_part_onplev = gp_part.interpolate("sig"=>va_plev)

  return gp_part_onplev

end


parser = GetoptLong.new

parser.set_options(
  ['--merge', '-m',              GetoptLong::NO_ARGUMENT],
  ['--planet',                   GetoptLong::REQUIRED_ARGUMENT],
#  ['--plev', '-p',               GetoptLong::REQUIRED_ARGUMENT],
)

$OPT_merge = false
$OPT_planet = "Earth"
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
  end
rescue
  exit(1)
end

a_plev = set_plev( $OPT_planet )

#print $OPT_merge, "\n"
#print $OPT_planet, "\n"
#print a_plev, "\n"
#exit


if ARGV.size < 3 then
  print "Usage : ruby s2p.rb Ps.nc IN.nc OUT.nc (options)\n"
  print "        Ps.nc  : NetCDF file for surface pressure\n"
  print "        IN.nc  : Input NetCDF file whose vertical level will be converged\n"
  print "        OUT.nc : Output NetCDF file\n"
  print "        Explanation for options can be found in the script.\n"
  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

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]
outncfn = ncfn_out

print "   Input (Ps) : ", ncfn_ps, "\n"
print "   Input      : ", ncfn, "\n"
print "   Output     : ", outncfn, "\n"
print "   plev       : ", a_plev, "\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 )
if $OPT_merge then
  url = ncfn[0..-4] + "_rank??????.nc@" + vname
else
  url = ncfn + "@" + vname
end
gp = GPhys::IO.open_gturl( url )

outfile = NetCDF.create(outncfn)
it = 0
GPhys::NetCDF_IO.each_along_dims_write(gp, outfile, -1) do |sub|
  gp_onplev = interpolate_on_p( a_plev, gp_ps[true,true,it..it], sub )
  gp_onplev.put_att('missing_value',[9.969209968386869e+36])
#    gp_onplev.put_att('_FillValue',[9.969209968386869e+36])
  it += 1
  [gp_onplev]
end
outfile.close