#!/usr/bin/env python3
#
############################################################################
#
# MODULE:    r.buffer.lowmem
# AUTHOR(S): Glynn Clements
# PURPOSE:   Low-memory replacement for r.buffer using r.grow.distance
#
# COPYRIGHT: (C) 2008, 2010 by Glynn Clements
#
#   This program is free software under the GNU General Public
#   License (>=v2). Read the file COPYING that comes with GRASS
#   for details.
#
#############################################################################

# %module
# % label: Creates a raster map showing buffer zones surrounding cells that contain non-NULL category values.
# % description: This is the low-memory alternative to the classic r.buffer module.
# % keyword: raster
# % keyword: buffer
# %end
# %flag
# % key: z
# % description: Ignore zero (0) data cells instead of NULL cells
# %end
# %option G_OPT_R_INPUT
# %end
# %option G_OPT_R_OUTPUT
# %end
# %option
# % key: distances
# % type: double
# % required: yes
# % multiple: yes
# % description: Distance zone(s)
# %end
# %option G_OPT_M_UNITS
# % options: meters,kilometers,feet,miles,nautmiles
# % description: Units of distance
# % answer: meters
# %end

import os
import atexit
import grass.script as grass
from grass.script.utils import encode


scales = {
    'meters': 1.0,
    'kilometers': 1000.0,
    'feet': 0.3048,
    'miles': 1609.344,
    'nautmiles': 1852.0
}

# what to do in case of user break:


def cleanup():
    if grass.find_file(temp_src)['file']:
        grass.run_command('g.remove', quiet=True, flags='fb', type='raster', name=temp_src)
    if grass.find_file(temp_dist)['file']:
        grass.run_command('g.remove', quiet=True, flags='fb', type='raster', name=temp_dist)


def main():
    global temp_dist, temp_src

    input = options['input']
    output = options['output']
    distances = options['distances']
    units = options['units']
    zero = flags['z']

    tmp = str(os.getpid())
    temp_dist = "r.buffer.tmp.%s.dist" % tmp
    temp_src = "r.buffer.tmp.%s.src" % tmp

    # check if input file exists
    if not grass.find_file(input)['file']:
        grass.fatal(_("Raster map <%s> not found") % input)

    scale = scales[units]

    distances = distances.split(',')
    distances1 = [scale * float(d) for d in distances]
    distances2 = [d * d for d in distances1]

    s = grass.read_command("g.proj", flags='j')
    kv = grass.parse_key_val(s)
    if kv['+proj'] == 'longlat':
        metric = 'geodesic'
    else:
        metric = 'squared'

    grass.run_command('r.grow.distance', input=input, metric=metric,
                      distance=temp_dist, flags='m')

    if zero:
        exp = "$temp_src = if($input == 0,null(),1)"
    else:
        exp = "$temp_src = if(isnull($input),null(),1)"

    grass.message(_("Extracting buffers (1/2)..."))
    grass.mapcalc(exp, temp_src=temp_src, input=input)

    exp = "$output = if(!isnull($input),$input,%s)"
    if metric == 'squared':
        for n, dist2 in enumerate(distances2):
            exp %= "if($dist <= %f,%d,%%s)" % (dist2, n + 2)
    else:
        for n, dist2 in enumerate(distances1):
            exp %= "if($dist <= %f,%d,%%s)" % (dist2, n + 2)
    exp %= "null()"

    grass.message(_("Extracting buffers (2/2)..."))
    grass.mapcalc(exp, output=output, input=temp_src, dist=temp_dist)

    p = grass.feed_command('r.category', map=output,
                           separator=':', rules='-')
    msg = "1:distances calculated from these locations\n"
    p.stdin.write(encode(msg))
    d0 = "0"
    for n, d in enumerate(distances):
        msg = "%d:%s-%s %s\n" % (n + 2, d0, d, units)
        p.stdin.write(encode(msg))
        d0 = d
    p.stdin.close()
    p.wait()

    grass.run_command('r.colors', map=output, color='rainbow')

    # write cmd history:
    grass.raster_history(output)

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