#!/bin/bash
#
############################################################################
#
# MODULE:       r.stream.watersheds
#
# AUTHOR(S):    Giuseppe Amatulli & Sami Domisch
#               based on "Domisch, S., Amatulli, G., Jetz, W. (2015)
#				Near-global freshwater-specific environmental variables for 
# 				biodiversity analyses in 1km resolution. Scientific Data 2:150073 doi:10.1038/sdata.2015.73"
#
# PURPOSE:      Calculation of contiguous stream-specific variables that account 
#				for the upstream environment (based on r.stream.watersheds).
#
# COPYRIGHT:    (C) 2001-2012 by the GRASS Development Team
#
#               This program is free software under the GNU General Public
#               License (>=v2). Read the file COPYING that comes with GRASS
#               for details.
#
#############################################################################

#%Module
#% description: Calculation of contiguous stream-specific variables that account for the upstream environment (based on r.stream.watersheds).
#% keywords: drainage, stream, hydrology
#%End

#%option
#% key: variable
#% type: string
#% key_desc: name
#% description: Name of raster to be converted into a stream-specific variable
#% required : yes
#%end

#%option
#% key: area
#% type: string
#% key_desc: string
#% multiple: no
#% options: watershed,stream
#% description: Area of aggregation: across the sub-watersheds or only across sub-streams
#% required : yes
#%end

#%option
#% key: folder 
#% type: string
#% key_desc: name
#% description: Provide the full folder path (same as for r.stream.watersheds)
#% required:no
#% answer: GISDBASE/folder_structure
#%end

#%option
#% key: out_folder 
#% type: string
#% key_desc: name
#% description: Provide the full folder path for the output stream-specific variable
#% required:no
#% answer: GISDBASE/stream-specific_variables
#%end

#%option
#% key: output
#% type: string
#% key_desc: method
#% multiple: yes
#% options: cells,min,max,range,mean,stddev,coeff_var,sum
#% description: Provide the output aggregation method for the stream-specific variable: upstream cells numbers, minimum, maximum, range, mean, standard deviation, coefficient of variation, sum. Output datatype is Int32
#% required:yes
#%end

#%option
#% key: scale
#% type: double
#% key_desc: value
#% description: Provide a scale factor to multiply or divide the final stream-specific variable (e.g if input raster values are between -1 and 1, use scale=1000 to inicrease the number of decimals - all outputs will be rounded to integers) 
#% answer: 1
#% required:no
#%end

#%option
#% key: cpu
#% type: double
#% description: Number of CPUs used for the parallel computation
#% answer: 1
#% required:no
#%end

if [ -z "$GISBASE" ] ; then
    echo "You must be in GRASS GIS to run this program." >&2
    exit 1
fi

if [ "$1" != "@ARGS_PARSED@" ] ; then
    exec g.parser "$0" "$@"
fi

#### check if we have awk
if [ ! -x "`which awk`" ] ; then
    g.message -e "awk required, please install awk or gawk first"
    exit 1
fi

# setting environment, so that awk works properly in all languages
unset LC_ALL
LC_NUMERIC=C
export LC_NUMERIC

#test:

if [ -z "$GIS_OPT_VARIABLE"  ] ; then
    g.message "Please provide the name of raster to be converted into a stream-specific variable."
exit 1
fi

if [ -z "$GIS_OPT_AREA" ] ; then
    g.message "Please provide area of aggregation: across the sub-watersheds or only across sub-streams."   
exit 1
fi

if [ -z "$GIS_OPT_OUTPUT"  ] ; then
    g.message "Please provide the output aggregation method for the stream-specific variable: upstream  min, max, range, mean, stddev, coeff_var, sum"
exit 1
fi

#check if drainage and stream file exists
eval `g.findfile element=cell file="$GIS_OPT_VARIABLE"`
if [ ! "$file" ] ; then
    g.message -e "<$GIS_OPT_VARIABLE> does not exist! Aborting."
    exit 1
fi


export GISDBASE=$(g.gisenv get=GISDBASE  separator=space)
export LOCATION_NAME=$(g.gisenv get=LOCATION_NAME  separator=space)
export GIS_OPT_AREA
export GIS_OPT_VARIABLE
export GIS_OPT_OUTPUT
export GIS_OPT_SCALE

if [ ${GIS_OPT_FOLDER} = "GISDBASE/folder_structure" ] ; then
       export GIS_OPT_FOLDER=$GISDBASE"/folder_structure"
else 
       export GIS_OPT_FOLDER=$GIS_OPT_FOLDER
fi

if [ ${GIS_OPT_OUT_FOLDER} = "GISDBASE/stream-specific_variables" ] ; then
      export GIS_OPT_OUT_FOLDER=$GISDBASE"/stream-specific_variables"
else 
      export GIS_OPT_OUT_FOLDER=$GIS_OPT_OUT_FOLDER
fi

mkdir  $GIS_OPT_OUT_FOLDER  2> /dev/null 

# what to do in case of user break:
exitprocedure()
{
echo ""
g.message -e 'Process aborted by user. All intermediate files have been deleted!'

# reset in the permanent mapset 
g.gisenv set="MAPSET=PERMANENT"
g.gisenv set="LOCATION_NAME=$LOCATION_NAME"
g.gisenv set="GISDBASE=$GISDBASE"

# delete intermediate files 
rm -fr  $GISDBASE/$LOCATION_NAME/sub_${GIS_OPT_AREA}* 
find    $GIS_OPT_FOLDER/ -maxdepth 1 -name  '*.txt' -delete 
find    $GIS_OPT_FOLDER/ -maxdepth 1 -name  '*.tif' -delete 
rm -fr  $GIS_OPT_FOLDER/blockfile/stat_*.txt 

exit 1
}

# shell check for user break (signal list: trap -l)
trap "exitprocedure" 2 3 9 15 19

echo ""
echo "##########################################################################"
echo "Stream-specific variable calculation based on the sub-watersheds and sub-streams"
echo ""
echo "Citation: Domisch, S., Amatulli, G., Jetz, W. (2015)"
echo "Near-global freshwater-specific environmental variables for"
echo "biodiversity analyses in 1km resolution. Scientific Data 2:150073 doi:10.1038/sdata.2015.73"
echo "##########################################################################"
echo ""

export GISRC_def=$GISRC

rm -fr   $GISDBASE/$LOCATION_NAME/sub_${GIS_OPT_AREA}*

echo Using  $( ls $GIS_OPT_FOLDER/blockfile/blockfile* | wc -l ) blocks in  $GIS_OPT_FOLDER/blockfile/

for file in  $GIS_OPT_FOLDER/*digit4/*digit3/*digit2/*digit1/blockfile*_sub_${GIS_OPT_AREA}.tar.gz    ; do 
    
    export DIRNAME=$(dirname   $file )
    filename=$(basename $file )
    cd $DIRNAME
    tar -xf $file 
    cd $GISDBASE

    export BLKID=$(echo $filename |  awk -v GIS_OPT_AREA=${GIS_OPT_AREA} '{ gsub("0000", "") ; gsub("blockfile","") ; gsub("_sub_","") ; gsub("GIS_OPT_AREA","") ; gsub(".tar.gz","") ;  print  }')
    export dir4d=${BLKID: -4:1} ; if [ -z $dir4d  ] ; then  dir4d=0 ; fi
    export dir3d=${BLKID: -3:1} ; if [ -z $dir3d  ] ; then  dir3d=0 ; fi  
    export dir2d=${BLKID: -2:1} ; if [ -z $dir2d  ] ; then  dir2d=0 ; fi
    export dir1d=${BLKID: -1:1} ; if [ -z $dir1d  ] ; then  dir1d=0 ; fi

g.gisenv set="GISDBASE=$GISDBASE"
g.gisenv set="LOCATION_NAME=$LOCATION_NAME"
g.gisenv set="MAPSET=PERMANENT"

echo  "Start the stream-specific variable aggregation for block $file"

ls  $DIRNAME/sub_${GIS_OPT_AREA}ID*.tif | awk '{ if (NR>0) { print $1 ,  NR} }'   | xargs  -n 2 -P $GIS_OPT_CPU  bash -c  $' 

file=$1
filename=$(basename $file .tif )
if [  $GIS_OPT_AREA =  "watershed"  ] ; then ID=${filename:15:99}; fi 
if [  $GIS_OPT_AREA =  "stream"  ] ; then ID=${filename:12:99}; fi 

NR=$2

if [ $( expr  $( expr $NR + 1 ) / 100)  -eq  $( expr $NR / 100 + 1 ) ] ; then 
echo  -en  "\r$(expr $NR / 100 + 1 )% done"
fi

# replicate the start-GISRC in a uniq GISRC

cp   $GISRC_def    $HOME/.grass7/rc$ID
export GISRC=$HOME/.grass7/rc$ID

g.mapset  -c   mapset=sub_${GIS_OPT_AREA}ID$ID   location=$LOCATION_NAME  dbase=$GISDBASE   --quiet

rm -f    $GISDBASE/$LOCATION_NAME/sub_${GIS_OPT_AREA}ID$ID/.gislock      

# export GISBASE=$( grep  gisbase   $(which grass78) | awk \'{ if(NR==2) { gsub ("\\"","" ) ; print $3 }  }\' )
export GISBASE=/usr/lib/grass78
export PATH=$PATH:$GISBASE/bin:$GISBASE/scripts
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:"$GISBASE/lib"
export GRASS_LD_LIBRARY_PATH="$LD_LIBRARY_PATH"
export PYTHONPATH="$GISBASE/etc/python:$PYTHONPATH"
export MANPATH=$MANPATH:$GISBASE/man

# echo  Load $DIRNAME/sub_${GIS_OPT_AREA}ID${ID}.tif 

r.in.gdal input=$DIRNAME/sub_${GIS_OPT_AREA}ID${ID}.tif       output=sub_${GIS_OPT_AREA}ID${ID}   --q 

g.region   rast=sub_${GIS_OPT_AREA}ID${ID}@sub_${GIS_OPT_AREA}ID${ID}  zoom=sub_${GIS_OPT_AREA}ID${ID}@sub_${GIS_OPT_AREA}ID${ID} --q 
r.mapcalc "sub_$GIS_OPT_VARIABLE = if ( sub_${GIS_OPT_AREA}ID${ID} == 1 , $GIS_OPT_VARIABLE@PERMANENT  , null())"    --o --q

echo $ID"|"$( r.univar -t --q  map=sub_$GIS_OPT_VARIABLE  |  awk  \'{ if (NR==2 ) print}\'   ) >   $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID$ID.txt

FULL=$(awk -F "|"  \'{if (NF==13) {print 1 } else {print 0} }\'   $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID$ID.txt)

if [ $FULL -eq 0 ] ; then 
rm   $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID$ID.txt   
else
rm -f $DIRNAME/sub_${GIS_OPT_AREA}ID${ID}.tif      
fi

rm -r $HOME/.grass7/rc$ID    $GISDBASE/$LOCATION_NAME/sub_${GIS_OPT_AREA}ID$ID

' _ 

g.gisenv set="MAPSET=PERMANENT"
g.gisenv set="LOCATION_NAME=$LOCATION_NAME"
g.gisenv set="GISDBASE=$GISDBASE"

echo -en  "\r100% done"
echo ""

echo "Check for missing files due to RAM overload"

if ls $DIRNAME/sub_${GIS_OPT_AREA}ID*.tif  1> /dev/null 2>&1 ; then

echo $(ls $DIRNAME/sub_${GIS_OPT_AREA}ID*.tif  | wc -l ) files had RAM issues

ls $DIRNAME/sub_${GIS_OPT_AREA}ID*.tif  | awk '{ if (NR>0) { print $1 ,  NR} }'   | xargs  -n 2 -P 1 bash -c  $' 

file=$1
filename=$(basename $file .tif )

echo Fix missing file $file using only 1 CPU

if [  $GIS_OPT_AREA =  "watershed"  ] ; then ID=${filename:15:99}; fi 
if [  $GIS_OPT_AREA =  "stream"  ] ; then ID=${filename:12:99}; fi 

NR=$2

if [ ! -f  $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID$ID.txt ] ; then 

if [ $( expr  $( expr $NR + 1 ) / 100)  -eq  $( expr $NR / 100 + 1 ) ] ; then 
echo  -en  "\r$(expr $NR / 100 + 1 )% done"
fi

# replicate the start-GISRC in a uniq GISRC

cp   $GISRC_def    $HOME/.grass7/rc$ID
export GISRC=$HOME/.grass7/rc$ID

g.mapset  -c   mapset=sub_${GIS_OPT_AREA}ID$ID   location=$LOCATION_NAME  dbase=$GISDBASE   --quiet

rm -f    $GISDBASE/$LOCATION_NAME/sub_${GIS_OPT_AREA}ID$ID/.gislock      

# export GISBASE=$( grep  gisbase   $(which grass78) | awk \'{ if(NR==2) { gsub ("\\"","" ) ; print $3 }  }\' )
export GISBASE=/usr/lib/grass78
export PATH=$PATH:$GISBASE/bin:$GISBASE/scripts
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:"$GISBASE/lib"
export GRASS_LD_LIBRARY_PATH="$LD_LIBRARY_PATH"
export PYTHONPATH="$GISBASE/etc/python:$PYTHONPATH"
export MANPATH=$MANPATH:$GISBASE/man

# echo  Load $DIRNAME/sub_${GIS_OPT_AREA}ID${ID}.tif 

r.in.gdal input=$DIRNAME/sub_${GIS_OPT_AREA}ID${ID}.tif       output=sub_${GIS_OPT_AREA}ID${ID}   --q 
rm -f $DIRNAME/sub_${GIS_OPT_AREA}ID${ID}.tif 

g.region   rast=sub_${GIS_OPT_AREA}ID${ID}@sub_${GIS_OPT_AREA}ID${ID}  zoom=sub_${GIS_OPT_AREA}ID${ID}@sub_${GIS_OPT_AREA}ID${ID} --q 
r.mapcalc "sub_$GIS_OPT_VARIABLE = if ( sub_${GIS_OPT_AREA}ID${ID} == 1 , $GIS_OPT_VARIABLE@PERMANENT  , null())"    --o --q

echo $ID"|"$( r.univar -t --q  map=sub_$GIS_OPT_VARIABLE  |  awk  \'{ if (NR==2 ) print}\'   ) >   $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID$ID.txt

FULL=$(awk -F "|"  \'{if (NF==13) {print 1 } else {print 0} }\'   $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID$ID.txt)

if [ $FULL -eq 0 ] ; then rm   $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID$ID.txt ; fi

rm -r $HOME/.grass7/rc$ID    $GISDBASE/$LOCATION_NAME/sub_${GIS_OPT_AREA}ID$ID

fi 

' _ 

else 
echo 0 files had RAM issues
fi 


cat $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID*.txt >  $DIRNAME/stat_${GIS_OPT_VARIABLE}.txt 
rm -f  $DIRNAME/stat_${GIS_OPT_VARIABLE}_ID*.txt 

done 

echo "Aggregating the final table to reclassify the raster ID"

echo "ID|non_null_cells|null_cells|min|max|range|mean|mean_of_abs|stddev|variance|coeff_var|sum|sum_abs" >  $GIS_OPT_FOLDER/blockfile/stat_${GIS_OPT_VARIABLE}.txt
cat  $GIS_OPT_FOLDER/*digit4/*digit3/*digit2/*digit1/stat_${GIS_OPT_VARIABLE}.txt  >>  $GIS_OPT_FOLDER/blockfile/stat_${GIS_OPT_VARIABLE}.txt
rm -f  $GIS_OPT_FOLDER/*digit4/*digit3/*digit2/*digit1/stat_${GIS_OPT_VARIABLE}.txt

echo Reclass the grid_id for the following output  ${GIS_OPT_OUTPUT//","/" "}

echo ${GIS_OPT_OUTPUT//","/" "} |  xargs -n 1 -P $GIS_OPT_CPU bash -c $'

if [ $1 = "cells" ]     ; then col=2 ; fi
if [ $1 = "min" ]       ; then col=4 ; fi
if [ $1 = "max" ]       ; then col=5 ; fi
if [ $1 = "range" ]     ; then col=6 ; fi
if [ $1 = "mean" ]      ; then col=7 ; fi 
if [ $1 = "stddev" ]    ; then col=9 ; fi 
if [ $1 = "coeff_var" ] ; then col=11 ; fi 
if [ $1 = "sum" ]       ; then col=12 ; fi 

awk -F "|" -v col=$col -v scale=$GIS_OPT_SCALE  \' { if (NR>1 ) {gsub("-nan","0",$col); gsub("inf","0",$col);  print $1" = "int($col * scale) }}\' $GIS_OPT_FOLDER/blockfile/stat_${GIS_OPT_VARIABLE}.txt  > $GIS_OPT_FOLDER/blockfile/stat_${GIS_OPT_VARIABLE}_$1.txt

' _ 

echo ${GIS_OPT_OUTPUT//","/" "} |  xargs -n 1 -P 1 bash -c $'

if [ $1 = "cells" ]     ; then col=2 ; fi
if [ $1 = "min" ]       ; then col=4 ; fi 
if [ $1 = "max" ]       ; then col=5 ; fi 
if [ $1 = "range" ]     ; then col=6 ; fi 
if [ $1 = "mean" ]      ; then col=7 ; fi 
if [ $1 = "stddev" ]    ; then col=9 ; fi 
if [ $1 = "coeff_var" ] ; then col=11 ; fi 
if [ $1 = "sum"  ]      ; then col=12 ; fi 

echo  Create stream-specific variable:  $GIS_OPT_VARIABLE $1

r.reclass input=grid_ID  output=${GIS_OPT_VARIABLE}_$1  rules=$GIS_OPT_FOLDER/blockfile/stat_${GIS_OPT_VARIABLE}_${1}.txt   --overwrite   --q
# rm -f $GIS_OPT_FOLDER/blockfile/stat_${GIS_OPT_VARIABLE}_${1}.txt

r.mapcalc "  ${GIS_OPT_VARIABLE}_$1  = ${GIS_OPT_VARIABLE}_$1 "  --o  

r.out.gdal -c  input=${GIS_OPT_VARIABLE}_${1}  nodata=-9999  output=$GIS_OPT_OUT_FOLDER/${GIS_OPT_VARIABLE}_${1}.tif  type=Int32  -c  createopt="COMPRESS=LZW,ZLEVEL=9"  --o --q  2> /dev/null
if test -f "$GIS_OPT_FOLDER/blockfile/${GIS_OPT_VARIABLE}_${1}.tif"; then echo "$GIS_OPT_FOLDER/blockfile/${GIS_OPT_VARIABLE}_${1}.tif has been created!";fi
#echo "$GIS_OPT_OUT_FOLDER/blockfile/${GIS_OPT_VARIABLE}_${1}.tif has been created!"

' _ 

exit 

echo "The full stream variable calculation has been done!"
echo "To list the stream-specific output variables: ls $GIS_OPT_OUT_FOLDER/*.tif"