#!/usr/bin/env python
############################################################################
#
# MODULE:       r.colors.stddev
# AUTHOR:       M. Hamish Bowman, Dept. Marine Science, Otago Univeristy,
#                 New Zealand
#               Converted to Python by Glynn Clements
# PURPOSE:      Set color rules based on stddev from a map's mean value.
#
# COPYRIGHT:    (c) 2007,2009-2010 Hamish Bowman, and 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: Sets color rules based on stddev from a raster map's mean value.
#% keyword: raster
#% keyword: color table
#%end
#% option G_OPT_R_MAP
#%end
#%flag
#% key: b
#% description: Color using standard deviation bands
#%end
#%flag
#% key: z
#% description: Force center at zero
#%end

import os
import atexit

import grass.script as gscript


def z(n):
    return mean + n * stddev


def cleanup():
    if tmpmap:
        gscript.run_command('g.remove', flags='f', type='raster',
                            name=tmpmap, quiet=True)


def main():
    global tmpmap
    tmpmap = None

    map = options['map']
    zero = flags['z']
    bands = flags['b']

    if not zero:
        s = gscript.read_command('r.univar', flags='g', map=map)
        kv = gscript.parse_key_val(s)
        global mean, stddev
        mean = float(kv['mean'])
        stddev = float(kv['stddev'])

        if not bands:
            # smooth free floating blue/white/red
            rules = '\n'.join(["0% blue",
                               "%f blue" % z(-2),
                               "%f white" % mean,
                               "%f red" % z(+2),
                               "100% red"])
        else:
            # banded free floating  black/red/yellow/green/yellow/red/black

            # reclass with labels only works for category (integer) based maps
            # r.reclass input="$GIS_OPT_MAP" output="${GIS_OPT_MAP}.stdevs" <<
            # EOF

            # >3 S.D. outliers colored black so they show up in d.histogram w/ white background
            rules = '\n'.join(["0% black",
                               "%f black" % z(-3),
                               "%f red" % z(-3),
                               "%f red" % z(-2),
                               "%f yellow" % z(-2),
                               "%f yellow" % z(-1),
                               "%f green" % z(-1),
                               "%f green" % z(+1),
                               "%f yellow" % z(+1),
                               "%f yellow" % z(+2),
                               "%f red" % z(+2),
                               "%f red" % z(+3),
                               "%f black" % z(+3),
                               "100% black"])
    else:
        tmpmap = "r_col_stdev_abs_%d" % os.getpid()
        gscript.mapcalc("$tmp = abs($map)", tmp=tmpmap, map=map)

        # data centered on 0  (e.g. map of deviations)
        info = gscript.raster_info(tmpmap)
        maxv = info['max']

        # current r.univar truncates percentage to the base integer
        s = gscript.read_command('r.univar', flags='eg', map=map,
                                 percentile=[95.45,
                                             68.2689,
                                             99.7300])
        kv = gscript.parse_key_val(s)

        stddev1 = float(kv['percentile_68_2689'])
        stddev2 = float(kv['percentile_95_45'])
        stddev3 = float(kv['percentile_99_73'])

        if not bands:
            # zero centered smooth blue/white/red
            rules = '\n'.join(["%f blue" % -maxv,
                               "%f blue" % -stddev2,
                               "0 white",
                               "%f red" % stddev2,
                               "%f red" % maxv])
        else:
            # zero centered banded  black/red/yellow/green/yellow/red/black

            # >3 S.D. outliers colored black so they show up in d.histogram w/ white background
            rules = '\n'.join(["%f black" % -maxv,
                               "%f black" % -stddev3,
                               "%f red" % -stddev3,
                               "%f red" % -stddev2,
                               "%f yellow" % -stddev2,
                               "%f yellow" % -stddev1,
                               "%f green" % -stddev1,
                               "%f green" % stddev1,
                               "%f yellow" % stddev1,
                               "%f yellow" % stddev2,
                               "%f red" % stddev2,
                               "%f red" % stddev3,
                               "%f black" % stddev3,
                               "%f black" % maxv, ])

    gscript.write_command('r.colors', map=map, rules='-', stdin=rules)

if __name__ == "__main__":
    options, flags = gscript.parser()
    atexit.register(cleanup)
    main()