#!/bin/bash
set -u

# This script takes EC-Earth outputs, reinterpolates them to the NEMO grid, and then generates BDY dta from it.
# Author: C. Pelletier

# Inclusive year range
year_start=1979
year_stop=2050

# Input prefix (with folder) for input data files (in .tar format).
in_prefix="/storage/elic/chamarro/data"
# Name of experience
name_exp="a1tq"
# Frequency
freq="1m"

# Working folder where pre-interpolated datafiles will be put
dest_folder="/storepelican/pelletie/nemo_bdy/bdydta"
dest_prefix="ECEarth_${name_exp}_${freq}_eORCA025-SO_BDYDTA_nbr1_"

pwd=`pwd`
# Python script for reformatting BDYDTA (final step)
reformat_script="${pwd}/generate_new_bdy_dta_onefile.py"

# min_lon, max_lon, min_lat, max_lat
# To precut the input data prior to interpolation (not required, but less expensive if your source code is much bigger than the destination one, which is the case from global to the Southern Ocean.
# Take a generous margin (4-5 nodes bigger than the destination grid limit)
latlon_cut="-180,180,-90,-28"

# # Script for converting EOS80 to TEOS-10 (look at the script comments for why it exists and what it does).
# not needed for EC-Earth, which is already in TEOS-10
# curr_dir=`pwd`
# eos10_script="${curr_dir}/convert_eos10_new.py"

# NEMO destination CDO grid description (see `cdo_get_griddes.sh`).
dest_gridfile_t="/storepelican/pelletie/nemo_bdy/grids/eORCA025-SO_grid_T.grid"
dest_gridfile_u="/storepelican/pelletie/nemo_bdy/grids/eORCA025-SO_grid_U.grid"
dest_gridfile_v="/storepelican/pelletie/nemo_bdy/grids/eORCA025-SO_grid_V.grid"

pid="$$"

mkdir -p ${dest_folder}
cd ${dest_folder}

for((year=year_start;year<=year_stop;year++)); do

    # Loop over input years.
    
    module purge
    module load 2018b CDO NCO

    # 1) Extract T, S and ssh (T grid).

    tar -xvf ${in_prefix}/${name_exp}/NEMO/MMO_${name_exp}_18500101_fc0_${year}0101-${year}1231.tar ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_T_3D.nc ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_T_2D.nc
    
    cdo selname,thetao,so ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_T_3D.nc tmp_pid${pid}.nc
    cdo selname,zos ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_T_2D.nc tmp2_pid${pid}.nc
    cdo merge tmp_pid${pid}.nc tmp2_pid${pid}.nc tmp3_pid${pid}.nc

    rm -f ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_T_3D.nc ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_T_2D.nc
    
    # precut
    cdo -O sellonlatbox,${latlon_cut} tmp3_pid${pid}.nc tmp_pid${pid}.nc

    ncrename -v olevel,deptht tmp_pid${pid}.nc
    ncrename -d olevel,deptht tmp_pid${pid}.nc
    
    # remap
    cdo remapbil,${dest_gridfile_t} tmp_pid${pid}.nc tmp2_pid${pid}.nc
    cdo setmisstonn tmp2_pid${pid}.nc tmp_pid${pid}.nc

    module purge
    module load ELIC/0.1-foss-2018b-Python-3.6.6

    input_full=`readlink -f tmp_pid${pid}.nc`
    output_full="${dest_folder}/${dest_prefix}grid_T_y${year}.nc"
    
    python ${reformat_script} ${input_full} t ${output_full}

    module purge
    module load CDO NCO    
    
    # 2) Extract u velocities.

    if [ "${freq}" == "1d" ]; then
	
	tar -xvf ${in_prefix}/${name_exp}/NEMO/MMO_${name_exp}_18500101_fc0_${year}0101-${year}1231.tar ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_U_3D.nc ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_U_2D.nc
	cdo selname,ubar ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_U_2D.nc tmp2_pid${pid}.nc

    else
	# ubar is not stored in monthly outputs, so take it from daily and monmean
	tar -xvf ${in_prefix}/${name_exp}/NEMO/MMO_${name_exp}_18500101_fc0_${year}0101-${year}1231.tar ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_U_3D.nc ${name_exp}_1d_${year}0101_${year}1231_opa_grid_U_2D.nc
	cdo selname,ubar ${name_exp}_1d_${year}0101_${year}1231_opa_grid_U_2D.nc tmp3_pid${pid}.nc
	cdo monmean tmp3_pid${pid}.nc tmp2_pid${pid}.nc
	
    fi    
    
    cdo selname,uo ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_U_3D.nc tmp_pid${pid}.nc
    cdo merge tmp_pid${pid}.nc tmp2_pid${pid}.nc tmp3_pid${pid}.nc

    rm -f ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_U_3D.nc ${name_exp}_??_${year}0101_${year}1231_opa_grid_U_2D.nc
    
    # precut
    cdo -O sellonlatbox,${latlon_cut} tmp3_pid${pid}.nc tmp_pid${pid}.nc

    ncap2 -O -s 'uobaroclinic = uo - ubar;' tmp_pid${pid}.nc tmp2_pid${pid}.nc
    ncks -x -v uo tmp2_pid${pid}.nc -O tmp_pid${pid}.nc
    ncrename -v ubar,uobarotropic tmp_pid${pid}.nc
    ncrename -v olevel,depthu -v olevel_bnds,depthu_bnds tmp_pid${pid}.nc
    ncatted -a bounds,depthu,o,c,"depthu_bnds" -a name,depthu,o,c,"depthu" tmp_pid${pid}.nc
    ncrename -d olevel,depthu tmp_pid${pid}.nc
    ncpdq -O -a time_counter,depthu,y,x tmp_pid${pid}.nc tmp2_pid${pid}.nc
    
    # remap
    cdo remapbil,${dest_gridfile_u} tmp2_pid${pid}.nc tmp_pid${pid}.nc
    cdo setmisstonn tmp_pid${pid}.nc tmp2_pid${pid}.nc

    module purge
    module load ELIC/0.1-foss-2018b-Python-3.6.6

    input_full=`readlink -f tmp2_pid${pid}.nc`
    output_full="${dest_folder}/${dest_prefix}grid_U_y${year}.nc"
    
    python ${reformat_script} ${input_full} u ${output_full}

    module purge
    module load CDO NCO    

    
    # 3) Extract v velocities.

    if [ "${freq}" == "1d" ]; then
	
	tar -xvf ${in_prefix}/${name_exp}/NEMO/MMO_${name_exp}_18500101_fc0_${year}0101-${year}1231.tar ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_V_3D.nc ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_V_2D.nc
	cdo selname,vbar ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_V_2D.nc tmp2_pid${pid}.nc

    else
	# vbar is not stored in monthly outputs, so take it from daily and monmean
	tar -xvf ${in_prefix}/${name_exp}/NEMO/MMO_${name_exp}_18500101_fc0_${year}0101-${year}1231.tar ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_V_3D.nc ${name_exp}_1d_${year}0101_${year}1231_opa_grid_V_2D.nc
	cdo selname,vbar ${name_exp}_1d_${year}0101_${year}1231_opa_grid_V_2D.nc tmp3_pid${pid}.nc
	cdo monmean tmp3_pid${pid}.nc tmp2_pid${pid}.nc
	
    fi    
    
    cdo selname,vo ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_V_3D.nc tmp_pid${pid}.nc
    cdo merge tmp_pid${pid}.nc tmp2_pid${pid}.nc tmp3_pid${pid}.nc

    rm -f ${name_exp}_${freq}_${year}0101_${year}1231_opa_grid_V_3D.nc ${name_exp}_??_${year}0101_${year}1231_opa_grid_V_2D.nc
    
    # precut
    cdo -O sellonlatbox,${latlon_cut} tmp3_pid${pid}.nc tmp_pid${pid}.nc

    ncap2 -O -s 'vobaroclinic = vo - vbar;' tmp_pid${pid}.nc tmp2_pid${pid}.nc
    ncks -x -v vo tmp2_pid${pid}.nc -O tmp_pid${pid}.nc
    ncrename -v vbar,vobarotropic tmp_pid${pid}.nc
    ncrename -v olevel,depthv -v olevel_bnds,depthv_bnds tmp_pid${pid}.nc
    ncatted -a bounds,depthv,o,c,"depthv_bnds" -a name,depthv,o,c,"depthv" tmp_pid${pid}.nc
    ncrename -d olevel,depthv tmp_pid${pid}.nc
    ncpdq -O -a time_counter,depthv,y,x tmp_pid${pid}.nc tmp2_pid${pid}.nc
    
    # remap
    cdo remapbil,${dest_gridfile_v} tmp2_pid${pid}.nc tmp_pid${pid}.nc
    cdo setmisstonn tmp_pid${pid}.nc tmp2_pid${pid}.nc 

    module purge
    module load ELIC/0.1-foss-2018b-Python-3.6.6

    input_full=`readlink -f tmp2_pid${pid}.nc`
    output_full="${dest_folder}/${dest_prefix}grid_V_y${year}.nc"
    
    python ${reformat_script} ${input_full} v ${output_full}

    module purge
    module load CDO NCO    

    
done


rm -rf ./*pid${pid}*