#!/bin/bash

# Usage documentation line
USAGE1=" Usage: "
USAGE2=" ./regrid.sh --shfile=test_sh.grb --ggfile=test_gg.grb --regfile=all_regular.grb --plevfile=all_plevels.grb --ncout=0 --grbedition=2 " 
USAGE="$USAGE1 $USAGE2"

# Default configuratons
VERBOSE=0
TMPDIR=../tmp/
GRBED=2
NCOUT=1

# Loop over all arguments
while [ $# -gt 0 ]; do
  case "$1" in
    -h|--help)
      echo $USAGE
      exit 1
      ;;
      
    --shfile=*) ## Input file
      SHFILE="${1#*=}"
      ;;
    
    --ggfile=*) ## Input file
      GGFILE="${1#*=}"
      ;;  
      
    -v|--verbose)
      VERBOSE=$(( $VERBOSE + 1 ))  # Each -v adds 1 to verbosity.
      ;;
      
    --regfile=*) # Takes an option argument; ensure it has been specified.
      REGOUT="${1#*=}"
      ;;
    
    --plevfile=*) # Takes an option argument; ensure it has been specified.
      PLEVOUT="${1#*=}"
      ;;
    
    --grbedition=*) # Which grib edition to use. "2" is best, and default. 
      GRBED="${1#*=}"
      ;;
    
    --ncout=*) # Output in netCDF format
      NCOUT="${1#*=}"
      ;;
      
    --tmp=*)
      TMPDIR="${1#*=}"
      ;;
    
    *)
      printf "***************************\n"
      printf "* Error: Invalid argument.*\n"
      printf "***************************\n"
      echo $USAGE
      exit 1
  esac
  shift
done

# Make tmp directory
mkdir -p $TMPDIR

# This script needs grib_copy and cdo
# Check if they exist
hash grib_set 2>/dev/null || { echo >&2 "I require grib_set but it's not installed or in the PATH.  Aborting."; exit 1; }
hash grib_copy 2>/dev/null || { echo >&2 "I require grib_copy but it's not installed or in the PATH.  Aborting."; exit 1; }
hash cdo 2>/dev/null || { echo >&2 "I require cdo but it's not installed or in the PATH.  Aborting."; exit 1; }

# An array with all files that we will glue together in the end
files=()

for FILE in $GGFILE $SHFILE; do 
   
   # Remove lnsp in isobaric coordinates. 
   # For some (strange) reason, grib_set converts it to hybrid coordinates but keeps level index,
   # so that lnsp at 1000 hPa becomes lnsp at hybrid level 100000, which does not exist! 
   #
   # Thus, we first make a new file where we exclude this variable. 
   if [ $VERBOSE -ge 1 ]; then  
      echo " Remove isobaric LNSP data from output. "
   fi
   grib_copy -w shortName!=lnsp -w typeOfLevel!=isobaricInhPa $FILE $TMPDIR/tmpfile_clean.grb
   
   # Set the file to use grib1 or grib2. 
   # CDO does not support mixed types. 
   if [ "$GRBED" == "1" ]; then
      if [ $VERBOSE -ge 1 ]; then  
         echo " Convert all GRIB2 data to GRIB1 (not lossless conversion!) "
      fi
      grib_set -w editionNumber=2 -s editionNumber=1 $TMPDIR/tmpfile_clean.grb $TMPDIR/tmpfile.grb
   fi 
   if [ "$GRBED" == "2" ]; then
      if [ $VERBOSE -ge 1 ]; then  
         echo " Convert all GRIB1 to GRIB2 (should be lossless) "
      fi
      grib_set -w editionNumber=1 -s editionNumber=2 $TMPDIR/tmpfile_clean.grb $TMPDIR/tmpfile.grb
   fi 
   rm -f tmpfile_clean.grb
   
   # Convert the different grid types to regular grid, and merge to one file 
   #
   # We support horizontal grids reduced_gg, regular_gg and sh
   # Anything else will not be included in the output. 
   #
   # Split into files for each horizontal grid type
   
   if [ $FILE == $SHFILE ]; then
      # In the SH file, all variables are spectral 
      if [ $VERBOSE -ge 1 ]; then  
         echo " Convert spectral data to regular grid using sp2gpl "
      fi 
      cdo -O sp2gpl $TMPDIR/tmpfile.grb $TMPDIR/tmpfile_shreg.grb
      rm -f $TMPDIR/${INFILE}_sh.grb1
      if [ -e "$TMPDIR/tmpfile_shreg.grb" ]; then
         files+=("$TMPDIR/tmpfile_shreg.grb")
      else
         echo " That is weird. No file with SH->GG data exists! "
         echo " I stop here. Something probably went wrong "
         exit 1
      fi 
   fi
   
   if [ $FILE == $GGFILE ]; then
      # In the GG files, there can be reduced_gg or regular_gg
      if [ $VERBOSE -ge 1 ]; then  
         echo " Split into one file for regular and one for reduced Gaussian data "
      fi 
      grib_copy -w gridType=reduced_gg $TMPDIR/tmpfile.grb $TMPDIR/tmpfile_reduced_gg.grb 
      grib_copy -w gridType=regular_gg $TMPDIR/tmpfile.grb $TMPDIR/tmpfile_regular_gg.grb
      
      # Convert from reduced_gg to regular_gg
      if [ -e "$TMPDIR/tmpfile_reduced_gg.grb" ]; then
         if [ $VERBOSE -ge 1 ]; then  
            echo " Regrid reduced Gaussian data to regular Gaussian "
         fi 
         cdo -O -setgridtype,regular $TMPDIR/tmpfile_reduced_gg.grb $TMPDIR/tmpfile_ggreg.grb
         files+=("$TMPDIR/tmpfile_ggreg.grb")
      else
         echo " WARNING: No reduced_gg data exists in this file: "$FILE
      fi 
      
      if [ -e "$TMPDIR/tmpfile_regular_gg.grb" ]; then
         files+=("$TMPDIR/tmpfile_regular_gg.grb")
      else
         echo " No regular_gg data exists in the original file " $FILE
      fi 
   fi 

done

# Write all file names to string
string=$( printf " %s " ${files[@]} )

# and merge to one grib file 
if [ $VERBOSE -ge 1 ]; then  
   echo " Merge all horizontal grids to one file "
fi 
cdo -O merge $string $TMPDIR/tmpfile.grb
rm -f $string

# Convert to netCDF
if [ $NCOUT -gt 0 ]; then
   if [ $VERBOSE -ge 1 ]; then  
      echo " Convert output to netCDF "
   fi 
   cdo -O -f nc -t ecmwf copy $TMPDIR/tmpfile.grb $TMPDIR/tmpfile.nc
   rm -f $TMPDIR/tmpfile.grb
fi 

# Split file into hybrid coordinates and not
if [ $NCOUT -gt 0 ]; then
   cdo selzaxis,hybrid $TMPDIR/tmpfile.nc $TMPDIR/tmpfile_hybrid.nc
else
   cdo selzaxis,hybrid $TMPDIR/tmpfile.grb $TMPDIR/tmpfile_hybrid.grb
fi 

# Convert hybrid data to pressure level data
plevels=100000,92500,85000,75000,60000,50000,40000,30000,20000,10000,5000,2000,1000
if [ $NCOUT -gt 0 ]; then
   cdo -O ml2plx,${plevels} $TMPDIR/tmpfile_hybrid.nc $TMPDIR/tmpfile_plevels.nc
   rm -f $TMPDIR/tmpfile_hybrid.nc 
else
   cdo -O ml2plx,${plevels} $TMPDIR/tmpfile_hybrid.grb $TMPDIR/tmpfile_plevels.grb
   rm -f $TMPDIR/tmpfile_hybrid.grb
fi 
   
# Make new file names 
if [ $NCOUT -gt 0 ]; then
   mv $TMPDIR/tmpfile.nc $REGOUT
   mv $TMPDIR/tmpfile_plevels.nc $PLEVOUT
else
   mv $TMPDIR/tmpfile.grb $REGOUT
   mv $TMPDIR/tmpfile_plevels.grb $PLEVOUT
fi