#! /bin/bash
#SBATCH --job-name=proc_fv3    # Specify job name
#SBATCH --partition=shared     # Specify partition name
#SBATCH --ntasks=5             # Specify max. number of tasks to be invoked
#SBATCH --mem-per-cpu=5120     # Specify real memory required per CPU in MegaBytes
#SBATCH --time=08:00:00        # Set a limit on the total run time
#SBATCH --mail-type=FAIL       # Notify user by email in case of job failure
#SBATCH --account=bb1072       # Charge resources on this project account
#SBATCH --output=my_job.o%j    # File name for standard output
#SBATCH --error=my_job.e%j     # File name for standard error output

# Author: Daniel Klocke
# Date: Aug 2018
# Purpose: process FV3 2D DYAMOND output to obtain time averaged fields on 0.1x0.1deg lat lon grid
#
echo "Executing proc_FV3.sh" `date`
source /sw/rhel6-x64/etc/profile.mistral
module swap cdo cdo/1.9.5-magicsxx-gcc64

## choose what you want to do. Files need to be remapped first. Remapping takes long.
do_weights='no'
do_remap='yes'
do_timeave='no'
do_zonmean='no'
do_ztmean='no'

## choose variables to processe.
#
## List of FV3 2D fields:
#                         cape, flds, cin, flus, flut, fsds, fsdt, fsus
#                         fsut, h500, intqg, intqi, lhflx, u10m, pr, intql,
#                         u200, ps, intqr, intqs, ustrs, intqv, v10m, rh500,
#                         v200, rh700, rh850, shflx, t2m, vstrs, ts
declare -a var_list=("ps") #"cape" "flds" "cin" "flus" "flut" "fsds" "fsdt" "fsus" "fsut" "intqg" "intqi" "lhflx" "u10m" "pr" "intql" "intqr" "intqs" "intqv" "v10m" "shflx" "t2m" "ustrs" "vstrs" "ts")

## Directories
tmp_dir='/scratch/m/m218027'
out_dir='/work/ka1081/Hackathon/GrossStats'
fv3_dir='/work/ka1081/DYAMOND/FV3-3.25km'

## Creat grid for original FV3 grid to latlon
if [ $do_weights == 'yes' ]
then
    echo 'creating remapping weights from FV3 cubed spere to 0.1x01.deg lat lon grid'
    for N in 1 2 3 4 5 6; do 
      cdo import_fv3grid /work/ka1081/DYAMOND/FV3-3.25km/2016082100/grid_spec.tile$N.nc ${tmp_dir}/FV3-3.25km/gridspec.tile$N.nc
    done
    cdo collgrid ${tmp_dir}/FV3-3.25km/gridspec.tile?.nc ${tmp_dir}/FV3-3.25km/gridspec.nc
    cdo -P 12 genycon,${out_dir}/0.10_grid.nc -setgrid,${tmp_dir}/FV3-3.25km/gridspec.nc -collgrid,gridtype=unstructured ${fv3_dir}/2016080100/lhflx_15min.tile?.nc  ${tmp_dir}/FV3-3.25km/fv3_weights_to_0.10deg.nc
fi


## remapping
if [ $do_remap == 'yes' ]
then
    echo 'remapping FV3 fields to 0.1x01.deg lat lon grid'
    cd ${tmp_dir}
    for var in "${var_list[@]}"; do
    echo 'remapping ' ${var}
    cdo -P 12 remap,${out_dir}/0.10_grid.nc,${tmp_dir}/FV3-3.25km/fv3_weights_to_0.10deg.nc -setgrid,${tmp_dir}/FV3-3.25km/gridspec.nc -collgrid,gridtype=unstructured ${fv3_dir}/2016080100/${var}_15min.tile?.nc ${tmp_dir}/FV3_${var}_1.nc
    cdo -P 12 remap,${out_dir}/0.10_grid.nc,${tmp_dir}/FV3-3.25km/fv3_weights_to_0.10deg.nc -setgrid,${tmp_dir}/FV3-3.25km/gridspec.nc -collgrid,gridtype=unstructured ${fv3_dir}/2016081100/${var}_15min.tile?.nc ${tmp_dir}/FV3_${var}_2.nc
    cdo -P 12 remap,${out_dir}/0.10_grid.nc,${tmp_dir}/FV3-3.25km/fv3_weights_to_0.10deg.nc -setgrid,${tmp_dir}/FV3-3.25km/gridspec.nc -collgrid,gridtype=unstructured ${fv3_dir}/2016082100/${var}_15min.tile?.nc ${tmp_dir}/FV3_${var}_3.nc
    cdo -P 12 remap,${out_dir}/0.10_grid.nc,${tmp_dir}/FV3-3.25km/fv3_weights_to_0.10deg.nc -setgrid,${tmp_dir}/FV3-3.25km/gridspec.nc -collgrid,gridtype=unstructured ${fv3_dir}/2016083100/${var}_15min.tile?.nc ${tmp_dir}/FV3_${var}_4.nc
    cdo cat ${tmp_dir}/FV3_${var}_?.nc ${out_dir}/AllTimeStepsOn0.1degGrid/fv3_${var}_0.10deg.nc
done
fi

## time means
if [ $do_timeave == 'yes' ]
then
    echo 'time averaging remapped FV3 fields'
    cd ${out_dir}/AllTimeStepsOn0.1degGrid

    for var in "${var_list[@]}"; do
      if [ ! -f fv3_${var}_0.10deg.nc ]; then
        echo 'file fv3_'"${var}"'_0.10deg.nc does not exist' 
        echo 'remap first!'
        exit
      fi
      cdo timmean fv3_${var}_0.10deg.nc ../40DayTimeAveragesOn0.1degGrid/fv3_0d_40d_${var}_0.10deg.nc
      cdo -timmean -seltimestep,1/960 fv3_${var}_0.10deg.nc ../10DayTimeAveragesOn0.1degGrid/fv3_0d_10d_${var}_0.10deg.nc
      cdo -timmean -seltimestep,961/1920 fv3_${var}_0.10deg.nc ../10DayTimeAveragesOn0.1degGrid/fv3_10d_20d_${var}_0.10deg.nc
      cdo -timmean -seltimestep,1921/2880 fv3_${var}_0.10deg.nc ../10DayTimeAveragesOn0.1degGrid/fv3_20d_30d_${var}_0.10deg.nc
      cdo -timmean -seltimestep,2881/3840 fv3_${var}_0.10deg.nc ../10DayTimeAveragesOn0.1degGrid/fv3_30d_40d_${var}_0.10deg.nc
    done
fi

## zonal means
if [ $do_zonmean == 'yes' ]
then
    echo 'zonally averaging remapped FV3 fields'
    cd ${out_dir}/AllTimeStepsOn0.1degGrid

    for var in "${var_list[@]}"; do
      if [ ! -f fv3_${var}_0.10deg.nc ]; then
        echo 'file fv3_'"${var}"'_0.10deg.nc does not exist' 
        echo 'remap first!'
        exit
      fi
      echo 'processing ' ${var}
      cdo -f nc4 -k auto zonmean fv3_${var}_0.10deg.nc ../ZonMeanOn0.1degGrid/fv3_zonmean_${var}_0.10deg.nc
    done
fi

## zonal time means
if [ $do_ztmean == 'yes' ]
then
    echo 'time averaging zonal averaged remapped FV3 fields'
    cd ${out_dir}/ZonMeanOn0.1degGrid

    for var in "${var_list[@]}"; do
      if [ ! -f fv3_zonmean_${var}_0.10deg.nc ]; then
        echo 'file fv3_zonmean_'"${var}"'_0.10deg.nc does not exist' 
        echo 'remap and zonally average first!'
        exit
      fi
      echo 'processing ' ${var}
      cdo timmean fv3_zonmean_${var}_0.10deg.nc fv3_zonmean_0d_40d_${var}_0.10deg.nc
      cdo timmean -seltimestep,1/960 fv3_zonmean_${var}_0.10deg.nc fv3_zonmean_0d_10d_${var}_0.10deg.nc
      cdo timmean -seltimestep,961/1920 fv3_zonmean_${var}_0.10deg.nc fv3_zonmean_10d_20d_${var}_0.10deg.nc
      cdo timmean -seltimestep,1921/2880 fv3_zonmean_${var}_0.10deg.nc fv3_zonmean_20d_30d_${var}_0.10deg.nc
      cdo timmean -seltimestep,2881/3840 fv3_zonmean_${var}_0.10deg.nc fv3_zonmean_30d_40d_${var}_0.10deg.nc
    done
fi

exit