#!/bin/sh
# Plugin to deal with cmorization of ecearth ocean and atmosphere outputs
#
# Written by Pierre-Antoine Bretonnière
# Institut Catala de Ciencies del Clima / Climate Forecasting Unit (IC3/CFU)
# Created: March 17, 2015
#
# List of functions:  *initialize_cmor
#                     *cmor_oce -> prepare_oce_input_dir
#                                  prepare_oce_namelist
#                                  cmor_oce_post
#                     *cmor_atm -> prepare_atm_input_dir
#                                  preprocess_grib_files
#                                  create_atm_namelist
######################################################################################

#######################################
# initialize_cmor
# Globals: CMORFAMILY, MEMBER, START_date, MODELS_DIR, CMORIZATION
# Arguments:
#   None
# Returns:
#   None
#######################################

function initialize_cmor () {

    experiment_id="${CMORFAMILY}"
    real=$(($(echo $MEMBER | cut -d 'c' -f 2)+1))
    init=1
    phys=1
    reftime=$(echo `echo $START_date | cut -c1-4 ``echo "-"``echo $START_date | cut -c5-6``echo "-"``echo $START_date | cut -c7-8``echo "(T00:00:00Z)"`) # yyyy-mm-dd(Thh:mm:ssZ)
   
} #end of function initialize_cmor

#######################################
# cmor_oce
# Globals: output_dir, cmoroi, cmoroo, EXPID, MODELS_DIR, CMORIZATION, Chunk_start_date, Chunk_end_date, RUN_months, 
# Arguments:
#   None
# Returns:
#   None
#######################################
function cmor_oce() {
  # available: rebuilt ocean output in *.nc (depends on iodef.xml defined)
  cmoroi=$output_dir/cmoro_input
  cmoroo=$output_dir/cmoro_output
  mkdir -p $cmoroi $cmoroo
  cd $cmoroi
  ln -sf $output_dir/${EXPID}_??_*.nc .
  ln -sf $PROJDIR/specs2cmor/common/Tables .
  startyear=$(echo $Chunk_start_date | cut -c1-4)
  endyear=$(echo $Chunk_end_date | cut -c1-4)
  startmonth=$(echo $Chunk_start_date | cut -c5-6)
  endmonth=$(echo $Chunk_end_date | cut -c5-6)
  set +e; i=1
  while [ $i -le $RUN_months ]; do
    ncra -d time_counter,$((($i-1)*$RUN_months+1)),$(($i*$RUN_months)) ${EXPID}_1d_${Chunk_start_date}_${Chunk_end_date}_diaptr.nc tmp${i}.nc
    i=$(($i+1))
  done
  rm -f ${EXPID}_1m_${Chunk_start_date}_${Chunk_end_date}_diaptr.nc
  ncrcat tmp* ${EXPID}_1m_${Chunk_start_date}_${Chunk_end_date}_diaptr.nc
  rm -f tmp*; set -e
  prepare_oce_namelist
  cmor-oce.exe namelist_cmoro
  cmor_oce_post
} #end of function cmor_oce

#######################################
# prepare_oce_namelist
# Globals: MODEL_ID, FORCING, real, startyear,endyear,startmonth,endmonth,init,INIT_DESCR,phys,PHYS_DESCR,ASSOC_MODEL,reftime,EXPID,ASSOCIATED_EXPERIMENT
# Arguments:
#   None
# Returns:
#   None
#######################################
function prepare_oce_namelist() {
  cat << EOF > namelist_cmoro
  &NAMEXP
  inpath='$cmoroi/',
  outpath='$cmoroo/',
  expname='$EXPID',
  experiment_id='$experiment_id',
  source='$MODEL_ID $VERSION official release from SMHI',
  model_id='$MODEL_ID',
  forcing='$FORCING',
  parent_experiment_id='N/A',
  parent_experiment_rip='N/A',
  institution='Barcelona Supercomputing Center, www.bsc.es',
  institute_id='BSC-CNS',
  realization=$real,
  history='',
  comment='',
  references='',
  contact='PA Bretonnière <pierre-antoine.bretonniere@bsc.es>',
  startyear=$startyear,
  endyear=$endyear,
  startmonth=$startmonth,
  endmonth=$endmonth,
  initialization_method=$init,
  initialization_description='$INIT_DESCR',
  physics_version=$phys,
  physics_description='$PHYS_DESCR',
  associated_model='$ASSOC_MODEL',
  forecast_reference_time='$reftime',
  prefix='$EXPID',
  associated_experiment='$ASSOCIATED_EXPERIMENT'
/
EOF
} #end of function prepare_oce_namelist

#######################################
# cmor_oce_post
# Globals: MODEL_ID, FORCING, real, startyear,endyear,startmonth,endmonth,init,INIT_DESCR,phys,PHYS_DESCR,ASSOC_MODEL,reftime,EXPID,associated_experiment
# Arguments:
#   None
# Returns:
#   None
#######################################
function cmor_oce_post() {
  cd $cmoroo
  tarfile="$output_dir/CMORO_${EXPID}_${START_date}_${MEMBER}_${Chunk_start_date}-${Chunk_end_date}.tar"
  if [[ ! -z $( ls ) ]]; then
    tar -cvf $tarfile *
    gzip -9 -f $tarfile; cd $output_dir
  else
    touch $tarfile ; cd $output_dir
  fi
  rm -rf $cmoroi $cmoroo
} #end of function cmor_oce_post 

#######################################
# cmor_atm
# Globals: output_dir, cmorai, cmorao, EXPID, MODELS_DIR, CMORIZATION, Chunk_start_date, Chunk_end_date, RUN_months, MEMBER
# Arguments:
#   None
# Returns:
#   None
#######################################
function cmor_atm() {
  # available: daily grib (3h or 6h) on monthly basis
  cmorai=$output_dir/cmora_input
  cmorao=$output_dir/cmora_output
  mkdir -p $cmorai $cmorao
  cd $cmorai
  ln -sf $PROJDIR/specs2cmor/common/Tables .
  preprocess_grib_files
  create_atm_namelist
  cmor-atm.exe namelist_cmora
  cmor_atm_post     
} #end of function cmor_atm

function preprocess_grib_files() {
   for file in $(ls $output_dir/ICM*grb); do
    ln -sf $file .
    sh=$( echo $file | sed -e s/GG/SH/g )
    ln -sf $sh .
  done
  NYEAR0=$(echo $START_date | cut -c1-4)
  NMONTH0=$(echo $START_date | cut -c5-6)
  list=""
  csy=$(echo $Chunk_start_date | cut -c1-4) # chunk start year
  csm=$(echo $Chunk_start_date | cut -c5-6) # chunk start month
  cs=$RUN_months # chunk size (in months)
  for c in $(seq 1 $cs); do
    csm=${csm#0}
    list=$( echo $list $(printf %4d%02d $csy $csm) )
    div=$((((csm+1))/13))
    rem=$((((csm+1))%13))
    csy=$((csy+div))
    csm=$rem
    [[ $csm -eq 0 ]] && csm=01
  done
  rm -rf tmp_dir_means
  mkdir -p tmp_dir_means
  cd tmp_dir_means
  file=fcdata
  ln -sf ../ICMSH${EXPID}*.grb .
  ln -sf ../ICMGG${EXPID}*.grb .
  for yyyymm in ${list}; do #199305 199306 199307 199308 ; do #${list[@]} ; do
    rm -rf tmp_dir_means fcdata fcdata0
    #means
    NYEAR=$(echo $yyyymm | cut -c1-4)
    NMONTH=$(echo $yyyymm | cut -c5-6)
    NMDAYS=$(cal $NMONTH $NYEAR | grep -v "[a-z]" | grep -v "^$" | wc -w)
    NSTP_DAY_R=$(( 24 / $NFRP ))
    NSTP_DAY=${NSTP_DAY_R%.*}
    param_6hr="151 167 168 164 165 166 129"
    param_day="167 165 166 151 164 168 169 177 179 228 201 202 130"
    param_mon="167 201 202 165 166 151 144 228 205 182 164 146 147 176 169 177 175 212 141 180 181 179 168 243 129 130 131 132 133"
    cdo_reftime=$( echo "${reftime:0:10},00:00" )
    echo FILTER_FILE ${file} ${yyyymm}
    echo "if ( dataDate >= ${yyyymm}01 ) { write ; }" > rules_files
     # get first timestep for each month from previous file (if possible)
    ym1=$(date -u -d "$NYEAR-$NMONTH-01 - 1 month" +%Y%m)
    for XX in SH GG; do
      rm -f ICM${XX}
      if [ -f ICM${XX}${EXPID}+${ym1}.grb ]; then
        grib_filter -o ICM${XX} rules_files \
                     ICM${XX}${EXPID}+${ym1}.grb ICM${XX}${EXPID}+${yyyymm}.grb
      else
        ln -sf ICM${XX}${EXPID}+${yyyymm}.grb ICM${XX}
      fi
    done
    # remap on regular Gauss grid
    cdo splitparam -sp2gpl ICMSH ${file}_
    cdo -R splitparam ICMGG ${file}_

    # total precipitation (remove negative values)
    cdo setcode,228 -setmisstoc,0 -setvrange,0,Inf -add ${file}_{142,143}.128.grb ${file}_228.128.grb

    # daily variables
    for param in $param_day; do
      new_units=""
      case $param in
        ( 169 | 177 | 179 )
        # radiation
        new_units="W m-2"
        cdo_operator="-divc,$((NFRP*3600)) -daymean -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 228 )
        # precipitation
        new_units="kg m-2 -s"
        cdo_operator="-mulc,1000 -divc,$((NFRP*3600)) -daymean -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 201 )
        # maximum
        cdo_operator="-daymax -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 202 )
        # minmimum
        cdo_operator="-daymin -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 130 )
        # 850 hPa
        cdo_operator="-daymean -sellevel,85000 -selmon,$NMONTH" ;;
        ( * )
        # default, plain daily mean
        cdo_operator="-daymean -selmon,$NMONTH" ;;
      esac
      cdo -f nc -t ecmwf setreftime,$cdo_reftime $cdo_operator ${file}_${param}.128.grb ${file}_${param}_day.nc

      if [ -n "$new_units" ]; then
        var=$(cdo showvar ${file}_${param}_day.nc | tr -d ' ')
        ncatted -O -a units,$var,m,c,"$new_units" ${file}_${param}_day.nc
      fi

      # concat all vars in one file for day
      ncks -A ${file}_${param}_day.nc ${EXPID}_day_${yyyymm}01-${yyyymm}${NMDAYS}.nc
    done

    # monthly variables
    for param in  $param_mon; do
      new_units=""
      case $param in
        ( 146 | 147 | 176 | 169 | 177 | 175 | 179 | 212 )
        # radiation/heat
        new_units="W m-2"
        cdo_operator="-divc,$((NFRP*3600)) -monmean -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 180 | 181 )
        # momentum flux
        new_units="N m-2"
        cdo_operator="-divc,$((NFRP*3600)) -monmean -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 144 | 228 | 205 | 182 )
        # precipitation/evaporation/runoff
        new_units="kg m-2 s-1"
        cdo_operator="-mulc,1000 -divc,$((NFRP*3600)) -monmean -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 201 )
        # mean daily maximum
        cdo_operator="-timmean -daymax -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 202 )
        # mean daily minmimum
        cdo_operator="-timmean -daymin -selmon,$NMONTH -shifttime,-${NFRP}hours" ;;
        ( 130 | 131 | 132 | 133 )
        # upper-air
        cdo_operator="-timmean -sellevel,5000,20000,50000,85000 -selmon,$NMONTH" ;;
        ( 129 )
        # upper-air geopotential
        new_units="m"
        cdo_operator="-divc,9.81 -timmean -sellevel,5000,20000,50000,85000 -selmon,$NMONTH" ;;
        ( * )
        # default, plain monthly mean
        cdo_operator="-timmean -selmon,$NMONTH" ;;
      esac

      cdo -f nc -t ecmwf setreftime,$cdo_reftime $cdo_operator ${file}_${param}.128.grb ${file}_${param}_mon.nc

      if [ -n "$new_units" ]; then
        var=$(cdo showvar ${file}_${param}_mon.nc | tr -d ' ')
        ncatted -O -a units,$var,m,c,"$new_units" ${file}_${param}_mon.nc
      fi

      # concat all vars in one file for mon
      ncks -A ${file}_${param}_mon.nc ${EXPID}_mon_${yyyymm}01-${yyyymm}${NMDAYS}.nc
    done

    # 6-hourly variables
    for param in  $param_6hr; do
      new_units=""
      case $param in
        ( 129 )
        # geopotential
        new_units="m"
        cdo_operator="-divc,9.81 -sellevel,50000 -selmon,$NMONTH" ;;
        ( * )
        # default, plain monthly mean
        cdo_operator="-selmon,$NMONTH" ;;
      esac

      cdo -f nc -t ecmwf setreftime,$cdo_reftime $cdo_operator \
      ${file}_${param}.128.grb ${file}_${param}_6hr.nc

      if [ -n "$new_units" ]; then
        var=$(cdo showvar ${file}_${param}_mon.nc | tr -d ' ')
        ncatted -O -a units,$var,m,c,"$new_units" ${file}_${param}_6hr.nc
      fi

      # concat all vars in one file for mon
      ncks -A ${file}_${param}_6hr.nc ${EXPID}_6hr_${yyyymm}01-${yyyymm}${NMDAYS}.nc
    done

    cp ${EXPID}_*nc ../.
    rm fcdata*grb
  done # loop over month
  cd ..
  date
  if [ "$NYEAR" == "1850" -a "$NMONTH" == 01 ]; then
    suff="1"
  else
    suff=""
  fi
} #end of function preprocess_grib_files 

#######################################
# prepare_atm_namelist
# Globals: MODEL_ID, FORCING, real, startyear,endyear,startmonth,endmonth,init,INIT_DESCR,phys,PHYS_DESCR,ASSOC_MODEL,reftime,EXPID,ASSOCIATED_EXPERIMENT
# Arguments:
#   None
# Returns:
#   None
#######################################
function create_atm_namelist() {
  cat << EOF > namelist_cmora
  &NAMEXP
  inpath='$cmorai/',
  outpath='$cmorao/',
  expname='$EXPID',
  experiment_id='$experiment_id',
  source='$MODEL_ID $VERSION official release from SMHI',
  model_id='$MODEL_ID',
  forcing='$FORCING',
  parent_experiment_id='N/A',
  parent_experiment_rip='N/A',
  institution='Barcelona Supercomputing Center, www.bsc.es',
  institute_id='BSC-CNS',
  realization=$real,
  history='',
  comment='',
  references='',
  contact='PA Bretonnière <pierre-antoine.bretonniere@bsc.es>',
  startyear=$NYEAR,
  endyear=$NYEAR,
  startmonth=$NMONTH,
  endmonth=$NMONTH,
  initialization_method=$init,
  initialization_description='$INIT_DESCR',
  physics_version=$phys,
  physics_description='$PHYS_DESCR',
  associated_model='$ASSOC_MODEL',
  forecast_reference_time='$reftime',
  prefix='$EXPID',
  associated_experiment='$ASSOCIATED_EXPERIMENT'
/
EOF
} #end of function create_atm_namelist

#######################################
# cmor_atm_post
# Globals: output_dir, cmorai, cmorao, EXPID, START_date, CMORIZATION, Chunk_start_date, Chunk_end_date, RUN_dir, chunk, MEMBER, PATHOUT_OUT
# Arguments:
#   None
# Returns:
#   None
#######################################
function cmor_atm_post() {
  unset yyyymms
  # post cmorization steps
  cd $cmorao
  tarfile="$output_dir/CMORA_${EXPID}_${START_date}_${MEMBER}_${Chunk_start_date}-${Chunk_end_date}.tar"
  if [[ ! -z $( ls ) ]]; then
    tar -cvf $tarfile *
    gzip -9 -f $tarfile; cd $output_dir
  else
    touch $tarfile ; cd $output_dir
  fi
  rm -rf $cmorai $cmorao
} #end of function cmor_atm_post

#######################################
# cmor_transfer
# Globals:
# Arguments: CP, EXPID, START_date, MEMBER, Chunk_start_date, Chunk_end_date, PATHOUT_OUT
#   None
# Returns:
#   None
#######################################
function cmor_transfer() {
  $CP CMOR?_${EXPID}_${START_date}_${MEMBER}_${Chunk_start_date}-${Chunk_end_date}.tar.gz $PATHOUT_OUT
} #end of cmor_transfer function

