import ascraster
import os
import sys

icold = [0,-1,0,1,-1,0,1,-1,0,1]
irowd = [0,1,1,1,0,0,0,-1,-1,-1]


def goto_nextcell(icell,direction,ncols,mask=None):
    '''
    Return the cell number of the first cell which is found in the direction of direction.
    When the grid is mask, this is difficult.
    When the grid is not masked, it is easy.
    '''
    if (mask == None):
        return icell + irowd[direction]*ncols + icold[direction]  
    else:
        # The grid is masked
        # icell_new is the number of the gridcell in the unmasked grid.
        icell_new = mask[icell] + irowd[direction]*ncols + icold[direction]
        try:
            if (icell_new > mask[icell]):
                # Find value icell_new in mask[icell:]
                id = mask.index(icell_new,icell)
            else:
                icell_new_pointer = icell + irowd[direction]*ncols + icold[direction]
                if (icell_new_pointer > mask[-1]): icell_new_pointer = mask[-1]
                if (icell_new_pointer < 0): icell_new_pointer = 0

                # Find value icell_new in mask[icell_new_pointer:]
                id = mask.index(icell_new,icell_new_pointer)
        except ValueError:        
            #print "The next cell is not in the mask for cell: " + str(icell) + " and direction " +str(direction)
            id = -1

        return id


def factor(icell,retentionmap):
    '''
    When retentionmap is given is will give the 1.0 - value of retentionmap in the cell icell.
    Otherwise it will return 1.
    '''
    if (retentionmap == None):
        return 1.0
    else:
        fac = retentionmap.get_data(icell,0.0)
        if (fac < 0.0):
            print "retentionmap is not correct. It contains negative value of " + str(fac) + " in cell: "+str(icell)
            # Change retention map.
            retentionmap.set_data(icell,0.0)
            return 0.0
        elif (fac > 1.0):    
            print "retentionmap is not correct. It contains positive value (> 1.0) of " + str(fac) + " in cell: "+str(icell)
            # Change retention map.
            retentionmap.set_data(icell,1.0)
            return 1.0
        else:
            return (1.0 - fac)
            
def accuflux(lddmap,fluxmap,retentionmap=None,negative_flux_possible=0):
    '''
    ACCUFLUX
    Based on Gerard van Drecht, RIVM, Bilthoven, February 2002
    April 2002: accu fraction flux and accu fraction state added
    
    PURPOSE:
    Accumulate flow of mass, given by a material map [grid],
    along a flowpath, given by a local drain direction(LDD) map [grid].
    The output map contains accumulated flux [grid].
    All files are ascii-grid formatted.
    
    Lddmap: network of local drain directions
    flow_direction can be coded as follows:
    UNH/GRDC                       PCraster(LDD)
    32 64 128   meaning: NW N NE       7 8 9
    16  -   1             W -  E       4 5 6
     8  4   2            SW S SE       1 2 3
    We use the PCraster LDD-format; negative and zero values are assumed
    '''
    
    # Make output rasters
    accufluxmap = ascraster.duplicategrid(fluxmap)
    # Make nodata on all cells.
    accufluxmap.add_values(accufluxmap.length*[accufluxmap.nodata_value])
           
    # Loop over all cells:    
    for icell in range(lddmap.length):
        # Get mass flux at the cell.
        flux = fluxmap.get_data(icell)
        if (flux == None):
            # No data value
            continue
        if (flux < 0.0):
            if (negative_flux_possible == 0):
                print "Negative mass flow of " + str(flux) + " at cell: " + str(icell)           
                # Go to next cell.
                continue
            
        end_of_path=False
        icell_loop = icell
        while not end_of_path:
            # Get direction of the ldd 
            direction = int(lddmap.get_data(icell_loop,-1))
            if (direction < 1):
                #print "Something is wrong with the lddmap (< 1). No river mouth found for icell: ",icell, " with load: ", flux, " and direction: ",direction
                # Go to next cell
                end_of_path=True
            elif (direction == 5):
                # Mouth of river basin is found.
                prev = accufluxmap.get_data(icell_loop,0.0)
                accufluxmap.set_data(icell_loop,prev + flux * factor(icell_loop,retentionmap))
                # Go to next cell
                end_of_path=True
            elif (direction > 9):
                #print "Something is wrong with the lddmap (> 9). No river mouth found for icell: ",icell, " with load: ", flux, " and direction: ",direction
                # Go to next cell
                end_of_path=True
            else:
                prev = accufluxmap.get_data(icell_loop,0.0)
                flux = flux * factor(icell_loop,retentionmap)
                accufluxmap.set_data(icell_loop,prev + flux)
                
                # Go to next cell
                icell_loop = goto_nextcell(icell_loop,direction,lddmap.ncols,mask=lddmap.mask)
                if (icell_loop < 0):
                    # Already printed a message. Go to next grid cell.
                    end_of_path=True

    return accufluxmap


def determine_order(lddmap):
    # Make output raster with order of the cell
    ordermap = ascraster.duplicategrid(lddmap)
    # Fill all cells with zero for the order and one for the valuemap.
    ordermap.add_values(ordermap.length*[0])

    # Loop over all cells:    
    for icell in range(lddmap.length):         
        end_of_path=False
        icell_loop = icell
        icell_loop_prev = icell
        while not end_of_path:
            lfirst = (icell_loop_prev == icell_loop)
            # Get direction of the ldd 
            direction = int(lddmap.get_data(icell_loop,-1))
            neighbour = int(ordermap.get_data(icell_loop_prev,0))
            if (direction < 1):
                # Go to next cell
                end_of_path=True
            elif (direction == 5):
                # Mouth of river basin is found.
                prev = int(ordermap.get_data(icell_loop,0))
                if (lfirst):
                    if (prev < neighbour):
                        ordermap.set_data(icell_loop,neighbour + 1)
                else:
                    if (prev < neighbour + 1):
                        ordermap.set_data(icell_loop,neighbour + 1)
                # Go to next cell
                end_of_path=True
            elif (direction > 9):
                # Go to next cell
                end_of_path=True
            else:
                prev = int(ordermap.get_data(icell_loop,0))
                if (lfirst):
                    if (prev < neighbour):
                        ordermap.set_data(icell_loop,neighbour + 1)
                else:
                    if (prev < neighbour + 1):
                        ordermap.set_data(icell_loop,neighbour + 1)
                
                # Go to next cell
                icell_loop_prev = icell_loop
                icell_loop = goto_nextcell(icell_loop,direction,lddmap.ncols,mask=lddmap.mask)

                if (icell_loop < 0):
                    # Already printed a message. Go to next grid cell.
                    end_of_path=True

    return ordermap

def accutime(lddmap,timemap,lake_icell_dict):
    '''
    ACCUTIME
   
    PURPOSE:
    Accumulate residence time for each grid cell to the mouth of the river.

    The output map contains accumulated residence time [grid].
    When a endoreic lake is encountered then, residence time is to the lake and not to the sea.
    All files are ascii-grid formatted.
    
    Lddmap: network of local drain directions

    flow_direction can be coded as follows:
    UNH/GRDC                       PCraster(LDD)
    32 64 128   meaning: NW N NE       7 8 9

    16  -   1             W -  E       4 5 6
     8  4   2            SW S SE       1 2 3
    We use the PCraster LDD-format; negative and zero values are assumed
    '''
    
    # Make output rasters
    accutimemap = ascraster.duplicategrid(timemap)
          
    # Loop over all cells:    
    for icell in range(lddmap.length):
        # Get residence time at the cell.
        restime = timemap.get_data(icell)
        if (restime == None):
            # No data value
            continue
        if (restime < 0.0):
            print "Negative residence time of " + str(restime) + " at cell: " + str(icell)           
            # Go to next cell.
            accutimemap.set_data(icell,0.0)
            continue
            
        end_of_path=False
        icell_loop = icell
        cumtime = restime
        while not end_of_path:
            # Get direction of the ldd 
            direction = int(lddmap.get_data(icell_loop,-1))
            # Check whether the cell is in a endoreic lake
            try:
                 lendo_lake = lake_icell_dict[icell].lendolake
            except KeyError:
                 lendo_lake = 0
            if (direction < 1 or direction > 9):
                #print "Something is wrong with the lddmap (< 1). No river mouth found for icell: ",icell, " with load: ", flux, " and direction: ",direction
                # Go to next cell
                end_of_path=True
            elif (direction == 5):
                # Mouth of river basin is found.
                end_of_path=True
            elif (lendo_lake == 1):
                # Endoreic lake found.
                end_of_path=True
            else:               
                # Go to next cell
                icell_loop = goto_nextcell(icell_loop,direction,lddmap.ncols,mask=lddmap.mask)
                if (icell_loop < 0):
                    # Already printed a message. Go to next grid cell.
                    end_of_path=True

            # Set accutime in grid or fetch new restime for next cell
            if (end_of_path):
                accutimemap.set_data(icell,cumtime)
            else:
                restime = timemap.get_data(icell_loop,0.0)
                cumtime += restime
       

    return accutimemap

