#!/bin/bash
set -evx 
# 
# This script interpolates vertically ocean restarts, extrapolates horizontally,
# fills up empty seas with climatologies, extrapolate vertically. 
#
# History : Virginie Guemas - Initial version                           2012 
#           Virginie Guemas - ORAS4 + vertical extrapolation + 
#                             commenting and clarifying         - March 2014
# 
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#                                Arguments
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
type='oras4' # 'glorys' / 'metoffice' : if you add another type, you'll have
             # to adapt the list of variables this restart contains below.
             # Look for the variable $type in the script.
fillinrot='TRUE' # Extrapolate or not the rotational rotn / rotb
                  # TRUE increases subtantially the interpolation time 
                  # FALSE has negligible effects ( u, v, t, s are initialized )
restart_in=$1 # Input restart file name without the .nc extension
meshmaskin='/cfu/autosubmit/con_files/mesh_mask_nemo.nemovar_O1L42_new.nc'
    # Meshmask of the input reanalysis. Example:
    #'/cfu/autosubmit/con_files/mesh_mask_nemo.glorys2v1_O25L75.nc'
confout='ORCA1L46' # output configuration
vertgridout='/cfu/autosubmit/con_files/mesh_mask_nemo.Ec3.0_O1L46.nc'
    # Meshmask of the output restart, i.e. model and grid to be used in
    # climate prediction 
Ptoextrap='/cfu/pub/scripts/prep_restarts/auxfiles/masks/Tofill_ecearth.v3.0.ORCA1L46.extrap.oras4.restarts.nc'
    # Locations of the points where to extrapolate horizontally the restarts after
    # interpolation = 1 in the netcdf file, 0 everywhere else. Locations available
    # for the tmask, umask, vmask and fmask grid. If you don't know about those
    # four variables read about the NEMO documentation about the grid.
    # To obtain the Ptoextrap netcdf file, it is necessary to run once this script
    # without extrapolation and filling of the empty seas and compare the output
    # restart with the output meshmask. To do so, set 0 everywhere in Ptoextrap
    # and use build_Tofill.bash afterwards.
Ptofillclim='/cfu/pub/scripts/prep_restarts/auxfiles/masks/Tofill_ecearth.v3.0.ORCA1L46.clim.nc'
    # Locations of the empty seas that need to be filled up with a climatology after
    # interpolation = 1 in the netcdf file, 0 everywhere else. Locations available
    # for the tmask, umask, vmask and fmask grid. If you don't know about those
    # four variables read about the NEMO documentation about the grid.
    # To obtain the Ptofillclim netcdf file, it is necessary to run once this script
    # without extrapolation and filling of the empty seas and compare the output
    # restart with the output meshmask. To do so, set 0 everywhere in Ptofillclim
    # and use build_Tofill.bash afterwards.
mon=$3 # Month of the restart on 2 digits MM
restart_out=$2 # Output restart file name without the .nc extension
cfutools='/home/vguemas/CFU_tools_new'
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

case $type in
  'glorys') 
# List of 3d variables = that need to be vertically interpolated 
   listvarsinterp=( 'en' 'avt' 'avmt' 'avmu' 'avmv' 'dissl' 'ub' 'vb' 'tb' 'sb' 'rotb' 'hdivb' 'un' 'vn' 'tn' 'sn' 'rotn' 'hdivn' 'rhop' )
# List of 2d variables = that only need extrapolation  
   listvarsextrap=( 'gcx' 'gcxb' 'gcxbfs' 'gru' 'grv' 'gsu' 'gsv' 'gtu' 'gtv' 'ssh_m' 'sshb' 'sshn' 'sss_m' 'sst_m' 'ssu_m' 'ssv_m' )
# List of 0d variables = that need to be copied to the output restart
   varscopy='nn_fsbc,rdt,rdttra1'
  ;;
  'oras4')
   listvarsinterp=( 'en' 'ub' 'vb' 'tb' 'sb' 'rotb' 'hdivb' 'un' 'vn' 'tn' 'sn' 'rotn' 'hdivn' )
   listvarsextrap=('gcx' 'gcxb' 'sshb' 'sshn' ) 
   varscopy='rdt,rdttra1,ndastp,adatrj'
  ;;
esac

source ${cfutools}/prep_restarts/library/library.bash

if [[ $fillinrot == 'FALSE' ]] ; then
 cp ${Ptoextrap} Ptoextrap.nc 
 ncks -v fmask Ptoextrap.nc tmp1.nc
 cdo mulc,0 tmp1.nc tmp2.nc
 ncks -A tmp2.nc Ptoextrap.nc
 rm -f tmp1.nc tmp2.nc
 Ptoextrap='Ptoextrap.nc'
fi

# Copy the 0d variables from the input restart to the output restart
ncks -O -v $varscopy ${restart_in}.nc ${restart_out}.nc
ncks -A -v nav_lon,nav_lat $meshmaskin ${restart_out}.nc
ncks -A -v nav_lev $vertgridout ${restart_out}.nc

# Vertical interpolation + horizontal extrapolation at each level + filling empty seas + vertical extrapolation
cdo vertsum -selvar,tmask,vmask,umask,fmask $meshmaskin mask2din.nc
for var in ${listvarsinterp[@]} ; do
  case $var in
   'avmu'|'un'|'ub') mask='umask' ; varlon='glamu' ; varlat='gphiu' ; fill='umask' ;;
   'avmv'|'vn'|'vb') mask='vmask' ; varlon='glamv' ; varlat='gphiv' ; fill='vmask' ;;
   'rotb'|'rotn') mask='fmask' ; varlon='glamf' ; varlat='gphif' ; fill='fmask' ;;
   'avt'|'avmt'|'en'|'dissl'|'hdivb'|'hdivn'|'rhop'|'sn'|'tn'|'sb'|'tb') mask='tmask' ; varlon='glamt' ; varlat='gphit' ;  fill='tmask' ;;
  esac
  # The vertical interpolation is performed below
  python ${cfutools}/interpolation/interp_vert.py ${restart_in}.nc $var $meshmaskin e3t_0 mask2din.nc $mask $vertgridout e3t_0 gdept_0 ${var}_${restart_out}_int.nc

  # The function extrap extrapolates horizontally and fill in the empty seas with a climatology
  extrap ${var}_${restart_out}_int.nc ${var} ${meshmaskin} ${varlon} ${varlat} ${Ptoextrap} ${Ptofillclim} ${fill} ${var}_${restart_out}_int2.nc 3d $cfutools $confout $mon 

  # The vertical extrapolation to empty levels is performed below
  python ${cfutools}/interpolation/vertextrap.py  ${var}_${restart_out}_int2.nc ${var} $vertgridout nav_lev ${var}_${restart_out}_int3.nc

  # Apply the mask
  applymask $vertgridout $mask ${var}_${restart_out}_int3.nc $var ${var}_${restart_out}.nc

  # Pasting the output to ${restart_out}.nc
  gathervar ${restart_out}.nc ${var}_${restart_out}.nc
  rm -f ${var}_${restart_out}.nc ${var}_${restart_out}_int.nc ${var}_${restart_out}_int2.nc ${var}_${restart_out}_int3.nc
done

# Horizontal extrapolation and filling empty seas of 2d variables
ncks -O -d lev,0 ${Ptoextrap} Ptoextrap_lev0.nc
ncks -O -d lev,0 ${Ptofillclim} Ptofillclim_lev0.nc
ncks -O -d z,0 $vertgridout mask_lev0.nc
ncks -O -d z,0 $meshmaskin maskin_lev0.nc
for var in ${listvarsextrap[@]} ; do
  case $var in
   'gru'|'gsu'|'gtu'|'ssu_m'|'u_io') varlon='glamu' ; varlat='gphiu' ; fill='umask' ; mask='umask' ;;
   'grv'|'gsv'|'gtv'|'ssv_m'|'v_io') varlon='glamv' ; varlat='gphiv' ; fill='vmask' ; mask='vmask' ;;
   'gcx'|'gcxb'|'gcxbfs'|'ssh_m'|'sshb'|'sshn'|'sss_m'|'sst_m'|'alb_ice'|'sst_io'|'sss_io') varlon='glamt' ; varlat='gphit' ; fill='tmask' ; mask='tmask' ;;
  esac

  # Apply the input mask
  applymask maskin_lev0.nc $mask ${restart_in}.nc $var ${var}_${restart_in}.nc

  # Horizontal extrapolation and fillup the empty seas with a climatology
  extrap ${var}_${restart_in}.nc ${var} ${meshmaskin} ${varlon} ${varlat} Ptoextrap_lev0.nc Ptofillclim_lev0.nc ${fill} ${var}_${restart_out}_int.nc 2d $cfutools $confout $mon
 
  # Apply the output mask
  applymask mask_lev0.nc $mask ${var}_${restart_out}_int.nc $var ${var}_${restart_out}.nc

  # Pasting the output to ${restart_out}.nc
  gathervar ${restart_out}.nc ${var}_${restart_out}.nc
  rm -f ${var}_${restart_out}.nc ${var}_${restart_out}_int.nc ${var}_${restart_in}.nc
done
rm -f Ptoextrap_lev0.nc Ptofillclim_lev0.nc mask2din.nc mask_lev0.nc tmpmask.nc maskin_lev0.nc
if [[ -e Ptoextrap.nc ]] ; then rm -f Ptoextrap.nc ; fi