###
#
# Description: The following code is used to convert hourly rainfall depth grids for a storm to  
# hyetograph format that can be input in HEC-HMS.  
#
#
#
#
#
###
from dbfread import DBF
import pandas as pd
from pandas import DataFrame
import numpy as np
import datetime
import matplotlib.pyplot as plt
from  matplotlib.dates import DayLocator
import arcpy
arcpy.CheckOutExtension("Spatial")
arcpy.env.overwriteOutput = True

Chey_Subs = "T:\\projects\\jdngroup\\XXXX\\20180925_Chey_Subs.shp" # shapefile of sub-basin network for basin
asp_reclass = "T:\\projects\\jdngroup\\XXXX\\chey_nf_sf.shp" # shapefile of NF and SF slopes
asp_subs = "T:\\projects\\jdngroup\\XXXX\\20180925_Chey_Subs_asp.shp" # shapefile of subbasins separated into NF and SF slopes

d_North=np.empty([1,3]) # change based on number of subbasins
d_South=np.empty([1,3]) # change based on number of subbasins
h=0 #initialize counter
t=datetime.datetime(2013,9,8,8)  #initialize time counter
tarray=np.empty([1,1])
xlabels=np.empty([1,1])
aspect=['NFS','SFS']

# use for loop to iterate through each hour of storm
for x in range(0,240):# change the upper limit in the range to equal the duration of the storm in hours
    hr=str(t.hour).zfill(2)+"00" #extract current hour
    day=str(t.day).zfill(2) #extract current day
    print(x)
    # use for loop to iterate through NFS and SFS
    for y in aspect:
        # Define precipitation grid for each hour
        precip=("T:\projects\jdngroup\XXXX\P_allsites_spas1302_"+str(h+1).zfill(3)+"_201309"+day+"_"+hr+"_utc_ascii.asc_proj.tif")  # change path to location of rainfall depth grids
        
        ## Calculate spatial average precip for each aspect for each hour in .dbf table format
        if y == 'NFS':
            table_North=("T:\projects\jdngroup\XXXX\mean_201309"+day+hr+"North"+".dbf") # change path to desired location for table to be output
            arcpy.MakeFeatureLayer_management(asp_subs,"asp_layer")
            arcpy.SelectLayerByAttribute_management ("asp_layer", "NEW_SELECTION", """{0}='{1}'""".format("ASPECT",'North'))
            tempEnvironment0 = arcpy.env.extent
            arcpy.env.extent = Chey_Subs
            arcpy.gp.ZonalStatisticsAsTable_sa("asp_layer","Name", precip,table_North,"DATA", "MEAN")
            arcpy.env.extent = tempEnvironment0
            
            ## Convert hourly tables into total storm hyetograph for each subbasin
            dbf_North=DBF(table_North)
            print(dbf_North)
            df_North=DataFrame(iter(dbf_North), columns=['NAME','COUNT','AREA','MEAN'])
            print (df_North)
            s_North=np.array(df_North['MEAN'])
            d_North=np.vstack([d_North,s_North])
        else:
            table_South=("T:\projects\jdngroup\XXXX\mean_201309"+day+hr+"South"+".dbf") # change path to desired location for table to be output
            tempEnvironment0 = arcpy.env.extent
            arcpy.env.extent = Chey_Subs
            arcpy.MakeFeatureLayer_management(asp_subs,"asp_layer")
            arcpy.SelectLayerByAttribute_management ("asp_layer", "NEW_SELECTION", """{0}='{1}'""".format("ASPECT",'South'))
            arcpy.gp.ZonalStatisticsAsTable_sa("asp_layer","Name", precip,table_South,"DATA", "MEAN")
            arcpy.env.extent = tempEnvironment0   
        
            ## Convert hourly tables into total storm hyetograph for each subbasin
            dbf_South=DBF(table_South)
            df_South=DataFrame(iter(dbf_South), columns=['NAME','COUNT','AREA','MEAN'])
            s_South=np.array(df_South['MEAN'])
            d_South=np.vstack([d_South,s_South])
    tarray=np.append(tarray,str(t.month)+"/"+str(t.day)+" "+str(t.hour)+":00")
    t=t+datetime.timedelta(hours=1)
    h+=1
    print (t)

## fill pandas with storm hyetograph arrays        
d_North=np.delete(d_North,0,0)
d_South=np.delete(d_South,0,0)
d_North=np.multiply(2.54,d_North) # convert precip from inches to cm
d_South=np.multiply(2.54,d_South)
tarray=np.delete(tarray,0,0)
clmns=['Subbasin 1','Subbasin 2','Subbasin 3']
hyeto_North=pd.DataFrame(d_North,index=tarray,columns=clmns)
hyeto_South=pd.DataFrame(d_South,index=tarray,columns=clmns)

## save hyetographs to Excel      
writer = pd.ExcelWriter(r"T:\projects\jdngroup\XXXX\aspect_subs_hyetograph.xlsx") # change path to desired location for excel hyetographs to be output
hyeto_North.to_excel(writer, sheet_name='North')
hyeto_South.to_excel(writer, sheet_name='South')
writer.save()  

hyeto = [hyeto_North, hyeto_South] 
## plot hyetographs
for y in range(0,3):
    s=0
    for x in hyeto: 
        
        ycum = x[str(clmns[y])].cumsum()
        ax1=x.plot(use_index=True,y=clmns[y], secondary_y=['ycum'],kind='bar', legend=None, stacked=True)
        ax2=ycum.plot(use_index=True, secondary_y=True, label='Cumulative Precip (cm)', linestyle='--',linewidth=.5, color='r')
        ax1.set_xticklabels(x.index,rotation=60)
        n = 24
        ticks = ax1.xaxis.get_ticklocs()
        ticklabels = [l.get_text() for l in ax1.xaxis.get_ticklabels()]
        ax1.xaxis.set_ticks(ticks[::n])
        ax1.xaxis.set_ticklabels(ticklabels[::n])
        
        ax1.set_ylabel('Precipitation (cm/hr)', fontsize=12)
        ax2.set_ylabel('Cumulative Precipitation (cm)', fontsize=12)

        plt.title(str(clmns[y])+' '+str(aspect[s])+' Hyetograph', fontsize=15) 
        plt.tight_layout()
        plt.savefig("T:\projects\jdngroup\XXXX\Hyeto_"+str(clmns[y])+' '+str(aspect[s])+".jpg") # change path to desired location for hyetograph image
        plt.close('all')
        s+=1
