#%% 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

#%% 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(30, 46)

# 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)
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)
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')

    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               = 50

    # 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')
    # 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)


#%% Plot using DFM functions
modelPlot  = [24]
# stormDirs  = [306, 135, 250]
stormDirs  = [3, 346, 342]

# fig, axes = plt.subplots(nrows=1, ncols=3,subplot_kw={'projection': ccrs.epsg(26915)}, figsize=(13, 6))
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

vmin        = 0 # Legend min
vmax        = 0.05 # Legend max
scale       = 60 # Size of arrows, smaller is bigger #40 # 15
qmax        = 0.025 # 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(1, 2):

    for i in modelPlot:
        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")

        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 = U[i]
        overGrid_V_Scale = V[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[i]] = np.nan
        overGrid_V_Plot[~regularMask[i]] = np.nan

        qv = axes[plotCount].quiver(X[i], Y[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('Depth Averaged Velocity [m/s]')

        axes[plotCount].set_title('Wind Direction: ' + str(stormDirs[tPlot]) + ' 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(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/Figures/' +
    'StormVelDisWide.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
i = 30
forebayPolyMask = [forebayPoly.iloc[0, 1].contains(Point(i[0], i[1]))
                   for i in np.array([facex[i], facey[i]]).T]


# %% Read in hist file
#
modelPlot = range(30, 46)

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)


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

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

    # 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')

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

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


#%% Plot Sediment  DFM functions
# modelPlot  = [30, 31, 32, 33, 34, 35]
#modelPlot  = [30, 35]
# modelPlot  = [35]
# stormDirs  = [306, 135, 250]
modelPlot = range(30, 46)

# 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()

plotCount = 1
sedIDX = 0
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, 6))
        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[plotCount].set_xlim(675524, 677027)
        # axes[plotCount].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 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')

        if i in [44, 45, 36, 37]:
            pc.set_clim([0, 2000])
        else:
            pc.set_clim([0, 200])

        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, :, 0] * area[i]
        massSedSum = np.sum(massSed)


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

        # 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])

        # # 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] + '.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
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6,6))
modelPlot = [27]

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']))

        # Add datetimes
        sourceSinkDat[-1].index = pd.to_datetime(sourceSinkDat[-1].time, unit='m', origin=pd.Timestamp('2005-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_sed']

        # 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()

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)
