#!/bin/sh

############################################################################
#
# MODULE:	v.rast.stats
# AUTHOR(S):	Markus Neteler, Otto Dassau
# PURPOSE:	Calculates univariate statistics from a GRASS raster map
#		only for areas covered by vector objects on a per-category base
# COPYRIGHT:	(C) 2005, 2011 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: Calculates univariate statistics from a raster map based on vector polygons and uploads statistics to new attribute columns.
#%  keywords: vector, raster, statistics
#%End
#%flag
#%  key: c
#%  description: Continue if upload column(s) already exist
#%END
#%flag
#%  key: e
#%  description: Calculate extended statistics
#%END
#%option
#% key: vector
#% type: string
#% key_desc: name
#% gisprompt: old,vector,vector
#% description: Name of vector polygon map
#% required : yes
#%End
#%option
#% key: layer
#% type: integer
#% gisprompt: old_layer,layer,layer
#% description: A single vector map can be connected to multiple database tables. This number determines which table to use
#% answer: 1
#% required : no
#%end
#%option
#% key: raster
#% type: string
#% key_desc: name
#% gisprompt: old,cell,raster
#% description: Name of raster map to calculate statistics from
#% required : yes
#%END
#%option
#% key: colprefix
#% type: string
#% description: Column prefix for new attribute columns
#% required : yes
#%end
#%option
#% key: percentile
#% type: integer
#% description: Percentile to calculate (requires extended statistics flag)
#% options: 0-100
#% answer: 90
#% 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

PROG=`basename "$0"`

#### check if we have awk
if [ ! -x "`which awk`" ] ; then
    g.message "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

#### setup temporary file
TEMPFILE="`g.tempfile pid=$$`"
if [ $? -ne 0 ] || [ -z "$TEMPFILE" ] ; then
    g.message -e "Unable to create temporary files" 
    exit 1
fi

### variables for temp files
SQLTMP="$TEMPFILE.sql"
STATSTMP="$TEMPFILE.csv"
COLNAMETMP="$TEMPFILE.col"

# we need a random name
TMPNAME=`basename "$TEMPFILE"`

cleanup()
{
   # restore settings:
   g.region region="$TMPNAME"
   g.remove region="$TMPNAME" --quiet
   g.remove rast="${VECTOR}_$TMPNAME" --quiet
   g.remove MASK --quiet 2>/dev/null
   if [ $MASKFOUND -eq 1 ] ; then
      g.message "Restoring previous MASK..."
      g.rename "${TMPNAME}_origmask",MASK --quiet
   fi
   rm -f "$TEMPFILE" "$TMPNAME" "$TEMPFILE.cats" "$SQLTMP"
}

# what to do in case of user break:
exitprocedure()
{
   g.message -e 'User break!'
   cleanup
   exit 1
}
# shell check for user break (signal list: trap -l)
trap "exitprocedure" 2 3 15

RASTER="$GIS_OPT_RASTER"
COLPREFIX="$GIS_OPT_COLPREFIX"

### setup enviro vars ###
eval `g.gisenv`
: ${GISBASE?} ${GISDBASE?} ${LOCATION_NAME?} ${MAPSET?}

# does map exist in CURRENT mapset?
eval `g.findfile element=vector file="$GIS_OPT_VECTOR" mapset="$MAPSET"`
VECT_MAPSET=`echo "$GIS_OPT_VECTOR" | grep '@' | cut -f2 -d'@'`
if [ -z "$VECT_MAPSET" ] ; then
   VECT_MAPSET="$MAPSET"
fi
if [ ! "$file" ] || [ "$VECT_MAPSET" != "$MAPSET" ] ; then
   g.message -e "Vector map <$GIS_OPT_VECTOR> not found in current mapset"
   exit 1
else
   VECTOR=`echo "$GIS_OPT_VECTOR" | cut -f1 -d'@'`
   VECTORFULL="${VECTOR}@${VECT_MAPSET}"
fi

# check the input vector map
eval `v.info -t map="$VECTORFULL" layer="$GIS_OPT_LAYER"`
if [ $areas -eq 0 ] ; then
   g.message -e "No areas found in vector map"
   exit 1
fi

# check the input raster map
eval `g.findfile element=cell file="$RASTER"`
if [ ! "$file" ] ; then
   g.message -e "Raster map <$RASTER> not found"
   exit 1
fi

# check presence of raster MASK, put it aside
MASKFOUND=0
eval `g.findfile element=cell file=MASK`
if [ "$file" ] ; then
   g.message "Raster MASK found, temporarily disabled"
   g.rename MASK,"${TMPNAME}_origmask" --quiet
   MASKFOUND=1
fi

# get RASTER resolution of map which we want to query:
# fetch separated to permit for non-square cells (latlong etc)
NSRES=`r.info -s "$RASTER" | grep nsres | cut -d'=' -f2`
if [ $? -ne 0 ] ; then
   g.message -e "An error occurred reading the input raster map resolution."
   cleanup
   exit 1
fi
EWRES=`r.info -s "$RASTER" | grep ewres | cut -d'=' -f2`

# save current settings:
g.region save="$TMPNAME" --quiet

# temporarily aligning current region to raster <$RASTER>
g.region align="$RASTER"

# create raster from vector map for r.univar 
v.to.rast in="$VECTORFULL" out="${VECTOR}_$TMPNAME" layer="$GIS_OPT_LAYER" use=cat type=area --quiet
if [ $? -ne 0 ] ; then
   g.message -e "An error occurred while converting vector to raster"
   cleanup
   exit 1
fi

# dump cats to file to avoid "too many argument" problem:
r.category "${VECTOR}_$TMPNAME" fs=';' --quiet | cut -d';' -f1 > "$TEMPFILE.cats"
# echo "List of categories found: $CATSLIST"
NUMBER=`cat "$TEMPFILE.cats" | wc -l | awk '{print $1}'`
if [ $NUMBER -lt 1 ] ; then
   g.message -e "No categories found in raster map"
   cleanup
   exit 1
fi

# check if DBF driver used, in this case cut to 10 chars col names:
DBFDRIVER=0
v.db.connect -gl map="$VECTOR" fs="|" layer="$GIS_OPT_LAYER" | \
   cut -d'|' -f5 | grep -i 'dbf' --quiet
if [ $? -eq 0 ] ; then
   DBFDRIVER=1
else
   # in case a table is missing, we'll trap a crash later...
   DBFDRIVER=0
fi

if [ $DBFDRIVER -eq 1 ] ; then
   # `wc -c` reports number of chars + 1 for the newline
   if [ `echo "$COLPREFIX" | wc -c` -gt 10 ] ; then
       g.message -e "Cannot create unique names for columns. \
         Either use a shorter column prefix or switch to another DB \
	 backend such as SQLite. DBF limits the length to 10 characters"
       cleanup
       exit 1
   fi
fi

# we need this for non-DBF driver:
DB_SQLDRIVER=`v.db.connect -gl map="$VECTOR" fs="|" layer="$GIS_OPT_LAYER" | cut -d'|' -f5`
g.message -d message="sql driver: [$DB_SQLDRIVER]"

DB_DATABASE="`v.db.connect -gl map=$VECTOR fs='|' layer=$GIS_OPT_LAYER | cut -d'|' -f4`"
g.message -d message="database: [$DB_DATABASE]"

# find out which table is linked to the vector map on the given layer
TABLE=`v.db.connect map="$VECTOR" -gl fs='|' layer="$GIS_OPT_LAYER" | cut -d'|' -f2`
if [ -z "$TABLE" ] ; then
   g.message -e 'There is no table connected to this map! Run v.db.connect or v.db.addtable first.'
   exit 1
fi
g.message -d message="table: [$TABLE]"

# get key column
KEYCOL=`v.db.connect map="$VECTOR" -gl fs='|' layer="$GIS_OPT_LAYER" | cut -d'|' -f3`

BASECOLS="n min max range mean stddev variance cf_var sum"

# do extended stats?
if [ $GIS_FLAG_E -eq 1 ] ; then
   # namespace is limited in DBF but the % value is important
   if [ "$DBFDRIVER" -eq 1 ] ; then
      PERCCOL="per$GIS_OPT_PERCENTILE"
   else
      PERCCOL="percentile_$GIS_OPT_PERCENTILE"
   fi
   EXTRACOLS="first_quartile median third_quartile $PERCCOL"
else
   EXTRACOLS=""
fi

unset ADDCOLS
unset COLNAMES
for i in $BASECOLS $EXTRACOLS ; do
  # check if column already present
  if [ $DBFDRIVER -eq 1 ] ; then
     CURRCOLUMN="`echo "${COLPREFIX}_$i" | cut -b1-10`"
  else
     CURRCOLUMN="${COLPREFIX}_$i"
  fi
  if [ -n "$COLNAMES" ] ; then
     COLNAMES="${COLNAMES},$CURRCOLUMN"
  else
     COLNAMES="$CURRCOLUMN"
  fi
  g.message -d message="Looking for <$CURRCOLUMN> in <$VECTOR> ..."
  v.info -c "$VECTORFULL" layer="$GIS_OPT_LAYER" --quiet | sed -e 's+^+|+' -e 's+$+|+' | \
    grep "|$CURRCOLUMN|" --quiet
  if [ $? -eq 0 ] ; then
    if [ "$GIS_FLAG_C" -ne 1 ] ; then
       g.message -e "Cannot create column <$CURRCOLUMN> (already present). \
         Use -c flag to update values in this column."
       cleanup
       exit 1
    fi
  else
    if [ -n "$ADDCOLS" ] ; then
       ADDCOLS="$ADDCOLS,"
    fi
    if [ "$i" = "n" ] ; then
       COLTYPE="INTEGER"
    else
       COLTYPE="DOUBLE PRECISION"
    fi
    ADDCOLS="${ADDCOLS}$CURRCOLUMN $COLTYPE"
  fi
done

if [ -n "$ADDCOLS" ] ; then
   g.message -v "Adding columns <$ADDCOLS>"
#   v.db.addcol map="$VECTORFULL" columns="$ADDCOLS" layer="$GIS_OPT_LAYER"
#   if [ $? -ne 0 ] ; then
#      g.message -e "Cannot continue (problem adding columns)."
#      cleanup
#      exit 1
#   fi

  # code borrowed from v.db.addcol
  colnum=`echo "$ADDCOLS" | awk -F, '{print NF}'`

  n=1
  while [ "$n" -le "$colnum" ]
  do
    col=`echo "$ADDCOLS" | cut -d',' -f$n`
    
    echo "ALTER TABLE $TABLE ADD COLUMN $col" | db.execute database="${DB_DATABASE}" driver="${DB_SQLDRIVER}"
    if [ $? -eq 1 ] ; then
      g.message -e "Cannot continue (problem adding column)."
      exit 1
    fi
    n=`expr $n + 1`
  done
fi

# calculate statistics for zones:
g.message -v "Processing data ..."

# get rid of any earlier attempts
rm -f "$SQLTMP" "$STATSTMP" "$COLNAMETMP"

# Processing statistics 
unset $BASECOLS $EXTRACOLS
if [ $GIS_FLAG_E -eq 1 ] ; then 
    r.univar -t -e map="$RASTER" \
    zones="${VECTOR}_$TMPNAME" percentile="$GIS_OPT_PERCENTILE" | \
    cut -f1,3,5-8,10-13,15-18 -d'|' | sed 's+nan+NULL+g' > "$STATSTMP"
else
    r.univar -t map="$RASTER" zones="${VECTOR}_$TMPNAME" | \
    cut -f1,3,5-8,10-13 -d'|' | sed 's+nan+NULL+g' > "$STATSTMP"
fi

if [ $GIS_FLAG_E -eq 1 ] && [ $DBFDRIVER -eq 1 ] ; then
    eval "$PERCCOL=\$percentile_$GIS_OPT_PERCENTILE"
fi

# create array with new column names
col1="`echo $COLNAMES | cut -f1 -d','`"
col2="`echo $COLNAMES | cut -f2 -d','`"
col3="`echo $COLNAMES | cut -f3 -d','`"
col4="`echo $COLNAMES | cut -f4 -d','`"
col5="`echo $COLNAMES | cut -f5 -d','`"
col6="`echo $COLNAMES | cut -f6 -d','`"
col7="`echo $COLNAMES | cut -f7 -d','`"
col8="`echo $COLNAMES | cut -f8 -d','`"
col9="`echo $COLNAMES | cut -f9 -d','`"
if [ $GIS_FLAG_E -eq 1 ] ; then
   col10="`echo $COLNAMES | cut -f10 -d','`"
   col11="`echo $COLNAMES | cut -f11 -d','`"
   col12="`echo $COLNAMES | cut -f12 -d','`"
   col13="`echo $COLNAMES | cut -f13 -d','`"
fi

# create SQL file for extended and normal statistics
g.message -v "Creating SQL file ..."
if [ $GIS_FLAG_E -eq 1 ] ; then
    sed -e '1d' "$STATSTMP" | awk -F "|" '{printf "\nUPDATE '${TABLE}' SET '${col1}' = %i , '${col2}' = %.2f , '${col3}' = %.2f , '${col4}' = %.2f , '${col5}' = %.2f , '${col6}' = %.2f , '${col7}' = %.2f , '${col8}' = %.2f , '${col9}' = %.2f , '${col10}' = %.2f , '${col11}' = %.2f , '${col12}' = %2f , '${col13}' = %.2f WHERE '${KEYCOL}' = %i;", $2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$1}' > "$SQLTMP"
else
    sed -e '1d' "$STATSTMP" | awk -F "|" '{printf "\nUPDATE '${TABLE}' SET '${col1}' = %i , '${col2}' = %.2f , '${col3}' = %.2f , '${col4}' = %.2f , '${col5}' = %.2f , '${col6}' = %.2f , '${col7}' = %.2f , '${col8}' = %.2f , '${col9}' = %.2f WHERE '${KEYCOL}' = %i;", $2,$3,$4,$5,$6,$7,$8,$9,$10,$1}' > "$SQLTMP"   
fi

g.message message="Updating the database ..."
db.execute input="$SQLTMP" database="${DB_DATABASE}" driver="$DB_SQLDRIVER"
EXITCODE=$?

cleanup

if [ "$EXITCODE" -eq 0 ] ; then
  g.message "Statistics calculated from raster map <$RASTER> and uploaded to \
    attribute table of vector map <$VECTOR>."
  g.message "Done."
fi

exit $EXITCODE
