###
# Accum_Time_function.py
#
# Description:
# This function accumulates the travel times to the subbasin outlet for each point in the subbasin. Two for loops are used to iterate 
# through each cell in the subbasin. For each cell, the algorithm routes through the subbasin based on the flow direction raster.  
# As the algorithm follows the flow direction raster, it accumulates the individual cell travel times - which are calculated in Tc_model.py - until it reaches the subbasin outlet.
# Refer to Appendix A of Woolridge (2019) for additional information on using the script
#
# Created by: Doug Woolridge
# Colorado State University
# 
###
import numpy as np
import arcpy
import math
from arcpy.sa import *
from arcpy.conversion import *
arcpy.CheckOutExtension("Spatial")
arcpy.CheckOutExtension("3D")
arcpy.env.overwriteOutput = True
 
def Accum_Time(outlet,sub,fdr,travelt,cellsize): 
    ## convert time raster to array
    tempEnvironment0 = arcpy.env.extent
    arcpy.env.extent = sub
    t_cellsub = ExtractByMask(travelt,sub)
    t_cellarr=arcpy.RasterToNumPyArray(t_cellsub, nodata_to_value=0)
    arcpy.env.extent = tempEnvironment0
    print('time array complete')

    ## Pre-process fdr raster and create empty array with same size as subbasin
    # clip fdr to watershed, replace outlet value with 0, and create array from raster
    tempEnvironment0 = arcpy.env.extent 
    tempEnvironment1 = arcpy.env.cellSize
    tempEnvironment2 = arcpy.env.snapRaster
    arcpy.env.extent = sub
    arcpy.env.cellSize = sub
    arcpy.env.snapRaster = sub
    fdr_sub = ExtractByMask(fdr,sub) # extract flow direction raster for subbasin
    outrast = "T:\\projects\\jdngroup\\XXXX\\outrast"
    arcpy.PointToRaster_conversion(outlet,"Id", outrast) # convert outlet point shapefile to a raster; the shapefile should have an "Id" field with a value of 0 so that the raster cells have a value of 0
    arcpy.env.workspace = "T:\\projects\\jdngroup\\XXXX\\ClarkUH" # change path to the same folder where fdr_sub and outrast are saved
    arcpy.gp.MosaicToNewRaster_management([fdr_sub, outrast], "T:\\projects\\jdngroup\\XXXX\\ClarkUH","fdrout","T:\\projects\\jdngroup\\XXXX\\20180925_Chey_Subs_asp.prj","8_BIT_UNSIGNED","9.1587504","1","LAST") # merge fdr raster with outlet point raster; this creates a flow direction raster with outlet cells of value=0
    fdrout = "T:\\projects\\jdngroup\\XXXX\\fdrout" # flow direction raster with outlet cells having a value of 0
    fdrarr=arcpy.RasterToNumPyArray(fdrout) # create array of fdrout raster
    print('fdr array complete')

    a=np.size(fdrarr,0) #rows # create empty array with same size as watershed
    b=np.size(fdrarr,1) #columns
    Travel_matrix=np.zeros((a,b))                                         
    arcpy.env.extent = tempEnvironment0
    arcpy.env.cellSize = tempEnvironment1
    arcpy.env.snapRaster = tempEnvironment2

    ## clean array    
    for i in np.arange(0,a,1):      #The for loop is used to set any negative data to zero for the travel time raster
        for j in np.arange(0,b,1):
            if t_cellarr[i,j] < 0 or math.isnan(t_cellarr[i,j]):
                t_cellarr[i,j] = 0
    print('raster cleaned')
                            
    ## print check to make sure all arrays have the same dimensions
    print (a, b)
    print (np.shape(fdrarr))
    print (np.shape(t_cellarr))
    print(np.shape(Travel_matrix))             
    fdrcsv = np.savetxt("T:\\projects\\jdngroup\\XXXX\\fdrarr.csv",fdrarr) # save fdr array to csv
    tcellcsv = np.savetxt("T:\\projects\\jdngroup\\XXXX\\tcellarr.csv",t_cellarr) # save travel time array to csv

    for i in np.arange(0,a,1):               # The for loop iterates through each row and column to calculate the travel time to the outlet for each cell                           
        for j in np.arange(0,b,1):                                      
            print ('cell: ', i,j)
            #print (fdrarr[i,j])
            if fdrarr[i,j]!=255:             # the flow direction raster has a value of 255 outside the extent of the subbasin; the if statement checks that the current cell is in the subbasin
                p=i                          # set p and q as index of current cell                           
                q=j                                                     
                Travel_matrix[i,j] = t_cellarr[p,q]           # Travel_matrix stores the travel time to the subbasin outlet of each cell. It is given an initial value of the travel time through itself.
                while fdrarr[p,q] != 0:                       # The while loop routes through subbasin by following the flow direction raster and adding the travel time through each cell from the t_cellarr array to the Travel_matrix
                    if fdrarr[p,q]==1:                        # as it passes through. When the fdrarr value of the current cell is 0 (i.e. the outlet), the final time is stored in the to cell i,j in the Travel_matrix               
                        p=p
                        q=q+1
                    if fdrarr[p,q]==2:
                        p=p+1
                        q=q+1                    
                    if fdrarr[p,q]==4:
                        p=p+1
                        q=q                    
                    if fdrarr[p,q]==8:
                        p=p+1
                        q=q-1                    
                    if fdrarr[p,q]==16:
                        p=p
                        q=q-1
                    if fdrarr[p,q]==32:
                        p=p-1
                        q=q-1
                    if fdrarr[p,q]==64:
                        p=p-1
                        q=q
                    if fdrarr[p,q]==128:
                        p=p-1
                        q=q+1
                    
                    Travel_matrix[i,j] = Travel_matrix[i,j]+t_cellarr[p,q]                                                                               

    tempEnvironment0 = arcpy.env.outputCoordinateSystem
    tempEnvironment1 = arcpy.env.snapRaster
    arcpy.env.outputCoordinateSystem = arcpy.SpatialReference(26913)
    arcpy.env.snapRaster = sub
    Accum_time = "T:\\projects\\jdngroup\\XXXX\\Accum_time"
    Accum_time = arcpy.NumPyArrayToRaster(Travel_matrix,arcpy.Point(fdr_sub.extent.XMin,fdr_sub.extent.YMin),cellsize,cellsize)
    Accum_time.save("T:\\projects\\jdngroup\\XXXX\\Accum_time")
    Accumcsv = np.savetxt("T:\\projects\\jdngroup\\XXXX\\accumarr.csv",Travel_matrix)
    arcpy.env.outputCoordinateSystem = tempEnvironment0
    arcpy.env.snapRaster = tempEnvironment1
    print ('complete')
    
    return Accum_time


