#!/bin/bash
set -u

# This script takes ORAS5 ocean reanalyses, reinterpolates them to the NEMO grid so that BDY data can be generated from it afterwards.
# As of July 2021, input ORAS5 can be downloaded from <https://www.cen.uni-hamburg.de/icdc/data/ocean/easy-init-ocean/ecmwf-oras5.html>
# Author: C. Pelletier

# Inclusive year range
year_start=1979
year_stop=2018

# Input folder where input data files (in .tar.gz format) are.
# This script expects member 'opa1' but you can change it below.
in_folder="/mfast/pelletie/oras5/raw"
# Folder where pre-interpolated datafiles will be 
dest_folder="/mfast/pelletie/oras5/preinterp"

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

# 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 ORAS5 to the Southern Ocean (we d
# Take a good 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).
eos10_script="/elic/home/pelletie/oras5/convert_eos10_new.py"

# NEMO destination CDO grid description (see `cdo_get_griddes.sh`).
dest_gridfile="/mfast/pelletie/oras5/eORCA025-SO_corners.grid"

pid="$$"

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

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

    # 1) Uncompress, concatenate, pre-cut and merge temperatures and salinities.
    
    tar -xvf ${in_folder}/votemper_ORAS5_1m_${year}_opa1.tar.gz
    cdo -O copy votemper_ORAS5_1m_${year}??_grid_T_02.nc tmp_pid${pid}.nc
    rm -f votemper_ORAS5_1m_${year}??_grid_T_02.nc
    cdo -O sellonlatbox,${latlon_cut} tmp_pid${pid}.nc tmp2_pid${pid}.nc

    tar -xvf ${in_folder}/vosaline_ORAS5_1m_${year}_opa1.tar.gz
    cdo -O copy vosaline_ORAS5_1m_${year}??_grid_T_02.nc tmp_pid${pid}.nc
    rm -f vosaline_ORAS5_1m_${year}??_grid_T_02.nc
    cdo -O sellonlatbox,${latlon_cut} tmp_pid${pid}.nc tmp3_pid${pid}.nc

    cdo -O merge tmp2_pid${pid}.nc tmp3_pid${pid}.nc tmp_pid${pid}.nc

    # 2) Convert EOS8 thermodynamics to TEOS-10 ones (look at the script to get why/what).
    module purge;
    module load ESMValTool;
    PYTHONPATH="/elic/home/pelletie/.local/lib/python3.6/site-packages:${PYTHONPATH}"
    python ${eos10_script} tmp_pid${pid}.nc tmp2_pid${pid}.nc

    # 3) Uncompress, concatenate, pre-cut and merge ssh with TEOS-10 temperatures and salinities.
    module purge
    module load 2018b CDO
    tar -xvf ${in_folder}/sossheig_ORAS5_1m_${year}_opa1.tar.gz
    cdo -O copy sossheig_ORAS5_1m_${year}??_grid_T_02.nc tmp_pid${pid}.nc
    rm -f sossheig_ORAS5_1m_${year}??_grid_T_02.nc
    cdo -O sellonlatbox,${latlon_cut} tmp_pid${pid}.nc tmp3_pid${pid}.nc
    cdo -O merge tmp3_pid${pid}.nc tmp2_pid${pid}.nc tmp_pid${pid}.nc


    # 4) Remap (bilinear) and drown all "T points" variables (temp/sal/ssh) to your NEMO grid.
    
    cdo remapbil,${dest_gridfile} tmp_pid${pid}.nc tmp2_pid${pid}.nc
    # For this interpolation, I have not put a src_gridfile, because this particular dataset (ORAS5) has 'good' netCDF attributes.
    # IF the input data has:
    # - a longitude variable with netCDF attributes standard_name="longitude", units="degrees East"
    # - a latitude variable with netCDF attributes standard_name="latitude", units="degrees North"
    # - at least one variable with a a netCDF attribute coordinates='<time> <depth> <lat> <lon>', where <time> is the name of the time variable; <depth> is the name of the depth variable (if the field is 3D); <lat> is the name of the latitude variable; <lon> is the name of the longitude variable.
    # THEN things go well, CDO recognizes the input grid.
    # And this happens to be the case with ORAS5 because it's a reanalysis from ECMWF, and they're standardized.
    # ELSE, you need to explicitly generate (see `get_cdo_griddes.sh`) and specify the input grid description in the remap call, which would have been: `cdo remapbil,${dest_gridfile} -setgrid,${src_gridfile} tmp_pid${pid}.nc tmp2_pid${pid}.nc`

    # Drown output fields. This is super lengthy because it computes distances (in kilometers).
    # There are other much more efficient ways to do it (discrete nearest-neighbor would be fine), but I'm too lazy.
    # I do have a function for doing it in nn_interp.py, though. Feel free to plug it in somehow.
    # Else, do it overnight/the weekend, I guess...    
    cdo setmisstonn tmp2_pid${pid}.nc ORAS5_preinterp_opa1_grid_T_${year}.nc
    
    module load NCO

    # 5) Uncompress, concatenate, pre-cut U velocities.
    tar -xvf ${in_folder}/vozocrtx_ORAS5_1m_${year}_opa1.tar.gz
    cdo -O copy vozocrtx_ORAS5_1m_${year}??_grid_U_02.nc tmp_pid${pid}.nc
    rm -f vozocrtx_ORAS5_1m_${year}??_grid_U_02.nc
    cdo -O sellonlatbox,${latlon_cut} tmp_pid${pid}.nc tmp3_pid${pid}.nc

    # 6) Get barotropic U velocities from full 3D velocities.
    # vozocrtx(t,z,y,x) = u_barotropic(t,y,x) + u_baroclinic(t,z,y,x), where u_barotropic is the vertical average of u_3d(t,z,y,x).
    # First, we get u_barotropic by vertical averaging, then we get u_baroclinic by removing it from vozocrtx.
    cdo vertmean tmp3_pid${pid}.nc tmp2_pid${pid}.nc
    ncks -C -v vozocrtx tmp2_pid${pid}.nc -O tmp_pid${pid}.nc
    ncwa -O -a depthu tmp_pid${pid}.nc tmp2_pid${pid}.nc
    ncrename -v vozocrtx,uobarotropic tmp2_pid${pid}.nc
    ncatted -a standard_name,uobarotropic,o,c,"barotropic_currents_in_x_direction" -a long_name,uobarotropic,o,c,"barotropic currents in x direction" tmp2_pid${pid}.nc
    ncks tmp2_pid${pid}.nc -A tmp3_pid${pid}.nc
    ncap2 -O -s 'uobaroclinic=vozocrtx-uobarotropic;' tmp3_pid${pid}.nc tmp2_pid${pid}.nc
    ncks -x -v vozocrtx tmp2_pid${pid}.nc -O tmp3_pid${pid}.nc
    ncatted -a standard_name,'baroclinic_currents_in_x_direction' -a long_name,uobaroclinic,o,c,"baroclinic currents in x direction" -a units,uobaroclinic,o,c,"m s-1" -a coordinates,uobaroclinic,o,c,"time_counter depthu nav_lat nav_lon" tmp3_pid${pid}.nc

    # 7) Remap both u velocities (barotropic and baroclinic) to destination grid, and drown.
    # Here they're remapped to T nodes, which is not correct (but it's not a big deal).
    # All interpolations and shifted by half a cell in the x direction, sorry.
    # If you want to correct that, you should generate a ${dest_gridfile} which points to U nodes.
    cdo remapbil,${dest_gridfile} tmp3_pid${pid}.nc tmp2_pid${pid}.nc
    cdo setmisstonn tmp2_pid${pid}.nc ORAS5_preinterp_opa1_grid_U_${year}.nc

    # 8) Uncompress, concatenate, pre-cut u velocities.
    tar -xvf ${in_folder}/vomecrty_ORAS5_1m_${year}_opa1.tar.gz
    cdo -O copy vomecrty_ORAS5_1m_${year}??_grid_V_02.nc tmp_pid${pid}.nc
    rm -f vomecrty_ORAS5_1m_${year}??_grid_V_02.nc
    cdo -O sellonlatbox,${latlon_cut} tmp_pid${pid}.nc tmp3_pid${pid}.nc

    # 9) Get barotropic V velocities from full 3D velocities.
    # vomecrty(t,z,y,x) = v_barotropic(t,y,x) + v_baroclinic(t,z,y,x), where v_barotropic is the vertical average of vomecrty(t,z,y,x).
    # First, we get v_barotropic by vertical averaging, then we get v_baroclinic by removing it from vomecrty.
    cdo vertmean tmp3_pid${pid}.nc tmp2_pid${pid}.nc
    ncks -C -v vomecrty tmp2_pid${pid}.nc -O tmp_pid${pid}.nc
    ncwa -O -a depthv tmp_pid${pid}.nc tmp2_pid${pid}.nc
    ncrename -v vomecrty,vobarotropic tmp2_pid${pid}.nc
    ncatted -a standard_name,vobarotropic,o,c,"barotropic_currents_in_x_direction" -a long_name,vobarotropic,o,c,"barotropic currents in x direction" tmp2_pid${pid}.nc
    ncks tmp2_pid${pid}.nc -A tmp3_pid${pid}.nc
    ncap2 -O -s 'vobaroclinic=vomecrty-vobarotropic;' tmp3_pid${pid}.nc tmp2_pid${pid}.nc
    ncks -x -v vomecrty tmp2_pid${pid}.nc -O tmp3_pid${pid}.nc
    ncatted -a standard_name,'baroclinic_currents_in_x_direction' -a long_name,vobaroclinic,o,c,"baroclinic currents in x direction" -a units,vobaroclinic,o,c,"m s-1" -a coordinates,vobaroclinic,o,c,"time_counter depthv nav_lat nav_lon" tmp3_pid${pid}.nc

    # 10) Remap both v velocities (barotropic and baroclinic) to destination grid, and drown.
    # Here they're remapped to T nodes, which is not correct (but it's not a big deal).
    # All interpolations and shifted by half a cell in the y direction, sorry.
    # If you want to correct that, you should generate a ${dest_gridfile} which points to V nodes.
    cdo remapbil,${dest_gridfile} tmp3_pid${pid}.nc tmp2_pid${pid}.nc
    cdo setmisstonn tmp2_pid${pid}.nc ORAS5_preinterp_opa1_grid_V_${year}.nc
    
done

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