#%% Plotting script for LSU D3D data

# Alexander Rey, 2022
import os
import pandas as pd
import geopandas as gp
import netCDF4 as nc
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.cm as cm

import datetime as datetime

from scipy.interpolate import LinearNDInterpolator, interp1d

from dfm_tools.get_nc import get_netdata, get_ncmodeldata, plot_netmapdata
from dfm_tools.get_nc_helpers import get_timesfromnc, get_ncfilelist, get_hisstationlist, get_ncvardimlist, get_timesfromnc
import cartopy.crs as ccrs
import contextily as ctx
from dfm_tools.regulargrid import scatter_to_regulargrid
import pathlib as pl

from shapely.geometry import Point, MultiPoint
import alphashape

import rasterio
from rasterio.plot import show as rioshow
from osgeo import osr #Spatial reference module

from os import listdir
from os.path import isfile, join
#%% Read Model Log
pth = pl.Path("//srv-mad3.baird.com/", "Projects", "13522.101 LSU Lakes", "06_Models", "Delft3DFM", "13522.678.W.SBV.Rev0_LSU Lakes D3D Model Log.xlsx")
runLog = pd.read_excel(pth.as_posix(), "ModelLog")

dataPath = "FlowFM_map.nc"

#%% Import using DFM functions
# modelPlot = range(53, 55) #30
# modelPlot = [66, 67, 68, 69, 70, 71]
# modelPlot = [44, 45]
modelPlot = [64, 65]

# stormTime =  [datetime.datetime(2005, 8,  29, 14, 0, 0),
#               datetime.datetime(2005, 9,  24, 14, 0, 0),
#               datetime.datetime(2005, 12, 28, 14, 0, 0)]
# stormTime =  [datetime.datetime(2005, 1, 14, 18, 0, 0),
#               datetime.datetime(2005, 2,  3, 18, 0, 0),
#               datetime.datetime(2005, 2,  10, 18, 0, 0)]
#stormTime =  [datetime.datetime(2005, 12, 31, 18, 0, 0)]
#stormTime =  [datetime.datetime(2005, 8, 7, 0, 0, 0)]


ugrid_all           = [None] * (max(modelPlot)+1)
X                   = [None] * (max(modelPlot)+1)
Y                   = [None] * (max(modelPlot)+1)
U                   = [None] * (max(modelPlot)+1)
V                   = [None] * (max(modelPlot)+1)
Xt                  = [None] * (max(modelPlot)+1)
Yt                  = [None] * (max(modelPlot)+1)
Ut                  = [None] * (max(modelPlot)+1)
Vt                  = [None] * (max(modelPlot)+1)
facex               = [None] * (max(modelPlot)+1)
facey               = [None] * (max(modelPlot)+1)
ux_mean             = [None] * (max(modelPlot)+1)
uy_mean             = [None] * (max(modelPlot)+1)
ux                  = [None] * (max(modelPlot)+1)
uy                  = [None] * (max(modelPlot)+1)
sed                 = [None] * (max(modelPlot)+1)
area                = [None] * (max(modelPlot)+1)
wd                  = [None] * (max(modelPlot)+1)
regularMask         = [None] * (max(modelPlot)+1)
regularMask_t       = [None] * (max(modelPlot)+1)
meshMask            = [None] * (max(modelPlot)+1)
alphaPolyList       = [None] * (max(modelPlot)+1)


for i in modelPlot:
    file_nc_map         = os.path.join(runLog['Run Location'][i], 'FlowFM', 'output', 'FlowFM_map.nc')
    # file_nc_map         = os.path.join(runLog['Run Location'][i], 'FlowFM', 'dflowfm', 'output', 'FlowFM_map.nc')
    # file_nc_map         = os.path.join(runLog['Run Location'][i], 'dflowfm', 'output', 'FlowFM_map.nc')

    tSteps              = get_timesfromnc(file_nc=file_nc_map, varname='time')

    # Find nearest time step to desired time
    # stormTime = [runLog['End Date'][i]]
    # tStep = []
    # for s in stormTime:
    #     abs_deltas_from_target_date = np.absolute(tSteps - s)
    #     tStep.append(np.argmin(abs_deltas_from_target_date))
    #
    tStep = len(tSteps) - 1
    # tStep               = 107 # Peak discharge tstep #40

    # Get Var info
    vars_pd, dims_pd = get_ncvardimlist(file_nc=file_nc_map)

    ugrid_all[i]        = get_netdata(file_nc=file_nc_map)#,multipart=False)
    facex[i]            = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_face_x')
    facey[i]            = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_face_y')
    # ux_mean[i]          = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucxa', timestep=tStep)
    # uy_mean[i]          = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucya', timestep=tStep)

    ux[i]               = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucx',
                                          timestep=tStep, layer='all')
    uy[i]               = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucy',
                                          timestep=tStep, layer='all')

    sed[i]              = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_bodsed',
                                          timestep=tStep, station='all')
    area[i]             = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_flowelem_ba')
    wd[i]               = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_waterdepth',
                                          timestep=tStep)

    ux_mean[i]          = ux[i].mean(axis=2)
    uy_mean[i]          = uy[i].mean(axis=2)

    X[i],Y[i],U[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                            ncellx=50, ncelly=70, values=ux_mean[i][0, :],
                                            maskland_dist=25, method='linear')
    X[i],Y[i],V[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                            ncellx=50, ncelly=70, values=uy_mean[i][0, :],
                                            maskland_dist=25, method='linear')

    Xt[i],Yt[i],Ut[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                            ncellx=75, ncelly=115, values=ux[i][0, :, 4],
                                            maskland_dist=25, method='linear')
    Xt[i],Yt[i],Vt[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                            ncellx=75, ncelly=115, values=uy[i][0, :, 4],
                                            maskland_dist=25, method='linear')


    # Find convex hull of pcolor data
    # Convert grid to multipoint

    # Create a list of grid points
    # alphaPoints = list(map(tuple, np.column_stack((facex[i][:], facey[i][:])).data))
    alphaPoints = list(map(tuple, np.column_stack((facex[i][wd[i][0, :] != 0],
                                                   facey[i][wd[i][0, :] != 0])).data))

    # Find the concave hull
    alphaPoly        = alphashape.alphashape(alphaPoints, 0.05)
    alphaPolyList[i] = alphaPoly

    # Create mask of regular grid inside concave hull
    regularMask[i] = [alphaPoly.contains(Point(i[0], i[1]))
                       for i in np.array([X[i].flatten(), Y[i].flatten()]).T]

    # Create mask of mesh grid points inside concave hull
    meshMask[i] = [alphaPoly.contains(Point(i[0], i[1]))
                       for i in np.array([facex[i], facey[i]]).T]

    # Reshape back to regular grid for mask
    regularMask[i] = np.array(regularMask[i]).reshape(X[i].shape)




    # Create mask of regular grid inside concave hull  of top later
    regularMask_t[i] = [alphaPoly.contains(Point(i[0], i[1]))
                       for i in np.array([Xt[i].flatten(), Yt[i].flatten()]).T]

    # Reshape back to regular grid for mask
    regularMask_t[i] = np.array(regularMask_t[i]).reshape(Xt[i].shape)

#%% Plot using DFM functions
modelPlot = [66, 67, 68, 69, 70, 71]
# stormDirs  = [306, 135, 250]
windCount  = 0
# fig, axes = plt.subplots(nrows=1, ncols=3,subplot_kw={'projection': ccrs.epsg(26915)}, figsize=(13, 6))

vmin        = 0 # Legend min
vmax        = 0.05 # Legend max
scale       = 100 # Size of arrows, smaller is bigger #25
qmax        = 0.02 # Scaling threshold. Above this value, arrows are scaled to 1 #0.1
amin        = 1 # Minimum arrow length. A dot is used below this value

for tPlot in range(0, 1):
    for i in modelPlot:
        fig, axes = plt.subplots(nrows=1, ncols=2, subplot_kw={'projection': ccrs.epsg(26915)}, figsize=(10, 6))
        fig.patch.set_facecolor('white')

        plt.tight_layout()

        plotCount = 1

        vel_mask = meshMask[i]

        # pc = plot_netmapdata(ugrid_all[i].verts[vel_mask],
        #                      values=np.sqrt(ux_mean[i][tPlot, :][vel_mask]**2 +\
        #                                     uy_mean[i][tPlot, :][vel_mask]**2),
        #                      ax=axes[plotCount], linewidth=0.5, cmap="viridis")
        pc = plot_netmapdata(ugrid_all[i].verts[vel_mask],
                             values=np.sqrt(ux[i][tPlot, :, 4][vel_mask]**2 +\
                                            uy[i][tPlot, :, 4][vel_mask]**2),
                             ax=axes[plotCount], linewidth=0.5, cmap="viridis")

        axes[plotCount].set_xlim(675524, 677027)
        axes[plotCount].set_ylim(3365191, 3367182)
        # axes[plotCount].set_xlim(676470, 677021)
        # axes[plotCount].set_ylim(3365939, 3366617)

        # axes.quiver(X[i], Y[i], U[i], V[i], color='w', scale=8)

        # Normalize velocity above limit

        # Using fixed grid
        overGrid_U_Scale = Ut[i]
        overGrid_V_Scale = Vt[i]

        # Using model grid
        # overGrid_U_Scale = ux_mean[i][tPlot, :]
        # overGrid_V_Scale = uy_mean[i][tPlot, :]

        r = np.power(np.add(np.power(overGrid_U_Scale, 2), np.power(overGrid_V_Scale, 2)), 0.5)
        curPlotMask = np.sqrt((overGrid_U_Scale) ** 2 + (overGrid_V_Scale) ** 2) > qmax

        overGrid_U_Plot = np.zeros(overGrid_U_Scale.shape)
        overGrid_V_Plot = np.zeros(overGrid_V_Scale.shape)

        # Normalize arrows to 1 where not zero
        overGrid_U_Plot[r != 0] = overGrid_U_Scale[r != 0] / r[r != 0]
        overGrid_V_Plot[r != 0] = overGrid_V_Scale[r != 0] / r[r != 0]

        # Scale arrows below qmax
        overGrid_U_Plot[~curPlotMask] = overGrid_U_Plot[~curPlotMask] * (r[~curPlotMask] / qmax)
        overGrid_V_Plot[~curPlotMask] = overGrid_V_Plot[~curPlotMask] * (r[~curPlotMask] / qmax)

        # Using fixed grid
        # Remove arrows outside of convex hull
        overGrid_U_Plot[~regularMask_t[i]] = np.nan
        overGrid_V_Plot[~regularMask_t[i]] = np.nan

        qv = axes[plotCount].quiver(Xt[i], Yt[i], overGrid_U_Plot, overGrid_V_Plot, scale=scale,
                          color='k', headlength=3.5, headwidth=3.5, headaxislength=3.5, width=0.0015, minlength=amin) #200

        # # Using model grid
        # # Remove arrows outside convex hull
        # overGrid_U_Plot[~np.array(meshMask[i])] = np.nan
        # overGrid_V_Plot[~np.array(meshMask[i])] = np.nan
        # qv = axes[plotCount].quiver(facex[i], facey[i], overGrid_U_Plot, overGrid_V_Plot, scale=scale,
        #                   color='k', headlength=3.5, headwidth=3.5, headaxislength=3.5, width=0.0015, minlength=amin) #200


        mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
        ctx.add_basemap(axes[plotCount], source=mapbox, crs='EPSG:26915')

        # extent = ax[plotCount].get_extent()
        # ax[plotCount].set_xticks(np.linspace(extent[0], extent[1], 4))
        # ax[plotCount].set_yticks(np.linspace(extent[2], extent[3], 4))


        # Add model bounds line
        # axes[plotCount].plot(*alphaPoly.exterior.xy, color='k', linewidth=1)

        pc.set_clim([vmin, vmax])
        if plotCount == 1:
            cbar = fig.colorbar(pc, ax=axes, shrink=0.95)
            cbar.set_label('Surface Velocity [m/s]')

        axes[plotCount].set_title(runLog['Bathymetry'][i] + ' Bathymetry; Wind from ' + str(runLog['Wind Dir'][i]) + ' Degrees')

        # Add rectangle to show storm arrow
        xLimDist = axes[plotCount].get_xlim()[1] - axes[plotCount].get_xlim()[0]
        yLimDist = axes[plotCount].get_ylim()[1] - axes[plotCount].get_ylim()[0]

        axes[plotCount].add_patch(patches.Rectangle((axes[plotCount].get_xlim()[0] + 0.7 * xLimDist,
                        axes[plotCount].get_ylim()[0] + 0.8 * yLimDist - xLimDist * 0.1),
                        xLimDist * 0.2, xLimDist * 0.2,
                        facecolor='white', edgecolor='w', linewidth=5))


        # Add wind arrow
        # Set at 80% of axis limits
        axes[plotCount].arrow(axes[plotCount].get_xlim()[0] + 0.8 * xLimDist,
                  axes[plotCount].get_ylim()[0] + 0.8 * yLimDist,
                  xLimDist * 0.03 * math.sin(math.radians(runLog['Wind Dir'][i]+180)),
                  xLimDist * 0.03 * math.cos(math.radians(runLog['Wind Dir'][i]+180)),
                  width=xLimDist * 0.015,
                  color='k')

        plotCount = plotCount + 1
        windCount = windCount + 1
        plt.show()

        fig.savefig(
            'C:/Users/arey/files/Projects/LSU/Modelling/Figures/Wide_' +
            runLog['Short Description'][i] + '.png', bbox_inches='tight', dpi=400)


# %% Setup forebay polygon
# forebayPoly = gp.read_file('C:/Users/arey/files/Projects/LSU/ForebayPoly.shp')
forebayPoly = gp.read_file('C:/Users/arey/files/Projects/LSU/ForebayPolyExpand.shp')

# Setup mask. Can only do this once if grids are the same
vol_FB = [None] * (max(modelPlot)+1)
i = 64
forebayPolyMask = [forebayPoly.iloc[0, 1].contains(Point(i[0], i[1]))
                   for i in np.array([facex[i], facey[i]]).T]

# Forebay volume
vol_FB[44] = np.sum(wd[i][0, forebayPolyMask] \
                * area[i][forebayPolyMask])

cityParkPoly = gp.read_file('C:/Users/arey/files/Projects/LSU/CityParkMaskNorth.shp')
cityParkMask = [cityParkPoly.iloc[0, 1].contains(Point(i[0], i[1]))
                   for i in np.array([facex[i], facey[i]]).T]

# %% Read in hist file
#
modelPlot = [44, 45]

cumSedCS = [None] * (max(modelPlot)+1)
cumSedCS_B = [None] * (max(modelPlot)+1)
SedCS_B = [None] * (max(modelPlot)+1)
cs_Names = [None] * (max(modelPlot)+1)
cumDis = [None] * (max(modelPlot)+1)
cumDisVol = [None] * (max(modelPlot)+1)
cumVol = [None] * (max(modelPlot)+1)
disVol = [None] * (max(modelPlot)+1)

for i in modelPlot:
    file_nc_hist = os.path.join(runLog['Run Location'][i], 'FlowFM', 'dflowfm', 'output', 'FlowFM_his.nc')
    # file_nc_hist = os.path.join(runLog['Run Location'][i], 'FlowFM', 'output', 'FlowFM_his.nc')

    # Get Var info
    vars_pd, dims_pd = get_ncvardimlist(file_nc=file_nc_hist)

    # Get station list
    station_list = get_hisstationlist(file_nc=file_nc_hist, varname='cross_section_discharge')

    # cumSedCS[i] = get_ncmodeldata(file_nc=file_nc_hist, varname='cross_section_cumulative_Sediment',
    #                               timestep='all', station='all')
    # SedCS_B[i] = get_ncmodeldata(file_nc=file_nc_hist, varname='cross_section_suspended_sediment_transport',
    #                               timestep='all', station='all')
    # cumSedCS_B[i] = SedCS_B[i].cumsum(axis=1)
    #
    # cs_Names[i] = get_ncmodeldata(file_nc=file_nc_hist, varname='cross_section_name',
    #                             station='all')
    # cumDis[i] = get_ncmodeldata(file_nc=file_nc_hist, varname='source_sink_prescribed_discharge',
    #                               timestep='all', station='all')

    disVol[i] = get_ncmodeldata(file_nc=file_nc_hist, varname='cross_section_discharge',
                                  timestep='all', station='all')

    cumVol[i] = get_ncmodeldata(file_nc=file_nc_hist, varname='cross_section_cumulative_discharge',
                                  timestep='all', station='all')

    # cumDisVol[i] = cumDis[i] * 5 * 60 * 129.6
    # cumDisVol[i] = cumDisVol[i].sum()

    # cumDisVol[i] = cumDis[i] * 5 * 60 * 757.2
    # cumDisVol[i] = cumDisVol[i].sum()

    # Get Time info
    tStepsHist = get_timesfromnc(file_nc=file_nc_hist, varname='time')

# %% Calculate 24h exchange
modelPlot = [44, 45]

fig = plt.figure(figsize=(10, 5))
dis24hRolling = []
for i in modelPlot:
    disVol_pd = pd.DataFrame(disVol[i], columns=station_list['cross_section_name'])
    disVol_pd['Time'] = tStepsHist
    disVol_pd.set_index('Time', inplace=True)

    # Multiple by 300 to go from m3/s to m3/ 5 minutes, then take the 24h rolling sum
    dis24hRolling.append(disVol_pd['Forebay1'].multiply(300).abs().rolling(window='24h').sum())

    # Calculate average exchange rate
    avgExchage = dis24hRolling[-1].divide(vol_FB[i]).multiply(100).mean()
    dis24hRolling[-1].divide(vol_FB[i]).multiply(100).plot(
        label=runLog['Short Description'][i] + '. Average Exchange: ' + str(round(avgExchage,2)) + '%')

# plt.ylim(0, 100)
plt.ylim(0, 2)

plt.ylabel('Percent of Forebay Volume Exchanged over 24h')
plt.legend()
plt.show()
# fig.savefig(
#     'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
#     'ForebayExchange_LT.png', bbox_inches='tight', dpi=400)

fig.savefig(
    'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
    'ForebayExchange_LT_0-2.png', bbox_inches='tight', dpi=400)

# Scatterplot

fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6, 6))

ax.scatter(dis24hRolling[1].divide(vol_FB[45]).multiply(100), dis24hRolling[0].divide(vol_FB[44]).multiply(100), s=10, c='r')
ax.set_xlabel('Percent of Forebay Volume Exchanged over 24h: Existing')
ax.set_ylabel('Percent of Forebay Volume Exchanged over 24h: Concept')
ax.plot([0, 100], [0, 100], 'k--')

plt.show()
fig.savefig(
    'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
    'ForebayExchange_Scatter.png', bbox_inches='tight', dpi=400)

#%% Plot Sediment  DFM functions
# modelPlot  = [30, 31, 32, 33, 34, 35]
#modelPlot  = [30, 35]
# modelPlot  = [35]
# stormDirs  = [306, 135, 250]
modelPlot = [64]
sedIDX = 2
# modelPlot = [65]
# sedIDX = 1


plotCount = 0

# fig, axes = plt.subplots(nrows=1, ncols=3,subplot_kw={'projection': ccrs.epsg(26915)}, figsize=(13, 6))
# fig, axes = plt.subplots(nrows=1, ncols=3,
#                          subplot_kw={'projection': ccrs.epsg(26915)},
#                          figsize=(6, 6))
# fig.patch.set_facecolor('white')
# plt.tight_layout()



Trapping_List = [None] * (max(modelPlot)+1)


for tPlot in range(0, 1):
    for i in modelPlot:
        # Plot as seperate figures
        fig, axes = plt.subplots(nrows=1, ncols=1,
                                 subplot_kw={'projection': ccrs.epsg(26915)},
                                 figsize=(6, 9.5))
        fig.patch.set_facecolor('white')

        vel_mask = meshMask[i]

        pc = plot_netmapdata(ugrid_all[i].verts[vel_mask],
                             values=sed[i][tPlot, :, sedIDX][vel_mask],
                             ax=axes, linewidth=0.5,
                             cmap="viridis")
        # pc = plot_netmapdata(ugrid_all[i].verts[vel_mask],
        #                      values=area[i][:][vel_mask],
        #                      ax=axes, linewidth=0.5,
        #                      cmap="viridis")

        # axes.set_xlim(675524, 677027)
        # axes.set_ylim(3365191, 3367182)
        # axes.set_xlim(676470, 677021)
        # axes.set_ylim(3365939, 3366617)
        axes.set_xlim(675524, 677027)
        axes.set_ylim(3365191, 3368135)

        mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
        ctx.add_basemap(axes, source=mapbox, crs='EPSG:26915')

        # extent = ax[plotCount].get_extent()
        # ax[plotCount].set_xticks(np.linspace(extent[0], extent[1], 4))
        # ax[plotCount].set_yticks(np.linspace(extent[2], extent[3], 4))


        # Add model bounds line
        # alphaPoly = alphaPolyList[i]
        # axes.plot(*alphaPoly.exterior.xy, color='k', linewidth=1)

        # Add forebay bounds line
        # forebayPoly.plot(ax=axes, linewidth=2,
        #                  facecolor='none', edgecolor='r')

        # Add City Park bounds line
        cityParkPoly.plot(ax=axes, linewidth=2,
                         facecolor='none', edgecolor='r')

        if i in [44, 45, 36, 37]:
            pc.set_clim([0, 2000])
        elif i < 15:
            pc.set_clim([0, 100])
        else:
            pc.set_clim([0, 100])

        cbar = fig.colorbar(pc, ax=axes, shrink=0.95)
        cbar.set_label('Available mass of sediment [$kg/m^2$]')

        # Calculate sediment in polygon
        massSed = sed[i][0, :, sedIDX] * area[i]
        massSedSum = np.sum(massSed)


        # massSed_FB = sed[i][0, forebayPolyMask, 0]\
        #              * area[i][forebayPolyMask]
        # massSedSum_FB = np.sum(massSed_FB)



        massSed_CP = sed[i][0, cityParkMask, sedIDX]\
                     * area[i][cityParkMask]
        massSedSum_CP = np.sum(massSed_CP)

        # Add title
        # axes.set_title('Existing Condition: ' + "{:0.1f}".format(
        #     (massSedSum_FB / massSedSum) * 100) + '% Trapping')
        # axes.set_title(runLog['Short Description'][i] + ': ' + "{:0.1f}".format(
        #     (massSedSum_FB / massSedSum) * 100) + '% Trapping')


        # plt.title(runLog['Short Description'][i] + ": {:0.1f}".format(
        #     (massSedSum_FB / cumDisVol[i]) * 100) + '% Trapping')


        plt.tight_layout()

        # Trapping_List[i] = (massSedSum_FB / cumDisVol[i])
        Trapping_List[i] = (massSedSum_CP / cumDisTimVol)
        plt.title(runLog['Short Description'][i] + ": {:0.1f}".format(
            (Trapping_List[i] * 100)) + '% Trapping')


        # # Add rectangle to show storm arrow
        # xLimDist = axes.get_xlim()[1] - axes.get_xlim()[0]
        # yLimDist = axes.get_ylim()[1] - axes.get_ylim()[0]
        #
        # axes.add_patch(patches.Rectangle((axes.get_xlim()[0] + 0.7 * xLimDist,
        #                 axes.get_ylim()[0] + 0.8 * yLimDist - xLimDist * 0.1),
        #                 xLimDist * 0.2, xLimDist * 0.2,
        #                 facecolor='white', edgecolor='w', linewidth=5))
        #
        #
        # # Add wind arrow
        # # Set at 80% of axis limits
        # axes.arrow(axes.get_xlim()[0] + 0.8 * xLimDist,
        #           axes.get_ylim()[0] + 0.8 * yLimDist,
        #           xLimDist * 0.03 * math.sin(math.radians(stormDirs[tPlot])),
        #           xLimDist * 0.03 * math.cos(math.radians(stormDirs[tPlot])),
        #           width=xLimDist * 0.015,
        #           color='k')

        plotCount = plotCount + 1
        plt.show()
        fig.savefig(
             'C:/Users/arey/files/Projects/LSU/Modelling/FiguresExpand/' +
            runLog['Run Name'][i] + '_RevB.png', bbox_inches='tight', dpi=400)

# plt.show()

# fig.savefig(
#     'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
#     'SedZoomPercent.png', bbox_inches='tight', dpi=400)

# %% Create summary table
modelPlot = range(30, 46)
trappingSumTable = pd.DataFrame()

trappingSumTable['Description'] = runLog['Short Description'][modelPlot]
trappingSumTable['Trapping'] = np.array(Trapping_List)[modelPlot]

# %% Calculate the total sediment discharged City Park Lake
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6,6))
modelPlot = [64]


sourceSinkTimDat = []
for i in modelPlot:
    inputFiles = [f for f in listdir(os.path.join(runLog['Run Location'][i], 'FlowFM', 'input')) if
                  isfile(join(os.path.join(runLog['Run Location'][i], 'FlowFM', 'input'), f))]

    for inputFile in inputFiles:

        if '.pli' in inputFile and 'source' in inputFile:
            # Read pli to find location
            sourceLocation_df = pd.read_csv(os.path.join(runLog['Run Location'][i], 'FlowFM', 'input', inputFile),
                                            skiprows=1, delim_whitespace=True)
            sourcePoint = Point(sourceLocation_df.iloc[0, 0], sourceLocation_df.iloc[0, 1])

            # If point in City Park Area
            if sourcePoint.within(cityParkPoly.iloc[0, 1]):
                # Read in the sediment discharge file (.tim)
                sourceSinkTimFile = os.path.join(runLog['Run Location'][i], 'FlowFM', 'input',
                                                         inputFile[0:-4] + '.tim')
                # sourceSinkTimDat.append(
                #     pd.read_csv(sourceSinkTimFile, sep='\s+', header=None, names=['time', 'Q', 'conc_sed', 'conc_fluff', 'conc_sec01']))
                sourceSinkTimDat.append(
                    pd.read_csv(sourceSinkTimFile, sep='\s+', header=None, names=['time', 'Q', 'conc_sed', 'conc_sec01']))

                # Add datetimes
                sourceSinkTimDat[-1].index = pd.to_datetime(sourceSinkTimDat[-1].time, unit='m', origin=pd.Timestamp('2005-01-01'))

                # Vol sediment at each time step, multiplied by the time step
                sourceSinkTimDat[-1]['SedVol'] = sourceSinkTimDat[-1]['Q'] *\
                                              sourceSinkTimDat[-1].loc[:, 'time'].diff() * 60 *\
                                              sourceSinkTimDat[-1]['conc_sec01']

                # Fix first time step
                sourceSinkTimDat[-1].iloc[0, -1] = 0

                # Sediment cumlative sum
                sourceSinkTimDat[-1].loc[:, 'CumSed'] = sourceSinkTimDat[-1].loc[:, 'SedVol'].cumsum()

                # sourceSinkTimDat[-1].loc[:, 'Q'].rolling(24*4).sum().plot()
                sourceSinkTimDat[-1].loc[:, 'Q'].multiply(
                    sourceSinkTimDat[-1].loc[:, 'conc_sec01']).multiply(
                    sourceSinkTimDat[-1].loc[:, 'time'].diff()).multiply(
                    60).cumsum().plot()

plt.show()

# %% Sum up city park lake

cumDisTimVol = 0

for s in range(0, len(sourceSinkTimDat)):
    cumDisTimVol = cumDisTimVol + sourceSinkTimDat[s].CumSed[-1]


# %% Calculate the total sediment discharged
# 
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6,6))
modelPlot = [58]

sourceSinkDat = []
for i in modelPlot:
    for sourceSink in range(1, 5):

        # Read in the sediment discharge file (.tim)
        sourceSinkFile = os.path.join(runLog['Run Location'][i], 'FlowFM', 'input', 'SourceSink0' + str(sourceSink) + '.tim')
        sourceSinkDat.append(pd.read_csv(sourceSinkFile, sep='\s+', header=None, names=['time', 'Q', 'conc_sed', 'conc_fluff', 'conc_sed01']))

        # Add datetimes
        sourceSinkDat[-1].index = pd.to_datetime(sourceSinkDat[-1].time, unit='m', origin=pd.Timestamp('2020-01-01'))

        # Vol sediment at each time step, multiplied by the time step
        sourceSinkDat[-1]['SedVol'] = sourceSinkDat[-1]['Q'] *\
                                      sourceSinkDat[-1].loc[:, 'time'].diff() * 60 *\
                                      sourceSinkDat[-1]['conc_sed01']

        # Fix first time step
        sourceSinkDat[-1].iloc[0, -1] = 0

        # Sediment cumlative sum
        sourceSinkDat[-1].loc[:, 'CumSed'] = sourceSinkDat[-1].loc[:, 'SedVol'].cumsum()

        # sourceSinkDat[-1].loc[:, 'Q'].rolling(24*4).sum().plot()
        sourceSinkDat[-1].loc[:, 'Q'].sum().plot()


plt.show()

# %% Read in 15 minute precipitation data
prep_data = pd.read_excel('//srv-mad3/Projects/13522.101 LSU Lakes/06_Models/PCSWMM Model/DataforHistoricRuns/15min_precip_formatted.xlsx',
                          sheet_name='Sheet2')

prep_data['DateTime'] = pd.to_datetime(prep_data['DT'])
prep_data.set_index('DateTime', inplace=True)

# %% Plot 24h precipitationto find events
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(12, 6))

prep_data.loc[:, 'per 15min'].rolling(24 * 4).sum().plot()
axes.set_ylabel('24h Precipiation (in)')

axes.set_xlim((pd.to_datetime('20050701'), pd.to_datetime('20060101')))

plt.show()

# fig.savefig(('//srv-mad3/Projects/13522.101 LSU Lakes/06_Models/PCSWMM Model/DataforHistoricRuns/24h_PrecpSum.png')

# %% Plot the sediment discharge and cumlative transport

fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6,6))

csIDX = 3

# Calculate the total sediment discharged
# SS 1/2 is SouthWest, 3/4 is SouthEast
sourceSinkDat[0].loc[:, 'CumSed'].multiply(2).plot(
        label='Discharged Sediment SW') # in kg
sourceSinkDat[2].loc[:, 'CumSed'].multiply(2).plot(
        label='Discharged Sediment NE') # in kg

totalSedDis = sourceSinkDat[0].loc[:, 'CumSed'].multiply(2) +\
              sourceSinkDat[2].loc[:, 'CumSed'].multiply(2)

# Plot the total sediment discharge
totalSedDis.plot(label='Total Discharged Sediment')

axes.plot(tStepsHist, cumSedCS[i][:, csIDX]*2650,
          label='Cross Section Sediment Transport', color='k') # In M3. Multiple by 2650 to get kg
# axes.plot(tStepsHist, cumSedCS_B[i][:, csIDX],
#           label='Cross Section Sediment Transport B', color='r') # In M3. Multiple by 2650 to get kg

axes.set_xlabel('Date')
axes.set_ylabel('Sediment [kg]')

trappedSed = (1 - (cumSedCS[i][-1, csIDX]*2650) /
             totalSedDis[-1])*100
trappedSedB = (1 - (cumSedCS_B[i][-1, csIDX]) /
             totalSedDis[-1])*100

axes.set_title('Sediment Transport: ' + "{:0.1f}".format(trappedSed) + '% Trapping')

plt.legend()
plt.show()

# fig.savefig(
#     'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
#     'SedTrap_Concept_NE.png', bbox_inches='tight', dpi=400)
# fig.savefig(
#     'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
#     'SedTrap_Concept.png', bbox_inches='tight', dpi=400)


#%% Plot Circulation  DFM functions
modelPlot = range(53, 55)

plotCount = 1
plotLayer = 0

stormDirs = [157.5, 157.5]

for tPlot in range(0, 1):
    for i in modelPlot:
        # Plot as seperate figures
        fig, axes = plt.subplots(nrows=1, ncols=1,
                                 subplot_kw={'projection': ccrs.epsg(26915)},
                                 figsize=(6, 6))
        fig.patch.set_facecolor('white')

        vel_mask = meshMask[i]

        pc = plot_netmapdata(ugrid_all[i].verts[vel_mask],
                             values=np.sqrt(ux[i][0, vel_mask, plotLayer] ** 2 + uy[i][0, vel_mask, plotLayer] ** 2),
                             ax=axes, linewidth=0.5,
                             cmap="viridis")

        # axes.set_xlim(675524, 677027)
        # axes.set_ylim(3365191, 3367182)
        axes.set_xlim(676470, 677021)
        axes.set_ylim(3365939, 3366617)

        mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
        ctx.add_basemap(axes, source=mapbox, crs='EPSG:26915')

        # extent = ax[plotCount].get_extent()
        # ax[plotCount].set_xticks(np.linspace(extent[0], extent[1], 4))
        # ax[plotCount].set_yticks(np.linspace(extent[2], extent[3], 4))

        # Add quiver arrows
        U_plot = ux[i][0, vel_mask, plotLayer] / np.sqrt(ux[i][0, vel_mask, plotLayer] ** 2 + uy[i][0, vel_mask, plotLayer] ** 2)
        V_plot = uy[i][0, vel_mask, plotLayer] / np.sqrt(ux[i][0, vel_mask, plotLayer] ** 2 + uy[i][0, vel_mask, plotLayer] ** 2)

        Q = plt.quiver(facex[i][vel_mask], facey[i][vel_mask], U_plot, V_plot, scale=55)
        # Add quiverkey
        qk = plt.quiverkey(Q, 0.8, 0.88, 2, r'$2 \frac{m}{s}$', labelpos='E',
                           coordinates='figure')

        # Add model bounds line
        alphaPoly = alphaPolyList[i]
        axes.plot(*alphaPoly.exterior.xy, color='k', linewidth=1)


        pc.set_clim([0, 0.1])

        cbar = fig.colorbar(pc, ax=axes, shrink=0.95)
        cbar.set_label('Surface Velocity [m/s]')

        # Add title
        plt.title(runLog['Short Description'][i])

        plt.tight_layout()

        # # Add rectangle to show storm arrow
        xLimDist = axes.get_xlim()[1] - axes.get_xlim()[0]
        yLimDist = axes.get_ylim()[1] - axes.get_ylim()[0]

        axes.add_patch(patches.Rectangle((axes.get_xlim()[0] + 0.7 * xLimDist,
                        axes.get_ylim()[0] + 0.8 * yLimDist - xLimDist * 0.1),
                        xLimDist * 0.2, xLimDist * 0.2,
                        facecolor='white', edgecolor='w', linewidth=5))


        # Add wind arrow
        # Set at 80% of axis limits
        axes.arrow(axes.get_xlim()[0] + 0.8 * xLimDist,
                  axes.get_ylim()[0] + 0.8 * yLimDist,
                  xLimDist * 0.03 * math.sin(math.radians(stormDirs[tPlot])),
                  xLimDist * 0.03 * math.cos(math.radians(stormDirs[tPlot])),
                  width=xLimDist * 0.015,
                  color='k')

        plotCount = plotCount + 1
        plt.show()
        fig.savefig(
             'C:/Users/arey/files/Projects/LSU/Modelling/FiguresExpand/Circ_Fig' +
            runLog['Run Name'][i] + '.png', bbox_inches='tight', dpi=400)

# plt.show()

# fig.savefig(
#     'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
#     'SedZoomPercent.png', bbox_inches='tight', dpi=400)


