#!/bin/bash
#
# computes and plots some simple ensemble statistics for the HIPPO_r0XXiXXXXpXfXXXX ensemble
# 
# ralf.hand@giub.unibe.ch 2020-03-25

set -ex 

# # normal time series
# 
dir=/mnt/climstor/ERC_PALAEO
dir=/tank/exports/ERC_PALAEO
#setlist="HIPPO_set_1420_1 HIPPO_set_1400_2"
setlist="HIPPO_set_1420_1  HIPPO_set_1420_2 HIPPO_set_1420_3 HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1780_3 HIPPO_set_1850_all_HadISST HIPPO_set_1850_1 HIPPO_set_1850_2  HIPPO_set_1850_3 HIPPO_set_1850_4"
#setlist="HIPPO_set_1850_1 HIPPO_set_1850_2 HIPPO_set_1850_3 HIPPO_set_1850_4"
#setlist="HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1780_3 HIPPO_set_1780_all"
#setlist="HIPPO_set_1420_2 HIPPO_set_1420_3 HIPPO_set_1420_LR_all "
#setlist="HIPPO_set_1850_all_HadISST HIPPO_set_1780_3 HIPPO_set_1780_all"

# varlist="hadc_nh_loc hadc_sh_loc hadc_nh_str hadc_sh_str"
# level="50000"
#tslist="mon year jja djf"
# tslist="mon year djf mam jja son"
infile_struct=${dir}/SET/RUN/timeseries/TS/RUN_VAR_TS.nc
outfile_struct=${dir}/SET/ensstat/timeseries/TS/SET_VAR_ENSSTAT_TS.nc

# setlist="HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1850_all_HadISST HIPPO_set_1850_20_members"
# setlist="HIPPO_set_1420_1 HIPPO_set_1420_2 HIPPO_set_1420_1_and_2 HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1780_1_and_2"
#setlist="HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1780_1_and_2 HIPPO_set_1780_all HIPPO_set_1850_1 HIPPO_set_1850_2 HIPPO_set_1850_3 HIPPO_set_1850_4"
#setlist="HIPPO_set_1780_3 HIPPO_set_1780_all"
varlist="tsurf temp2 totprec slp u v omega geopoth"
level=surf
tslist="mon"
# infile_struct=${dir}/SET/RUN/timeseries/TS/RUN_VAR_TS.nc
# outfile_struct=${dir}/SET/ensstat/timeseries/TS/SET_VAR_ENSSTAT_TS.nc

# # # region means
# # 
# dir=/scratch3/ralf/experiments/PALAEO-RA
# # # setlist="HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1850_all_HadISST HIPPO_set_1850_20_members"
# setlist="HIPPO_set_1420_1 HIPPO_set_1420_2 HIPPO_set_1420_1_and_2 HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1780_1_and_2"
# setlist="HIPPO_set_1780_1_and_2"
# setlist="HIPPO_set_1420_1 HIPPO_set_1420_2 HIPPO_set_1400_2 HIPPO_set_1420_LR_all"
# setlist="HIPPO_set_1420_LR_all"
# setlist="HIPPO_set_1420_1 HIPPO_set_1420_2 HIPPO_set_1420_3 HIPPO_set_1420_LR_all HIPPO_set_1780_1 HIPPO_set_1780_2 HIPPO_set_1780_3 HIPPO_set_1780_all HIPPO_set_1850_all_HadISST HIPPO_set_1850_1 HIPPO_set_1850_2  HIPPO_set_1850_3 HIPPO_set_1850_4"
# # # tslist="coldseas son djf mam jja mon year"
# # # tslist="ondj fmam"
# tslist="year djf mam jja son"
# # #level=surf
# tslist="year"
# # #varlist="temp2 slp albedo albedo_vis"
# varlist="slp temp2 totprec"
# varlist="temp2 tsurf"
# regionlist="NEU MED EAS WSB ENA CNA WNA NCA global"
regionlist="global"
# # #infile_struct=${dir}/SET/RUN/timeseries/TS/RUN_VAR_TS.nc
# # #outfile_struct=${dir}/SET/ensstat/timeseries/TS/SET_VAR_ENSSTAT_TS.nc
# # 
# # plot=true
# # # 
# infile_struct=${dir}/SET/RUN/indices/region_fldmean/TS/RUN_VAR_TS_fldmean_REGION.nc
# outfile_struct=${dir}/SET/ensstat/indices/region_fldmean/TS/SET_VAR_TS_fldmean_REGION_ENSSTAT.nc

# infile_struct=${dir}/SET/RUN/indices/region_fldmean/TS/RUN_VAR_TS_fldmean_REGION_land.nc
# outfile_struct=${dir}/SET/ensstat/indices/region_fldmean/TS/SET_VAR_TS_fldmean_REGION_land_ENSSTAT.nc


oldpwd=$(pwd)

for set in ${setlist}; do

   case ${set} in

      HIPPO_set_1420_LR_all)
           selyear=1420/1780
           ;;
                         
      *)
           selyear=all
           ;;

   esac  # set


   for var in ${varlist}; do
      for ts in ${tslist}; do
         for region in ${regionlist}; do

            tmp=$(mktemp -d ${TMP_BASEDIR}/tmp.ensstat.XXX)
            cd ${tmp}
            
            set +e; filelist=$(ls $(echo ${infile_struct} | sed -e "s&RUN&HIPPO_r\*i\*p\*f\*&g" -e "s&SET&${set}&g" -e "s&TS&${ts}&g" -e "s&VAR&${var}&g" -e "s&REGION&${region}&g") 2>/dev/null); set -e
         
            echo -e "ensemble mean will be computed from: \n ${filelist}"
         
            nfiles=$(echo ${filelist} | wc -w)
            echo "number of infiles: ${nfiles}"
   
            ensstat_fullname=$(eval echo ${outfile_struct} | sed -e "s&SET&${set}&g" -e "s&TS&${ts}&g" -e "s&VAR&${var}&g" -e "s&REGION&${region}&g")
            ensstat_dir=$(echo ${ensstat_fullname} | sed 's|\(.*\)/.*|\1|')
            ensstat_filename=$(basename ${ensstat_fullname})
   
            if [ ! -e ${ensstat_dir} ]; then mkdir -p ${ensstat_dir}; fi
         
            echo "outfile will be written to: ${ensstat_fullname}"
   
            for file in ${filelist}; do
               case ${selyear} in

                  all) 
                        ln -sf ${file} .
                        ;;

                  *) 
                        cdo selyear,${selyear} ${file} $(basename ${file})
                        ;;

               esac # selyear
            done # file


            cdo -O -s ensmin HIPPO_r*i*p*f*.nc ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensmin/g" )
            cdo -O -s ensmax HIPPO_r*i*p*f*.nc ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensmax/g" )
            cdo -O -s ensmean HIPPO_r*i*p*f*.nc ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensmean/g" )
            cdo -O -s ensstd HIPPO_r*i*p*f*.nc ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensstd/g" )
            cdo -O -s enspctl,12.5 HIPPO_r*i*p*f*.nc ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/enspctl125/g" )
            cdo -O -s enspctl,87.5 HIPPO_r*i*p*f*.nc ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/enspctl875/g" )
   
            ln -s ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensmin/g") ensmin.allyears.nc
            ln -s ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensmax/g") ensmax.allyears.nc
            ln -s ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensmean/g") ensmean.allyears.nc
            ln -s ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensstd/g") ensstd.allyears.nc
            ln -s ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/enspctl125/g") enspctl125.allyears.nc
            ln -s ${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/enspctl875/g") enspctl875.allyears.nc   

            cdo setrtoc,-9e33,9e33,1 ensmean.allyears.nc ones.nc
   
            case ${level} in
               surf) sellevel="";;
               *)    sellevel="-sellevel,${level}";;
            esac # level
   
            cdo -s mul -timmean ensmean.allyears.nc ones.nc ensmean.allyears.timmean.nc
   
            outfile=${ensstat_dir}/$(echo ${ensstat_filename} | sed -e "s/ENSSTAT/ensstat/g" -e "s/\.nc/.png/g")
   
            if ${plot}; then
         
               cat << __EOF__ | sed -e "s/^ *//" > plot.py
                   import xarray as xr
                   import numpy as np
                   import matplotlib.pyplot as plt
               
                   dsmean=xr.open_dataset('ensmean.allyears.nc')
                   dstimmean=xr.open_dataset('ensmean.allyears.timmean.nc')
                   dsmin=xr.open_dataset('ensmin.allyears.nc')
                   dsmax=xr.open_dataset('ensmax.allyears.nc')
                   dsstd=xr.open_dataset('ensstd.allyears.nc')
            
                   #-- create figure and axes instances
                   dpi = 100
                   fig = plt.figure(figsize=(1100/dpi, 1100/dpi), dpi=dpi)
                   ax  = fig.add_axes([0.1,0.1,0.8,0.9])
            
                   mean=dsmean.${var}.isel(lon=0,lat=0)
                   timmean=dstimmean.${var}.isel(lon=0,lat=0)
                   std=dsstd.${var}.isel(lon=0,lat=0)
                   time=dsmean.time
                   spreadmax=mean+std
                   spreadmin=mean-std
                   plt.plot(time,mean,linewidth=2,color='red')
                   plt.plot(time,timmean,linestyle='dashed',linewidth=2,color='black')
#                   plt.plot(time,spreadmin,linestyle='dashed',linewidth=2,color='red')
#                   plt.plot(time,spreadmax,linestyle='dashed',linewidth=2,color='red')
      
                   plt.fill_between(time.values,dsmin.${var}.isel(lon=0,lat=0),dsmax.${var}.isel(lon=0,lat=0),facecolor='gray')
                   plt.fill_between(time.values,spreadmin,spreadmax,facecolor='lightcoral')
            
                   plt.xlabel('time', fontsize=30)
                   plt.ylabel('${var}', fontsize=30)
             
                   # locs, labels=plt.xticks()
                   # print(locs)
             
                   ax.tick_params(axis='both', labelsize=20)
            
                   plt.savefig('${outfile}', bbox_inches='tight', dpi=dpi)
               
__EOF__
               
#               python3 plot.py
               echo "outfile written to ${outfile}"
   
            fi # plot?
         
            cd ${oldpwd}
            rm -rf ${tmp}

         done # region
      done # ts
   done # var
done # set

exit 0
