#! /usr/env/python

"""
route_flow_dn.py:
    
Given a ModelGrid, calculates single-path (steepest direction) flow directions,
drainage area, and (optionally) discharge.

The "dn" in the name means that this is a generalization of the D8 algorithm,
for a grid in which a node has N neighbors (N might happen to be 8, or not).

Created GT Nov 2013
Modified to save data to grid directly, DEJH March 2014
"""

import landlab
#from landlab import RasterModelGrid
from landlab.components.flow_routing import flow_direction_DN
#reload(flow_direction_DN)
from landlab.components.flow_accum import flow_accum_bw
from landlab import FieldError, Component
from landlab import ModelParameterDictionary
import numpy
#from scipy.weave.build_tools import CompileError

#output_suppression_flag = True

class FlowRouter(Component):
    """
    This class implements single-path (steepest direction) flow routing, and 
    calculates flow directions, drainage area, and (optionally) discharge. 
    
    It initializes with a reference to a ModelGrid of any kind. Optionally, it
    can also take *input_params*, the string which is the name of a text input
    file. The input file can optionally contain 'runoff_rate', a float giving
    a (spatially constant) runoff rate. This is equivalent to the optional 
    input field 'water__volume_flux_in', and will override it if both are set.
    If neither are set, value will default to 1.
    
    The primary method of this class is :func:`route_flow`.
    """
    
    _name = 'DNFlowRouter'
    
    _input_var_names = set(['topographic_elevation',
                            'water__volume_flux_in',
                            ])
    
    _output_var_names = set(['drainage_area',
                             'flow_receiver',
                             'topographic__steepest_slope',
                             'water__volume_flux',
                             'upstream_ID_order',
                             'links_to_flow_receiver',
                             'flow_sinks',
                             ])
                             
    _var_units = {'topographic_elevation' : 'm',
                  'water__volume_flux_in' : 'm**3/s',
                  'drainage_area' : 'm**2',
                  'flow_receiver' : '-',
                  'topographic__steepest_slope' : '-',
                  'water__volume_flux' : 'm**3/s',
                  'upstream_ID_order' : '-',
                  'links_to_flow_receiver' : '-',
                  'flow_sinks' : '-',
                  }
    
    _var_mapping = {'topographic_elevation' : 'node',
                    'water__volume_flux_in' : 'node',
                    'drainage_area' : 'node',
                    'flow_receiver' : 'node',
                    'topographic__steepest_slope' : 'node',
                    'water__volume_flux' : 'node',
                    'upstream_ID_order' : 'node',
                    'links_to_flow_receiver' : 'node',
                    'flow_sinks' : 'node',
                    }
    
    _var_defs = {'topographic_elevation' : 'Land surface topographic elevation',
                 'water__volume_flux_in' : 'External volume water input to each node (e.g., rainfall)',
                 "drainage_area" : "Upstream accumulated surface area contributing to the node's discharge",
                 'flow_receiver' : 'Node array of receivers (node that receives flow from current node)',
                 'topographic__steepest_slope' : 'Node array of steepest *downhill* slopes',
                 'water__volume_flux' : 'Discharge of water through each node',
                 'upstream_ID_order' : 'Node array containing downstream-to-upstream ordered list of node IDs',
                 'links_to_flow_receiver' : 'ID of link downstream of each node, which carries the discharge',
                 'flow_sinks' : 'Boolean array, True at local lows',
                  }
    
    
    def __init__(self, model_grid, input_params=None):
        
        # We keep a local reference to the grid
        self._grid = model_grid
        self.value_field = 'topographic_elevation'
        
        if input_params:
            if type(input_params) == str:
                input_dict = ModelParameterDictionary(input_params)
            else:
                assert type(input_params) == dict
                input_dict = input_params
        
        # We'll also keep track of the active links; if raster, then these are
        # the "D8" links; otherwise, it's just activelinks
        if type(model_grid) is landlab.grid.raster.RasterModelGrid:
            dal, d8f, d8t = model_grid.d8_active_links()
            self._active_links = dal
            self._activelink_from = d8f
            self._activelink_to = d8t
        else:
            self._active_links = model_grid.active_links
            self._activelink_from = model_grid.activelink_fromnode
            self._activelink_to = model_grid.activelink_tonode
        
        #test input variables are present:
        model_grid.at_node['topographic_elevation']
        try:
            model_grid.at_node['water__volume_flux_in']
            self.field_for_runoff = True
        except FieldError:
            self.field_for_runoff = False
            #build the input array into the grid. This is important in case variable values appear during model run
            model_grid.add_ones('node', 'water__volume_flux_in')
            
        if not self.field_for_runoff:
            try:
                model_grid.at_node['water__volume_flux_in'].fill(input_dict['runoff_rate'])
            except (KeyError, UnboundLocalError):
                model_grid.at_node['water__volume_flux_in'].fill(1.)
        else:
            try:
                model_grid.at_node['water__volume_flux_in'].fill(input_dict['runoff_rate'])
            except (KeyError, UnboundLocalError):
                pass
            else:
                print "WARNING: Both a field and input parameter are available for runoff value. Was this intentional?? Taking the input parameter value..."
        
        # Keep track of the following variables:
        #   - drainage area at each node
        #   - receiver of each node
        self.drainage_area = model_grid.add_zeros('node', 'drainage_area')
        self.receiver = model_grid.create_node_array_zeros('flow_receiver')
        self.steepest_slope = model_grid.create_node_array_zeros('topographic__steepest_slope')
        self.discharges = model_grid.create_node_array_zeros('water__volume_flux')
        self.upstream_ordered_nodes = model_grid.create_node_array_zeros('upstream_ID_order')
        self.links_to_receiver = model_grid.create_node_array_zeros('links_to_flow_receiver')
        
        self.weave_flag = model_grid.weave_flag
        
        
    def route_flow(self):
        """
        Routes surface-water flow by (1) assigning to each node a single 
        drainage direction, and then (2) adding up the number of nodes that
        contribute flow to each node on the grid (including the node itself).
                
        Stores as ModelGrid fields:
            - Node array of receivers (nodes that receive flow): 
              *'flow_receiver'*
            - Node array of drainage areas: *'drainage_area'*
            - Node array of discharges: *'water__volume_flux'*
            - Node array of steepest downhill slopes: *'topographic__steepest_slope'*
            - Node array containing downstream-to-upstream ordered list of node
              IDs: *'upstream_ID_order'*
            - Node array containing ID of link that leads from each node to its
              receiver (or ITS OWN ID if there is no receiver):
              *'links_to_flow_receiver'*
            - Boolean node array of all local lows: *'flow_sinks'*
        
        Returns:
            - the modified grid object
        
        Examples
        --------
        >>> import numpy as np
        >>> from landlab import RasterModelGrid
        >>> from landlab.components.flow_routing.route_flow_dn import FlowRouter
        >>> mg = RasterModelGrid(5, 4, 1.0)
        >>> elev = np.array([0.,  0.,  0., 0.,
        ...                  0., 21., 10., 0.,
        ...                  0., 31., 20., 0.,
        ...                  0., 32., 30., 0.,
        ...                  0.,  0.,  0., 0.])
        >>> _ = mg.add_field('node','topographic_elevation', elev)
        >>> mg.set_closed_boundaries_at_grid_edges(False, True, True, True)
        >>> fr = FlowRouter(mg)
        >>> mg = fr.route_flow()
        >>> mg.at_node['flow_receiver']
        array([ 0,  1,  2,  3,  4,  1,  2,  7,  8,  6,  6, 11, 12, 10, 10, 15, 16,
               17, 18, 19])
        >>> mg.at_node['drainage_area']
        array([ 1.,  2.,  6.,  1.,  1.,  1.,  5.,  1.,  1.,  1.,  3.,  1.,  1.,
                1.,  1.,  1.,  1.,  1.,  1.,  1.])

        Now let's change the cell area and the runoff rates:
        
        >>> mg = RasterModelGrid(5, 4, 10.) #so cell area==100.
        >>> _ = mg.add_field('node','topographic_elevation', elev) #put the data back into the new grid
        >>> mg.set_closed_boundaries_at_grid_edges(False, True, True, True)
        >>> fr = FlowRouter(mg)
        >>> runoff_rate = np.arange(mg.number_of_nodes)
        >>> _ = mg.add_field('node', 'water__volume_flux_in', runoff_rate)
        >>> mg = fr.route_flow()
        >>> mg.at_node['water__volume_flux']
        array([    0.,   600.,  5400.,   300.,   400.,   500.,  5200.,   700.,
                 800.,   900.,  3700.,  1100.,  1200.,  1300.,  1400.,  1500.,
                1600.,  1700.,  1800.,  1900.])
        
        """
        
        #if elevs is not provided, default to stored grid values, which must be provided as grid
        elevs = self._grid['node'][self.value_field]
        
        node_cell_area = self._grid.forced_cell_areas
            
        
        # Calculate the downhill-positive slopes at the d8 active links
        #TODO: generalize to use EITHER D8, if raster, or just active links,
        # otherwise.
        link_slope = -self._grid.calculate_gradients_at_d8_active_links(elevs)
        # Find the baselevel nodes
        (baselevel_nodes, ) = numpy.where(numpy.logical_or(self._grid.node_status==1, self._grid.node_status==2))

        # Calculate flow directions
        receiver, steepest_slope, sink, recvr_link  = \
            flow_direction_DN.flow_directions(elevs, self._active_links, 
                                         self._activelink_from,
                                         self._activelink_to, link_slope, 
                                         grid=self._grid,
                                         baselevel_nodes=baselevel_nodes, 
                                         use_weave=self.weave_flag)
#############grid=None???
        
        # TODO: either need a way to calculate and return the *length* of the
        # flow links, OR the caller has to handle the raster / non-raster case.
        
        #print 'sinks:', sink

        # Calculate drainage area, discharge, and ...
        a, q, s = flow_accum_bw.flow_accumulation(receiver, sink,
                                                  node_cell_area=node_cell_area, 
                                                  runoff_rate=self._grid.at_node['water__volume_flux_in'],
                                                  use_weave=self.weave_flag)
                                                  
        #added DEJH March 2014:
        #store the generated data in the grid
        self._grid['node']['drainage_area'] = a
        self._grid['node']['flow_receiver'] = receiver
        self._grid['node']['topographic__steepest_slope'] = steepest_slope
        self._grid['node']['water__volume_flux'] = q
        self._grid['node']['upstream_ID_order'] = s
        self._grid['node']['links_to_flow_receiver'] = recvr_link
        self._grid['node']['flow_sinks'] = numpy.zeros_like(receiver, dtype=bool)
        self._grid['node']['flow_sinks'][sink] = True
        
        return self._grid

    @property
    def node_drainage_area(self):
        return self._grid['node']['drainage_area']

    @property
    def node_receiving_flow(self):
        return self._grid['node']['flow_receiver']

    @property
    def node_steepest_slope(self):
        return self._grid['node']['topographic__steepest_slope']

    @property
    def node_water_discharge(self):
        return self._grid['node']['water__volume_flux']

    @property
    def node_order_upstream(self):
        return self._grid['node']['upstream_ID_order']

    @property
    def link_to_flow_receiving_node(self):
        return self._grid['node']['links_to_flow_receiver']


if __name__ == '__main__':
    import doctest
    doctest.testmod()
