#!/bin/bash

set -eu

# experiment name
expname=$1

# read configuration file
. $2

tmpdir=$tmpdir_root/ifs
mkdir -p $tmpdir

start_yyyymm=${start_yyyy}01
end_yyyymm=${end_yyyy}12

d=$resdir

nprocs=$nprocs_ifs

# frequency of PP output in IFS (needed to normalise accumulated variables)
ifs_pp_freq=$((6*3600))

# area for each ocean gridcell (undef over land)
area_file=/nobackup/rossby17/sm_wyser/ece2esmvaltool/area_a255.nc

# list of requested output variables
var_2d_list="clt pr pr-mmday rlut rsut tas psl hfls hfss tauu tauv zg-500 ta-200 ta-850 ua-200 ua-850 va-200 va-850 snd"
var_1d_list="clt tas pr-mmday pme-mmday rnds rndt"

# create parameter list
param_list=
for var in $var_2d_list $var_1d_list
do
    awk -v var=$var -f split_ptab.awk $ifs_par > $tmpdir/ifs-par-$var
    eval $(sed -n "/param[[:blank:]]*=/ p" $tmpdir/ifs-par-$var | tr -d [:blank:])
    param_list=$param_list" "$param
done
# remove duplicates
param_list=$(echo $param_list | tr '[:blank:]' '\n' | sort | uniq | tr '\n' ' ')

# create new parameter list, add parameters for derived variables
param_list_ifs=
for p in $param_list
do
    case $p in
    ( 228.128 | 103.228 ) param_list_ifs+=",142.128,143.128" ;;
    ( 101.128 ) param_list_ifs+=",142.128,143.128,182.128" ;;
    ( 97.128 ) param_list_ifs+=",178.128,212.128" ;;
    ( 98.128 ) param_list_ifs+=",179.128" ;;
    ( 99.128 ) param_list_ifs+=",178.128,179.128" ;;
    ( 100.128 ) param_list_ifs+=",176.128,177.128,146.128,147.128" ;;
    ( 102.128 ) param_list_ifs+=",141.128,33.128" ;;
    ( * ) param_list_ifs+=",$p" ;;
    esac
done
# remove duplicates
param_list_ifs=$(echo $param_list_ifs | tr ',' '\n' | sort | uniq | tr '\n' ',')
# remove leading and trailing comma
param_list_ifs=${param_list_ifs:1:${#param_list_ifs}-2}


# create list with filenames to be included
# NB: the approach with "find" may not be the best possible choice in case of
# very long simulations
rm -f  $tmpdir/icmgg_list $tmpdir/icmsh_list
for fname in $(find $d/ifs -regextype posix-extended -regex '.*/ICMGG.*+[0-9]{6}$' | sort )
do
    yyyymm=${fname:${#fname}-6:6}
    [ $yyyymm -ge $start_yyyymm ] && [ $yyyymm -le $end_yyyymm ] && echo $fname >> $tmpdir/icmgg_list
done

# create similar list for ICMSH files
sed -e's/ICMGG/ICMSH/g' $tmpdir/icmgg_list > $tmpdir/icmsh_list

# loop over all ICMSH files; execute in parallel
icount=0
while read fname
do
    (( icount+=1 ))
    echo "...now processing $(basename $fname)"

    yyyymm=${fname:${#fname}-6:6}
    yyyy=${yyyymm:0:4}
    mm=${yyyymm:4:2}

    # extract variables and compute monthly mean
    ( cdo -s splitparam -sp2gpl -setdate,$yyyy-$mm-01 -timmean \
        -selparam,130.128,131.128,132.128 -sellevel,20000,85000 \
        $fname $tmpdir/x_${yyyymm}_

      cdo -s splitparam -sp2gpl -setdate,$yyyy-$mm-01 -timmean \
        -selparam,129.128 -sellevel,50000 \
        $fname $tmpdir/x_${yyyymm}_ ) &

    if (( icount%nprocs == 0 )); then wait; fi 
done < $tmpdir/icmsh_list
wait

# loop over all ICMGG files; execute in parallel
icount=0
while read fname
do
    (( icount+=1 ))
    echo "...now processing $(basename $fname)"

    yyyymm=${fname:${#fname}-6:6}
    yyyy=${yyyymm:0:4}
    mm=${yyyymm:4:2}

    # extract variables and compute monthly mean
    cdo -s splitparam -setdate,$yyyy-$mm-01 -timmean \
        -selparam,$param_list_ifs $fname $tmpdir/x_${yyyymm}_ &

    if (( icount%nprocs == 0 )); then wait; fi 
done < $tmpdir/icmgg_list
wait

# concatenate all files
icount=0
for param in $(echo $param_list_ifs | tr ',' ' ')
do
    (( icount+=1 ))
    rm -f $tmpdir/x_${param}.grb
    cdo -s cat $tmpdir/x_*_${param}.grb $tmpdir/x_${param}.grb &
    if (( icount%nprocs == 0 )); then wait; fi 
done
wait

# change units
icount=0
for param in $param_list
do
    (( icount+=1 ))
    echo "...now processing parameter $param"

    ( case $param in
      ( 146.128 | 147.128 | 176.128 | 169.128 | 177.128 | 175.128 | 179.128 | 212.128 | 95.128 | 96.128 | 209.128 | 180.128 | 181.128 )
          # radiation/heat
          cdo_operator="divc,$ifs_pp_freq" ;;
      ( 97.128 )
          # upwelling SW radiation
          cdo -s setparam,97.128 \
              -divc,$ifs_pp_freq \
              -sub $tmpdir/x_212.128.grb $tmpdir/x_178.128.grb \
              $tmpdir/x_97.128.grb 
          cdo_operator="" ;;
      ( 98.128 )
          # upwelling LW radiation
          cdo -s setparam,98.128 \
              -divc,$ifs_pp_freq \
              -mulc,-1 $tmpdir/x_179.128.grb $tmpdir/x_98.128.grb
          cdo_operator="" ;;
      ( 99.128 )
          # net TOA radiation
          cdo -s setparam,99.128 \
              -divc,$ifs_pp_freq \
              -add $tmpdir/x_178.128.grb $tmpdir/x_179.128.grb \
              $tmpdir/x_99.128.grb
          cdo_operator="" ;;
      ( 100.128 )
          # net SFC energy flux
          cdo -s setparam,100.128 \
              -divc,$ifs_pp_freq \
              -add $tmpdir/x_146.128.grb \
              -add $tmpdir/x_147.128.grb \
              -add $tmpdir/x_176.128.grb $tmpdir/x_177.128.grb \
              $tmpdir/x_100.128.grb
          cdo_operator="" ;;
      ( 142.128 | 143.128 | 144.128 | 205.128 | 182.128 )
          # snow/precipitation/evaporation/runoff
          cdo_operator="mulc,1000 -divc,$ifs_pp_freq" ;;
      ( 228.128 ) 
          # total precipitation rate in kg m-2 s-1
          cdo -s setparam,228.128 \
              -mulc,1000 -divc,$ifs_pp_freq \
              -add $tmpdir/x_142.128.grb $tmpdir/x_143.128.grb \
              $tmpdir/x_228.128.grb 
          cdo_operator="" ;;
      ( 103.128 ) 
          # total precipitation in mm/day
          cdo -s setparam,103.228 \
              -mulc,86.4e6 -divc,$ifs_pp_freq \
              -add $tmpdir/x_142.128.grb $tmpdir/x_143.128.grb \
              $tmpdir/x_103.128.grb 
          cdo_operator="" ;;
      ( 101.128 ) 
          # precipitation minus evaporation in mm/day
          cdo -s setparam,101.128 \
              -mulc,86.4e6 -divc,$ifs_pp_freq \
              -add $tmpdir/x_182.128.grb \
              -add $tmpdir/x_142.128.grb $tmpdir/x_143.128.grb \
              $tmpdir/x_101.128.grb 
          cdo_operator="" ;;
      ( 164.128 ) 
          # cloud cover
          cdo_operator="mulc,100" ;;
      ( 129.128 ) 
          # geopotential meters
          cdo_operator="divc,9.81" ;;
      ( 102.128 ) 
          # snow depth (make sure we have snow density (33.128) or an estimate)
          [ -f $tmpdir/x_33.128.grb ] \
            && div_dens="-div $tmpdir/x_141.128.grb $tmpdir/x_33.128.grb" \
            || div_dens="-divc,350 $tmpdir/x_141.128.grb"
          cdo -s setparam,102.128 \
              -mulc,1000 $div_dens $tmpdir/x_102.128.grb
          cdo_operator="" ;;
      ( * )
          # default
          cdo_operator="" ;;
      esac

      if [ -n "$cdo_operator" ]
      then
          cdo -s $cdo_operator \
              $tmpdir/x_${param}.grb $tmpdir/y_${param}.grb
      else
          mv $tmpdir/x_${param}.grb $tmpdir/y_${param}.grb
      fi ) &
    if (( icount%nprocs == 0 )); then wait; fi 
done
wait

# split levels (if necessary) and save as compressed netCDF
icount=0
for var in $var_2d_list
do
    (( icount+=1 ))
    echo "...now saving $var"
    ( # find param in unique parameter table
      eval $(sed -n "/param[[:blank:]]*=/ p" $tmpdir/ifs-par-$var | tr -d [:blank:])

      # level and field type for ESMValTool
      case $var in
      ( ta-[0-9]* | ua-[0-9]* | va-[0-9]* | zg-[0-9]* )
          sellev="-sellevel,$((100*${var##*[a-z,A-Z]-}))"
          evt_ftype='T2M' ;;
      ( * )
          sellev=
          evt_ftype='T2Ms' ;;
      esac

      # name of output file
      outfile=$outdir/CMIP5_Amon_historical_${expname}_r1i1p1_${evt_ftype}_${var}_${start_yyyy}-${end_yyyy}.nc
      rm -f $outfile

      if [ $evt_ftype = 'T2M' ]
      then
          cdo -s -f nc4c -z zip -R setpartabp,$tmpdir/ifs-par-$var \
              -setreftime,$refdate,00:00:00,days \
              -invertlat $sellev \
              $tmpdir/y_${param}.grb $tmpdir/z_${var}.nc
          # remove unused vertical dimension
          ncwa -a lev $tmpdir/z_${var}.nc $outfile
      else
          cdo -s -f nc4c -z zip -R setpartabp,$tmpdir/ifs-par-$var \
              -setreftime,$refdate,00:00:00,days \
              -invertlat $sellev \
              $tmpdir/y_${param}.grb $outfile
      fi ) &
    if (( icount%nprocs == 0 )); then wait; fi 
done
wait

# split levels (if necessary) and save as compressed netCDF
icount=0
for var in $var_1d_list
do
    (( icount+=1 ))
    echo "...now saving $var"
    ( # find param in unique parameter table
      eval $(sed -n "/param[[:blank:]]*=/ p" $tmpdir/ifs-par-$var | tr -d [:blank:])

      # level and field type for ESMValTool
      case $var in
      ( rndt )
          evt_ftype='T1M' ;;
      ( * )
          evt_ftype='T1Ms' ;;
      esac

      # name of output file
      outfile=$outdir/CMIP5_Amon_historical_${expname}_r1i1p1_${evt_ftype}_${var}_${start_yyyy}-${end_yyyy}.nc
      rm -f $outfile

      cdo -s -f nc4c -z zip -R setpartabp,$tmpdir/ifs-par-$var \
          -setreftime,$refdate,00:00:00,days \
          -fldmean \
          $tmpdir/y_${param}.grb $outfile

      # ocean only
      case $var in
      ( rnds | pme-mmday | pr-mmday )
          outfile_oce=${outfile/.nc/_ocean.nc}
          rm -f $outfile_oce
          cdo -s -f nc4c -z zip setpartabp,$tmpdir/ifs-par-$var \
              -setreftime,$refdate,00:00:00,days \
              -div -fldsum -mul $tmpdir/y_${param}.grb $area_file \
              -fldsum $area_file $outfile_oce
      esac ) &
    if (( icount%nprocs == 0 )); then wait; fi 
done
wait

# clean up
rm -r $tmpdir

echo "!!! ece2evt-ifs done !!!"
exit 0