#!/bin/ksh

# ICON
#
# ------------------------------------------
# Copyright (C) 2004-2026, DWD, MPI-M, DKRZ, KIT, ETH, MeteoSwiss
# Contact information: icon-model.org
# See AUTHORS.TXT for a list of authors
# See LICENSES/ for license information
# SPDX-License-Identifier: BSD-3-Clause
# ------------------------------------------

#
# script to make mean over season or year and multiple ensemble members
#
# Martin Koehler April 2004
#--------------------------------------------------------------#

#. $HOME/.kshrc               # for troubleshooting ecp

set -ex                      # exit if error, verbose

#############################################################
# Setup
#############################################################

# --- argument list processing

while getopts a:l:f:d: option
do
   case $option in
    a) EXPVER=$OPTARG;;      # experiment ID  of EXP (e.g. e8s7)
    l) NLEV=$OPTARG;;        # number of levels (e.g. 60 or 91)
    f) FCLENGTH=$OPTARG;;    # forecast length [h]
    d) DATESTR=$OPTARG;;     # multiple initial times (e.g. 19870501/19870502/19870503)
   \?) errflg=1;;
   esac
done

if [ $errflg -ne 0 ] ; then
  echo "incorrect usage!" >&2
  exit 1
fi


# --- data input directory
#     (DATA directory for sms-job, scratch for interactive)

#DATADIR=${DATA:-/scratch/rd/$USER}
#DATADIR=${DATA}
#export DATADIR


#############################################################
# get data
#############################################################

# --- get time information for plotting
dayl=24
ppfrq=24
rcorioi=1.0
if [[ -f $SMSINCLUDE/ifsnam.fc.h && $FSFAMILY = /gfc ]] ; then
   rco=$(grep RCORIOI= $SMSINCLUDE/ifsnam.fc.h | sed 's/RCORIOI=//' |sed 's/.$//')
   if [[ $rco > 0 ]]; then
      rcorioi=$rco
   fi
fi
if [[ -f $SMSINCLUDE/config.fc.h && $FSFAMILY = /gfc ]] ; then
   pfr=$(grep PPFRQ= $SMSINCLUDE/config.fc.h | sed 's/PPFRQ=//')
   if [[ $pfr > 0 ]]; then
      ppfrq=$pfr
   fi
fi
dayl=$((dayl * rcorioi))
if [[ $dayl -lt 1 ]]; then
      dayl=1
fi


set -A nperiods `metview -b dates.met $FCLENGTH $DATESTR $dayl | grep nmon     | sed 's/nmon//'     | sed 'y/[],/   /'`
set -A yy_start `metview -b dates.met $FCLENGTH $DATESTR $dayl | grep yy_start | sed 's/yy_start//' | sed 'y/[],/   /'`
set -A mm_start `metview -b dates.met $FCLENGTH $DATESTR $dayl | grep mm_start | sed 's/mm_start//' | sed 'y/[],/   /'`
set -A mm_len   `metview -b dates.met $FCLENGTH $DATESTR $dayl | grep mm_len   | sed 's/mm_len//'   | sed 'y/[],/   /'`

set -A inidates `echo ${DATESTR} | sed 's/\// /g'`
nens=${#inidates[*]}


# --- read monthly means from experiment (multiple ensemble members)

#  for inidate in ${inidates[*]} ; do
#    currdir=`pwd`
#    cd $DATA
#    ecp -o ec:$FSROOT/${EXPVER}/${inidate}12/mmean_${EXPVER}_${inidate}.grb.tar .
#    tar xvf mmean_${EXPVER}_${inidate}.grb.tar
#    \rm mmean_${EXPVER}_${inidate}.grb.tar
#    cd $currdir
#  done


# --- loop over periods

iperiod=0
while (( ${iperiod} < ${nperiods} )) ; do


# --- loop over vartypes (sfc pl ml)

  set -A vartypes sfc pl ml
# set -A vartypes sfc
  for vartype in ${vartypes[*]}  ; do


    if (( ${mm_start[${iperiod}]} < 10 )) ; then
      iniyearmonth=${yy_start[${iperiod}]}"0"${mm_start[${iperiod}]}
    else
      iniyearmonth=${yy_start[${iperiod}]}${mm_start[${iperiod}]}
    fi
    yearmonth=${iniyearmonth}

# --- loop over months

    rm -f ${DATA}/mmean_${EXPVER}_${vartype}_period.grb

    imon=0
    while (( ${imon} < ${mm_len[${iperiod}]} )) ; do

# --- combine period data from all ensemble members (cat)

      cat   ${DATA}/mmean_${EXPVER}_${vartype}_ini*mon${yearmonth}.grb \
        >>  ${DATA}/mmean_${EXPVER}_${vartype}_period.grb

      month=`echo $yearmonth | cut -c 5-6`
      if [[ $month = 12 ]] ; then
        let yearmonth=yearmonth+100-11
      else
        let yearmonth=yearmonth+1
      fi

      let imon=imon+1
    done  # end of loop over months (getting data)

#   rm -f ${DATA}/mmean_${EXPVER}_${vartype}_ini*mon*.grb


#############################################################
# mean over ensemble and period
#############################################################

    currdir=`pwd`
    cd $DATA
    infile=mmean_${EXPVER}_${vartype}_period.grb
    meanfile=mmean_${EXPVER}_${vartype}_ensmean_n${nens}_${iniyearmonth}_n${mm_len[${iperiod}]}.grb
    rm -f ${meanfile}

# attention: the # of characters replaced by sed / / / have to be identical
    set -A DATE     `/usr/local/lib/metaps/bin/grib2request -1 -f $infile | grep DATE     | sed 'y/DATE=,/      /'`
    set -A STEP     `/usr/local/lib/metaps/bin/grib2request -1 -f $infile | grep STEP     | sed 'y/STEP=,/      /'`
    set -A PARAM    `/usr/local/lib/metaps/bin/grib2request -1 -f $infile | grep PARAM    | sed 'y/PARAM=,/       /'`
    set -A LEVELIST `/usr/local/lib/metaps/bin/grib2request -1 -f $infile | grep LEVELIST | sed 'y/LEVELIST=,/          /'`

#   set -A PARAM    `echo ${PARAM}    | sed 's/\// /g'`  # REPLACE / WITH SPACE
    set -A LEVELIST `echo ${LEVELIST} | sed 's/\// /g'`  #  - " -

    if [ ${vartype} = 'sfc' ] ; then
      LEVELIST=off
    fi

    echo 'date:     '$DATE
    echo 'step:     '$STEP
    echo 'param:    '$PARAM
    echo 'levelist: '$LEVELIST


    for LEVELIST1 in ${LEVELIST[*]} ; do
#   for PARAM1 in ${PARAM[*]} ; do
      PARAM1=$PARAM

      cat > mars.request1 << EOF
      retrieve,
        class=rd,
        type=fc,
        expver=$EXPVER,
        time=12,
        date=$DATE,
        step=$STEP,
        grid=2.5/2.5,
        param=$PARAM1,
        levtype=$vartype,
        levelist=$LEVELIST1,
        source="$infile",
        target="[param].temp.grb"
EOF
      mars mars.request1 || true

      for temp in *temp.grb ; do

        cat > mars.request2 << EOF
        read,
          source="$temp",
          fieldset=onefld
        compute,
          formula="mean(onefld)",
          fieldset=meanfld
        write,
          fieldset=meanfld,
          target=meanfld.grb
EOF
        mars mars.request2

        cat meanfld.grb >> ${meanfile}
#       rm -f temp.grb
        rm -f $temp

      done
    done

    cd $currdir

    rm -f ${DATA}/mmean_${EXPVER}_${vartype}_period.grb

  done    # end of loop over vartypes


# --- save to ecfs
#     (attention: first single ensemble member junk data will
#                 at the end be overwritten with the full data)

  currdir=`pwd`
  cd $DATA
  tar -cvf mmean_${EXPVER}_ensmean_${iniyearmonth}_${mm_len[${iperiod}]}.grb.tar \
         mmean_${EXPVER}_*_ensmean_${iniyearmonth}_${mm_len[${iperiod}]}.grb
  PUT      mmean_${EXPVER}_ensmean_${iniyearmonth}_${mm_len[${iperiod}]}.grb.tar
  \rm      mmean_${EXPVER}_ensmean_${iniyearmonth}_${mm_len[${iperiod}]}.grb.tar \
         mmean_${EXPVER}_*_ensmean_${iniyearmonth}_${mm_len[${iperiod}]}.grb
  cd $currdir


  let iperiod=iperiod+1
done    # end of loop over periods

# --- clean up


exit
