'''
cascade.py

Description:    This is first cut script to help users of the USGS code GSFLOW to calculate PRMS Cascade direction that make use of irregular polygons.
Created by:     C. Justin Mayers (cjmayers@usgs.gov)
Created:        4/15/2010
Last Modified:  5/14/2010
'''

#import system modules
import arcgisscripting
import datetime
import os
import sys
import time
import traceback

print 'Spam...'
startTime = time.time()

#local variables...
gp = arcgisscripting.create(9.3)
gp.AddToolbox("C:/ArcGIS/ArcToolbox/Toolboxes/Data Management Tools.tbx")
gp.overwriteoutput = 1

#input variables
inputDIR = os.path.dirname(sys.argv[1])  #parse input directory
inFC = os.path.basename(sys.argv[1])     #parse input FC name
fieldID = sys.argv[2]                    #identify unique field name
fieldElev = sys.argv[3]                  #identify elevation field name
outputDIR = os.path.dirname(sys.argv[4]) #parse output directory
outFC = os.path.basename(sys.argv[4])    #parse output FC name

gp.workspace = inputDIR  #set workspace

try:
    #make a feature layer of entire input feature class
    inFC_lyr = os.path.basename(inFC)[:-4] + '_lyr'
    txt = 'Make Feature Layer: ' + inFC_lyr
    print(txt)
    gp.AddMessage(txt)
    gp.MakeFeaturelayer_management(inFC, inFC_lyr)
    
    #collect information about the feature class
    desc = gp.Describe(inFC)             #create object containing feature class information
    coordsys = desc.SpatialReference     #coordinate system of feature class
    shapefieldname = desc.ShapeFieldName #feature class shape field name
    FCchar = '"' #character to enclose feature class field - assume shapefile for now

    #gather info about each polygon in the feature class    
    txt = 'Gather information about individual polygons'
    print(txt)
    gp.AddMessage(txt)
    polygonInfo = {}  #empty dictionary to store polygon information
    rows = gp.SearchCursor(inFC)
    row = rows.Next()
    
    #cycle through each polygon record
    while row: 
        ID_va = row.getvalue(fieldID)                            #unique ID of polygon
        geometry = row.getvalue(shapefieldname)                  #create object with geometry information
        polygonInfo[ID_va] = {}
        polygonInfo[ID_va][fieldElev] = row.getvalue(fieldElev)  #elevation value
        polygonInfo[ID_va]['x'] = geometry.truecentroid.x        #x coordinate value
        polygonInfo[ID_va]['y'] = geometry.truecentroid.y        #y coordinate value
        polygonInfo[ID_va]['CascadeID'] = ''                     #place holder for ID value of polygon that will receive surface-water cascade
        row = rows.Next()
    del rows, row
    
    #cycle through individual polygons 
    for ID_va in polygonInfo.iterkeys():
        #select individual polygon
        where_clause = FCchar + fieldID + FCchar + ' = ' + str(ID_va)
        selection_type = 'NEW_SELECTION'
        gp.SelectLayerByAttribute_management (inFC_lyr, selection_type, where_clause) 
        
        #select all polygons that touch individual polygon (inclusive)
        overlap_type = 'SHARE_A_LINE_SEGMENT_WITH'
        selection_type = 'NEW_SELECTION'
        gp.SelectLayerByLocation_management (inFC_lyr, overlap_type, inFC_lyr,'#', selection_type) 
        selectedRows = gp.SearchCursor(inFC_lyr)
        selectedRow = selectedRows.Next()

        #collect elevation data and assign CascadeID
        elev_info = []
        while selectedRow:
            adj_HRU = selectedRow.getvalue(fieldID) #adjacent HRU... will also look at itself (may be a low spot)
            elev_va = selectedRow.getvalue(fieldElev) #elevation value
            elev_info.append((elev_va,adj_HRU)) #update elevation list
            selectedRow = selectedRows.Next()   
        del selectedRows, selectedRow
        elev_info.sort() # sort elevation list
        polygonInfo[ID_va]['CascadeID'] = elev_info[0][1] #assign Cascade ID - second part of elev_info tupple
        
    #create feature class with polylines with cascades
    txt = 'Create feature class: ' + outFC
    print(txt)
    gp.AddMessage(txt)
    gp.workspace = outputDIR
    if gp.exists(outFC): #check for existing feature class
        gp.delete_management(outputDIR+'\\'+outFC)
    gp.CreateFeatureclass(outputDIR, outFC, "POLYLINE")
    gp.DefineProjection_management(outFC, coordsys)
    
    #add fields to the new feature class
    for fieldName, fieldType in [('HRU_Up','LONG'),('HRU_Down','LONG'),('SEG_Down','LONG'),
                                 ('Percent','FLOAT'),('Swale','LONG')]:
        gp.AddField_management (outFC, fieldName, fieldType) 
    rows = gp.InsertCursor(outFC)
    row = rows.Next()
    lineArray = gp.CreateObject("Array") #line array object
    point = gp.CreateObject("Point")     #point object
    
    #create polylines and populate fields
    for ID_va in polygonInfo.iterkeys():
        #set the x and y coordinates for origin vertex.
        point.x = polygonInfo[ID_va]['x']
        point.y = polygonInfo[ID_va]['y']
        lineArray.add(point) # add point to line array
        
        #set the x and y coordinates for destination vertex
        cascadeID = polygonInfo[ID_va]['CascadeID']
        point.x = polygonInfo[cascadeID]['x']
        point.y = polygonInfo[cascadeID]['y']
        lineArray.add(point) # add point to line array

        #insert the new polyline into the feature class.
        newFeature = rows.NewRow()
        newFeature.shape = lineArray
        if ID_va == cascadeID: #check for source == destination
            newFeature.Swale = 1
        else:
            newFeature.Swale = 0
        newFeature.HRU_Up = ID_va
        newFeature.HRU_Down = cascadeID
        newFeature.Percent = 1.0
        rows.InsertRow(newFeature)
        lineArray.RemoveAll() #clear line array
    del rows, row, point, lineArray
    gp.DeleteField_management (outFC, 'Id') #delete unnecessary  field

except:
    #catch and report errors modified from ESRI site
    tb = sys.exc_info()[2]
    tbinfo = traceback.format_tb(tb)[0]
    pymsg = "PYTHON ERRORS:\nTraceback Info:\n" + tbinfo + "\nError Info:\n     " +        str(sys.exc_type) + ": " + str(sys.exc_value) + "\n"
    msgs = "GP ERRORS:\n" + gp.GetMessages(2) + "\n"
    gp.AddError(msgs)
    gp.AddError(pymsg)
    print(msgs)
    print(pymsg)
    gp.AddMessage(gp.GetMessages(1))
    print(gp.GetMessages(1))

endTime = time.time()
print '...Eggs - ', (endTime - startTime)/60, 'minutes'