#!/bin/bash
# 
# ##########################
# # computes SST indices   #
# ##########################
# 
# dependencies: cdo, nco, eofsvd_nc.bash, cor_reg_fields_new.bash
# 
# computes several climate indices, following the definitions provided in UCAR's climate data guide.
# See: https://climatedataguide.ucar.edu/climate-data/overview-climate-indices for definitions and reference datasets for comparison. 
# 
# For proper working, the input fields should be as follows:
# 
# for AMV, NINO indices and PDO: sst (use option "--varname alternative_varname" if the variable is called differently in your input
# 
# 
# usage:	compute_indices.bash -i infile -o outfile [further options]
# 
# options:      -i,--infile  		specify name of the infile.
#               -o,--outfile 		specify name of the outfile. String "INDEX" will be replaced by the choosen index (e.g. "--outfile HadISST_INDEX.nc")
# 		--amv-region flon,llon,flat,llat	specify region for AMV averaging. Default: 80W:0,0:65N
# 		--interpolate gridfile	interpolate to the grid specified in gridfile (could either be a cdo readable grid definition file, or any netcdf file with the desired target grid.
# 		--mask_land slm		apply a land sea mask given in file "lsm". to work properly, this file should containing a variable "sea" where every ocean grid point is set to 1,
# 					while all other grid points are set to missing.
# 					WARNING: automatically sets "--interpolate" to the grid of the land/sea mask!. Only really necessary, if land is not yet set as missvals.
# 		--pdo-region flon,llon,flat,llat	specify	region for PDO EOF deconstruction. Default: 120E:100W,20N:70N
# 		-x,--index indexname	specify, which index to be computed. Possible indices: PDO, AMV, NINO12, NINO3, NINO34, NINO4.
# 					If none specified, all implemented indices are computed.
#               -v             		set -ex
#               -h,--help     		prints this help
# 
# copyright     Ralf Hand (ralf.hand@giub.unibe.ch), 23.05.2020
# 
# __EOH__

function show_help {
   sed "/# __EOH__/q" ${0} | sed -e "s/# //" -e "s&#!/bin/bash&&" -e "s/__EOH__//"
}

# Reset all variables that might be set
infile=
outfile=
index=
interpolate=false
mask_land=false
varname=
verbose=0 # Variables to be evaluated as shell arithmetic should be initialized to a default or validated beforehand.

# set defaults (eventually overwritten by options!)

declare -A regions
regions[amv]=-80,0,0,65
regions[pdo]=120,-100,20,70
regions[nino12]=-90,-80,-10,0
regions[nino3]=-150,-90,-5,5
regions[nino34]=-170,-120,-5,5
regions[nino4]=160,-150,-5,5

sst_ready=false

oldpwd=$(pwd)

while :; do
    case $1 in
        -h|-\?|--help)   # Call a "show_help" function to display a synopsis, then exit.
            show_help
            exit
            ;;
        -i|--infile)       # Takes an option argument, ensuring it has been specified.
            if [ -n "$2" ]; then
                infile=$2
                shift 2
                continue
            else
                printf 'ERROR: "--infile" requires a non-empty option argument.\n' >&2
                exit 1
            fi
            ;;
        --infile=?*)
            infile=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --infile=)         # Handle the case of an empty --infile=
            printf 'ERROR: "--infile" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        -o|--outfile)       # Takes an option argument, ensuring it has been specified.
            if [ -n "$2" ]; then
                outfile=$2
                shift 2
                continue
            else
                printf 'ERROR: "--outfile" requires a non-empty option argument.\n' >&2
                exit 1
            fi
            ;;
        --outfile=?*)
            outfile=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --outfile=)         # Handle the case of an empty --outfile=
            printf 'ERROR: "--outfile" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        --amv-region)       # Takes an option argument, ensuring it has been specified.
            if [ -n "$2" ]; then
                regions[amv]=$2
                shift 2
                continue
            else
                printf 'ERROR: "--amv-region" requires a non-empty option argument.\n' >&2
                exit 1
            fi
            ;;
        --amv-region=?*)
            regions[amv]=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --amv-region=)         # Handle the case of an empty --amv-region=
            printf 'ERROR: "--amv-region" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        --interpolate)       # Takes an option argument, ensuring it has been specified.
            if [ -n "$2" ]; then
                interpolate=true
                grid=$2
                shift 2
                continue
            else
                printf 'ERROR: "--interpolate" requires a non-empty option argument.\n' >&2
                exit 1
            fi
            ;;
        --interpolate=?*)
            grid=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --interpolate=)         # Handle the case of an empty --amv-region=
            printf 'ERROR: "--interpolate" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        --mask_land)       # Takes an option argument, ensuring it has been specified.
            if [ -n "$2" ]; then
                mask_land=true
		lsm=$(readlink -f $2)
                interpolate=true
                grid=${lsm}
                shift 2
                continue
            else
                printf 'ERROR: "--mask_land" requires a non-empty option argument.\n' >&2
                exit 1
            fi
            ;;
        --mask_land=?*)
            lsm=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --mask_land=)         # Handle the case of an empty --amv-region=
            printf 'ERROR: "--mask_land" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        --pdo-region)       # Takes an option argument, ensuring it has been specified.
            if [ -n "$2" ]; then
                regions[pdo]=$2
                shift 2
                continue
            else
                printf 'ERROR: "--pdo-region" requires a non-empty option argument.\n' >&2
                exit 1
            fi
            ;;
        --pdo-region=?*)
            regions[pdo]=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --pdo-region=)         # Handle the case of an empty --pdo-region=
            printf 'ERROR: "--pdo-region" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        --varname)       # Takes an option argument, ensuring it has been specified.
            if [ -n "$2" ]; then
                varname=$2
                shift 2
                continue
            else
                printf 'ERROR: "--varname" requires a non-empty option argument.\n' >&2
                exit 1
            fi
            ;;

        --varname=?*)
            outfile=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --varname=)         # Handle the case of an empty --varname=
            printf 'ERROR: "--varname" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        -x|--index)
            if [ -n "$2" ]; then
               index=$2
               shift 2
               continue
            else
               printf 'ERROR: "--index" requires a non-empty option argument.\n' >&2
               exit 1
            fi
            ;;
        --index=?*)
            index=${1#*=} # Delete everything up to "=" and assign the remainder.
            ;;
        --index=)         # Handle the case of an empty --index=
            printf 'ERROR: "--index" requires a non-empty option argument.\n' >&2
            exit 1
            ;;
        -v|--verbose)
            verbose=$((verbose + 1)) # Each -v argument adds 1 to verbosity.
            ;;
        --)              # End of all options.
            shift
            break
            ;;
        -?*)
            printf 'WARN: Unknown option (ignored): %s\n' "$1" >&2
            ;;
        *)               # Default case: If no more options then break out of the loop.
            break
    esac

    shift
done


if [ ${verbose} -ge 1 ]; then set -ex; fi

infile=$(readlink -f ${infile})
outfile=$(readlink -f ${outfile})

echo "infile: ${infile}"
echo "outfile: ${outfile}"
echo "index: ${index}"

# eof="bash /mnt/climpa37/shell_scripts/eofsvd_nc.bash"
eof="eofsvd_nc.bash"

cdo=cdo

# tmp_pref=~/spielwiese
tmp_pref=/scratch3/ralf/spielwiese

tmp=$(mktemp -d -p ${tmp_pref} tmp.compute_indices.XXXX)

function prepare_sst {

   if ! ${sst_ready}; then

      if ${interpolate}; then 
         ${cdo} remapbic,${grid} -setrtoc,-9e9,-1.8,-1.8 -selvar,${varname} ${infile} infile.nc
      else
         ${cdo} setrtoc,-9e9,-1.8,-1.8 -selvar,${varname} ${infile} infile.nc
      fi # interpolate?

      if ${mask_land}; then
         ${cdo} -mul -selvar,sea ${lsm} infile.nc sst.nc
      else
         ln -sf infile.nc sst.nc
      fi # mask land?

   fi # are SSTs already prepared?

   sst_ready=true
}

cd ${tmp}

#######
# PDO #
#######

if [[ x${index} =~ (^x|[[:space:]])PDO($|[[:space:]]) || x${index} == x ]]; then
   
   if [ "${varname}" == "" ]; then varname=sst; fi

   prepare_sst

   ${cdo} enlarge,sst.nc -fldmean sst.nc globmean.nc
   ${cdo} sub sst.nc globmean.nc globmeansub.nc
   
   ${cdo} ymonsub globmeansub.nc -ymonmean globmeansub.nc anomalies.nc
   ${cdo} sellonlatbox,${regions[pdo]} anomalies.nc anomalies.pacific.nc
   
   ${eof} -v -n=1 --outfile_eof=pdo.pattern.nc --outfile_pc=pdo.ts.nc --outfile_expvar=pdo.expvar.nc anomalies.pacific.nc 
   
   ${cdo} yearmean pdo.ts.nc pdo.ts.yearly.nc

   ncatted -O -h -a comment,global,c,c,"PDO index produced by command ${0} $*" pdo.pattern.nc ${outfile/INDEX/PDO_pattern} 
   ncatted -O -h -a comment,global,c,c,"PDO index produced by command ${0} $*" pdo.ts.nc ${outfile/INDEX/PDO_index}
   ncatted -O -h -a comment,global,c,c,"PDO index produced by command ${0} $*" pdo.ts.yearly.nc ${outfile/INDEX/PDO_index_yearly}
   ncatted -O -h -a comment,global,c,c,"PDO index produced by command ${0} $*" pdo.expvar.nc ${outfile/INDEX/PDO_expvar}

   # compute power spectrum

fi # index is PDO (or all)



#######
# AMV #
#######

if [[ x${index} =~ (^x|[[:space:]])AMV($|[[:space:]]) || x${index} == x ]]; then
   
   if [ "${varname}" == "" ]; then varname=sst; fi  	

   prepare_sst

   ${cdo} yearmean sst.nc yearmean.nc
   ${cdo} fldmean yearmean.nc globmean.nc

   ${cdo} sub -fldmean -sellonlatbox,${regions[amv]} yearmean.nc globmean.nc globmeansub.nc
   ${cdo} sub globmeansub.nc -timmean globmeansub.nc amv.ts.yearly.nc
   ${cdo} lowpass,.1 amv.ts.yearly.nc amv.ts.nc
  
   cor_reg_fields_new.bash --normalize --nocorconfid --nocorsignif yearmean.nc amv.ts.nc amv.pattern.nc

   ncatted -O -h -a comment,global,c,c,"AMV index produced by command ${0} $*" amv.pattern.nc ${outfile/INDEX/AMV_pattern} 
   ncatted -O -h -a comment,global,c,c,"AMV index produced by command ${0} $*" amv.ts.yearly.nc ${outfile/INDEX/AMV_index_yearly}
   ncatted -O -h -a comment,global,c,c,"AMV index produced by command ${0} $*" amv.ts.nc ${outfile/INDEX/AMV_index}

   # compute power spectrum

fi # index is AMV (or all)


########
# ENSO #
########

if [[ x${index} == x || x${index} =~ (^x|[[:space:]])ENSO($|[[:space:]]) ]]; then 
   ninolist="NINO12 NINO3 NINO34 NINO4"
elif [[ "${index}" =~ "(^|[[:space:]])NINO*($|[[:space:]])" ]]; then
   ninolist=$(echo ${index} | tr ' ' "\n" | sed -e "/^NINO/! s/.*//" | tr "\n" ' ')
fi

function compute_nino {

   nino_kind=$1

   if [ "${varname}" == "" ]; then varname=sst; fi

   prepare_sst

   ${cdo} fldmean -sellonlatbox,${regions[$(echo ${nino_kind} | tr "[:upper:]" "[:lower:]")]} sst.nc nino.fldmean.nc
   ${cdo} runmean,5 -ymonsub nino.fldmean.nc -ymonmean nino.fldmean.nc nino.anomalies.smoothed.nc
   ${cdo} div nino.anomalies.smoothed.nc -timstd nino.anomalies.smoothed.nc nino.ts.nc
   ${cdo} yearmean nino.ts.nc nino.ts.yearly.nc

   last=$(( $(${cdo} -s ntime nino.ts.nc) + 2 ))
   ${cdo} seltimestep,3/${last} sst.nc sst.cutted.nc

   cor_reg_fields_new.bash --normalize --nocorconfid --nocorsignif sst.cutted.nc nino.ts.nc nino.pattern.nc
   
   ncatted -O -h -a comment,global,c,c,"${nino_kind} index produced by command ${0} $*" nino.pattern.nc ${outfile/INDEX/${nino_kind}_pattern} 
   ncatted -O -h -a comment,global,c,c,"${nino_kind} index produced by command ${0} $*" nino.ts.yearly.nc ${outfile/INDEX/${nino_kind}_index_yearly}
   ncatted -O -h -a comment,global,c,c,"${nino_kind} index produced by command ${0} $*" nino.ts.nc ${outfile/INDEX/${nino_kind}_index}

}

for nino_kind in ${ninolist}; do

   compute_nino ${nino_kind}

done # nino in ninolist

pwd
cd ${oldpwd}
# rm -rf ${tmp}

exit 0
