#!/usr/bin/env python
"""
plume_buffer.py
GRASS GIS model to generate plumes of material transport based on vector attributes.

Requires a raster ocean mask named ocean in the mapset.  Ideally, place this 
mask in the PERMANENT folder of the mapset.

Authors: Matthew T. Perry, Shaun C. Walbridge
""" 

import sys
import os
import math
import glob

def getArgs():
    try:
        vnamevname = sys.argv[1]
        if len(sys.argv) == 3:
            attribs = [sys.argv[2]]
        else:
            # default attributes to analyze
            attribs = ['SUM_IMPV','SUM_PESTC','SUM_FERTC']
        return vnamevname,attribs
    except:
        print sys.argv[0] + " usage: <vname name> {attribute}"
        sys.exit(1)

def getCatList(vname,attribs):
    catlist = {}

    cmd = "v.db.select -c map=%s column=cat,basin_id,%s" % (vname,','.join(attribs))
    lines = os.popen(cmd).read().strip().split('\n')
    for i in lines:
        pair = i.split('|')
        cat = pair[0]
        basin = pair[1]

        catlist[cat] = {'basin_id': basin, 'attribs' : {}}
        if len(attribs) == 1:
            catlist[cat]['attribs'][attribs[0]] = pair[2]
            catlist[cat]['max'] = float(pair[2])
        else:
            for i in range(len(attribs)):
                name = getName(attribs[i])
                # find max value
                if catlist[cat].has_key('max'):
                    catlist[cat]['max'] = max(float(pair[i+2]),   \
                                              catlist[cat]['max'])
                else:
                    catlist[cat]['max'] = float(pair[i+2])

                catlist[cat]['attribs'][name] = float(pair[i+2])

    return catlist

def cleanHouse(vname,categories):
    for (c, c_data) in categories.items():
        basin_id = c_data['basin_id']
        rasters  = ['%s%s%s' % (vname, r, basin_id) for \
                   r in ['_', '_buff_', '_cost_', '_dw_']]
        cmd = "g.remove vect=%s_%s" % (vname, basin_id)
        os.popen(cmd)
        
        cmd = "g.remove rast=%s" % (','.join(rasters))
        os.popen(cmd)

    return True

def getName(name):
    # map column names to clean ones
    names = {'SUM_FERTC': 'fert',
             'SUM_PESTC': 'pest',
             'SUM_IMPV' : 'impv',
             'z_inorg' : 'inorganic',
             'z_nutri' : 'nutrients',
             'z_organ' : 'organic',
             'Sed_Increa' : 'sed_increase',
             'Sed_Decrea' : 'sed_decrease',
             'z_ports' : 'ports',
             'beach_pop' : 'beach_pop'
             }
    if names.has_key(name):
        return names[name]
    else:
        return name

def getEnv():
    cmd = "g.gisenv"
    lines = os.popen(cmd).read().rstrip().replace("'","").replace(';','').split("\n")
    env = {}
    for o in [i.split('=') for i in lines]:
        env[o[0]] = o[1]

    return env

def getPath():
    # return raster path
    env = getEnv()
    return "%s/%s/%s/cellhd/" % (env['GISDBASE'],env['LOCATION_NAME'],env['MAPSET'])
    
def getMaxDist(max):
    # exponents to distance mappings
    exp_distances = { -5 : 5,
                      -4 : 10,
                      -3 : 10,
                      -2 : 80,
                      -1 : 120,
                       0 : 200,
                       1 : 220,
                       2 : 230,
                       3 : 240,
                       4 : 250,
                       5 : 260,
                       6 : 270,
                       7 : 280,
                       8 : 300,
                    }

    max_sn = '%e' % max
    exp    = int(max_sn.split('e')[1])

    return exp_distances[exp]

def getLimit(column):
    default = 0.0001 # distribution limit in non-standard columns

    # limits generated by plume_distribution.R
    limits = { 'FERTC_SUM': 0.0006000, \
               'PESTC_SUM': 0.0001,    \
               'IMPV_SUM' : 0.0003438557, \
               'z_inorg': 0.00015,      \
               'z_nutri': 0.000003,      \
               'z_ports': 0.005,      \
               'z_organ': 0.00015,      \
               'Sed_Increa': 0.001,     \
               'Sed_Decrea': 0.002,     \
               'beach_pop' : 0.75       \
             }

    if limits.has_key(column):
        return limits[column]
    else:
        return default

def processCategory(c, c_data, vname, log):
    basin_id = c_data['basin_id']
    maxdist = getMaxDist(c_data['max'])
    columns = c_data['attribs']
    mask = 'ocean'
    
    # HACK: test for existing pours for restarting semi-complete jobs
    # don't use g.mlist as it drags ass
    plumes_create = []
    plumes = [os.path.basename(g) for g in glob.glob('%s/plume_*_%s' % (getPath(), basin_id))]

    for col in columns.keys():
        pn = 'plume_%s_%s' % (col, basin_id)
        if pn not in plumes:
            plumes_create.append(pn)

    if len(plumes_create) == 0:
        print "\n Plumes exist for %s, skipping.\n" % basin_id
        log.write("%s: all plumes exist, skipping." % basin_id)
        return

    print "\n Processing basin %s \n" % basin_id
    log.write("%s,%s,%s\n" % (c, basin_id, columns))

    # Initialize region to entire map
    cmd = 'g.region rast=%s' % mask
    os.popen(cmd)

    # Extract the single point
    pour = '%s_%s' % (vname, basin_id)
    cmd = 'v.extract input=%s output=%s where="cat = %s" new=1' % \
           (vname, pour, c)
    os.popen(cmd)

    # Subset region down to the narrowest possible buffer map
    cmd = "g.region vect=%s align=%s" % (pour, mask)
    os.popen(cmd)

    # Find current point region
    cmd = "g.region -gp"
    regionsplit = os.popen(cmd).read().rstrip().split('\n')
    region = {}
    for r in regionsplit:
       pv = r.strip().split('=')
       if (pv[0] != ''):
           region[pv[0]] = float(pv[1])

    # Extend region based on input value
    cellsize = region['nsres']
    regionbuff = cellsize * maxdist
    n = region['n'] + regionbuff
    s = region['s'] - regionbuff
    w = region['w'] - regionbuff
    e = region['e'] + regionbuff
    cmd = "g.region n=%s s=%s w=%s e=%s align=%s" % (n, s, w, e, mask)
    os.popen(cmd)

    # Convert to raster
    cmd = "v.to.rast input=%s output=%s use=cat" % \
           (pour, pour) 
    os.popen(cmd)

    # Buffer point to assure pour point hits coast
    pourbuff = '%s_buff_%s' % (vname, basin_id)
    cmd ="r.buffer input=%s out=%s distances=3.5 units=kilometers" % \
          (pour, pourbuff)
    os.popen(cmd)

    dw = '%s_dw_%s' % (vname, basin_id)
    cost = '%s_cost_%s' % (vname, basin_id)

    # Calculate cost distance
    cmd = "r.cost -k input=%s max_cost=%s output=%s start_rast=%s" % \
          (mask, maxdist, cost, pourbuff)
    os.popen(cmd)

    # Mask out non-ocean cells
    cmd = 'r.mapcalc %s = "if( %s, %s)"' % (dw, mask, cost) 
    os.popen(cmd)

    # Area Weighted distribution of sediment 
    cmd = 'r.stats -c %s' % dw
    area_info = os.popen(cmd).read().rstrip().split('\n')
    area_list = [i.split(' ') for i in area_info][:-1]

    pct = 0.005  # percentage of material deposited at each buffer ring

    if len(area_list) == 0:
        log.write("%s,%s,Area list is zero length, all values are null.\n" % (basin_id,c))
        print "area list null."
        return

    for (col, value) in columns.items():
        init  = float(value)
        limit = getLimit(col) 
        sum, percell, remain = 0., 0., 0.
        recode_table = ''

        plume = 'plume_%s_%s' % (col, basin_id)
        if init == ('0' or ''):
            log.write("%s,%s=%s has a value of 0, skipping.\n" % basin_id, col, init)
            break

        orig = init
        for (dist,count) in area_list:
            if init * pct < limit:
                remain = 0
                percell = init / float(count)
            else:
                percell = init * pct
                sum = percell * float(count)
                remain = init - sum
            init = remain
            
            # log.write("%2f,%2f,%2f,%2f,%2f\n" % (dist, count, percell, sum, remain))
            recode_table += "%s:%s:%s\n" % (dist, dist, percell) 

        # Reclass the dw map and recreate the plume map
        cmd = "r.recode input=%s output=%s <<EOF \n%sEOF" % (dw,plume,recode_table)
        os.popen(cmd)

    # Clean up
    tmp = {}
    tmp[c] = c_data
    cleanHouse(vname, tmp)
    log.flush()

def handle(cmd):
    try:
        handle = os.popen(cmd, 'r', 1)
        return handle
    finally:
        handle.close()

def addPlumes(outputFile, column):
    plumes = glob.glob("%s/plume_%s*" % (getPath(), getName(column)))
    (name, ext) = os.path.splitext(outputFile)

    os.putenv('PYTHONPATH', '/opt/gdal/pymod') # hack
    addCmd = os.path.realpath(os.path.dirname(sys.argv[0])) + '/gdal_add.py'

    pl = len(plumes)
    if pl == 0:
        print "No data found for column %s" % column
        return False

    batch = 60
    tempids = []
    
    for i in range(0,pl,batch):
        start = i 
        end = i + batch - 1 # don't double count edges
        id = "plume_%s_%s_%s" % (name, start, end)
        cmd = "%s -o %s.tif -ot Float32 %s" % \
              (addCmd, id, ' '.join(plumes[start:end]))
        print id
        handle(cmd)

        # gdal_translate the results to minimize file footprint
        cmd = "gdal_translate %s.tif -co COMPRESS=PACKBITS %s.tiff" % (id, id)
        handle(cmd)

        tempids.append(id)
        os.remove("%s.tif" % id)
    
if __name__ == '__main__':
    vname, attrib = getArgs()
    categories = getCatList(vname,attrib)

    # create error log file
    logfile = "plume_%s.log" % (vname)
    log = open(logfile,'w')
    log.write('# plume_buffer log output\n')
    log.write('basin_id,category,message\n')

    i = 1
    for (c, c_data) in categories.items():
        print "========  %s of %s (%s percent) ================" % \
               (i,len(categories),int(100 * (float(i)/float(len(categories)))))
        processCategory(c, c_data, vname, log)
        #log.write("%s,%s,%s\n" % (cat,catlist[cat][0],catlist[cat][1]))
        i = i + 1
    """
    for att in attrib:
        addPlumes("%s_%s_total.tif" % (vname, att), att)
    """
    log.close()
