#!/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)
    PATH=$MODELS_DIR/$CMORIZATION/models/$MODEL/bin:${PATH}
   
} #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() {
  if [[ ! -a $output_dir/cmoro ]]; then
    # 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 $MODELS_DIR/$CMORIZATION/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
  fi #if ocean cmorization done
} #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='Institut Catala de Ciences del Clima, www.ic3.cat',
  institute_id='IC3',
  realization=$real,
  history='',
  comment='',
  references='',
  contact='PA Bretonnière <pierre-antoine.bretonniere@ic3.cat>',
  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"
  tar -cvf $tarfile *
  gzip -9 -f $tarfile; cd $output_dir
  rm -rf $cmoroi $cmoroo
  touch $output_dir/cmoro

} #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() {
  if [[ ! -a $output_dir/cmora ]]; then
    # 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 $MODELS_DIR/$CMORIZATION/common/Tables .
    preprocess_grib_files
    create_atm_namelist
    cmor-atm.exe namelist_cmora
    cmor_atm_post     
  fi #atmospheric cmorization done
} #end of function cmor_atm

function preprocess_grib_files() {
#   for file in $(ls /cfunas/exp/ecearth/gribfiles/$EXPID/${START_date}/${MEMBER}/outputs/*grb ); do
   for file in $(ls $outputdir/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
  done
  mkdir 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 \
       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 && $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='Institut Catala de Ciences del Clima, www.ic3.cat',
  institute_id='IC3',
  realization=$real,
  history='',
  comment='',
  references='',
  contact='PA Bretonnière <pierre-antoine.bretonniere@ic3.cat>',
  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',
/
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"
  tar -cvf $tarfile *
  gzip -9 -f $tarfile; cd $output_dir
  rm -rf $cmorai $cmorao
  touch $output_dir/cmora
} #end of function cmor_atm_post

function cmor_transfer() {

if [[ $CMORIZATION != '' && ! -a ${RUN_dir}/cmor_${chunk} ]]; then
   $CP CMOR?_${EXPID}_${START_date}_${MEMBER}_${Chunk_start_date}-${Chunk_end_date}.tar.gz $PATHOUT_OUT
   touch ${RUN_dir}/cmor_${chunk}
fi

} #end of cmor_transfer function
