## Plotting Mike21 SW results for SouthShore
# Author: AJMR
# December 22, 2021

# %%  Setup Project
from mikeio import Dfsu, Mesh, Dfs2, Dfs0

import pandas as pd
import pathlib as pl
import numpy as np
import geopandas as gp

import matplotlib.pyplot as plt
import contextily as ctx
import matplotlib.font_manager as fm

# %% Read Model Log

pth = pl.Path("//srv-mad3.baird.com/", "Projects", "13632.101 South Shore Breakwater", "06_Models",
              "13632.101.W.AJMR.Rev0_ModelLog.xlsx")
runLog = pd.read_excel(pth.as_posix(), "HD")

# %% Read Model Results
modelPlot = range(6, 7)

curSpeed_list   = [None] * (max(modelPlot) + 1)
curElm_list     = [None] * (max(modelPlot) + 1)
curU_list       = [None] * (max(modelPlot) + 1)
curV_list       = [None] * (max(modelPlot) + 1)
dfs_list        = [None] * (max(modelPlot) + 1)
ds_list         = [None] * (max(modelPlot) + 1)
wl_list         = [None] * (max(modelPlot) + 1)
wd_list         = [None] * (max(modelPlot) + 1)
bed_list        = [None] * (max(modelPlot) + 1)

for m in modelPlot:
    dfsIN       = Dfsu(pl.Path(runLog['Run Location'][m],
                         runLog['Run Name'][m] + '.m21fm - Result Files', 'FullDomain.dfsu').as_posix())

    dfs_list[m]      = dfsIN
    ds_list[m]       = dfs_list[m].read()
    curSpeed_list[m] = ds_list[m]["Current speed"][-1]
    curU_list[m]     = ds_list[m]["U velocity"][-1]
    curV_list[m]     = ds_list[m]["V velocity"][-1]
    curElm_list[m]   = dfs_list[m].element_coordinates
    wl_list[m]       = ds_list[m]["Surface elevation"][-1]
    wd_list[m]       = ds_list[m]["Total water depth"][-1]
    bed_list[m]      = wl_list[m] - wd_list[m]
    # # Read Mesh for first model
    # mikeMeshIN = dfsIN.to_shapely()
    # # Convert to list of polygons
    # mikeMeshIN_list = [p for p in mikeMeshIN]

    # Add to list
    # hmo_list[m] = gp.GeoDataFrame(hm0, geometry=mikeMeshIN_list)

    # Cleanup unnecessary variables
    # del mikeMeshIN
    # del mikeMeshIN_list
    del dfsIN

# %% Plotting
# Shaded Water
mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
x, y, arrow_length = 0.93, 0.95, 0.12
fontprops = fm.FontProperties(size=12)
modelPlot = range(6, 7)

vmin    = 0
vmax    = 1
scale   = 1 #5

for m in modelPlot:
    if m < 4 or m == 6:
        vmax  = 5
        qmax  = 4
        scale = 50
        amin  = 1
    else:
        vmax  = 1
        qmax  = 0.25
        scale = 75
        amin  = 3

    fig, axes = plt.subplots(figsize=(8, 8))

    # Convert to feet and ignore missing
    cur_plot    = curSpeed_list[m]*3.28084
    cur_nan     = np.isnan(cur_plot)

    axDHI = dfs_list[m].plot(cur_plot[~cur_nan], plot_type='contourf', show_mesh=False, cmap='viridis', ax=axes, levels=11,
               vmin=vmin, vmax=vmax, label='Current Speed (ft/s)',
                     elements=dfs_list[m].element_ids[~cur_nan])

    # Plot arrows on overset grid
    overGrid     = dfs_list[m].get_overset_grid(dx=50)
    interpolant  = dfs_list[m].get_2d_interpolant(overGrid.xy, n_nearest=1)

    # Interpolate velocity at last time step by selecting time from dataset
    overGrid_Cur = dfs_list[m].interp2d(ds_list[m][ds_list[m].time[-1]:], *interpolant, shape=(overGrid.ny, overGrid.nx))

    overGrid_U = overGrid_Cur["U velocity"][-1]
    overGrid_V = overGrid_Cur["V velocity"][-1]

    # Normalize velocity above limit
    overGrid_U_Scale = overGrid_U*3.28084
    overGrid_V_Scale = overGrid_V*3.28084

    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
    overGrid_U_Plot = overGrid_U_Scale / r
    overGrid_V_Plot = overGrid_V_Scale / r

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

    qv = axDHI.quiver(overGrid.x, overGrid.y, 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

    axes.set_xlim(left=427800, right=431000)
    axes.set_ylim(bottom=4758000, top=4762000)

    ctx.add_basemap(axes, source=mapbox, crs='EPSG:32616')
    # axes.title.set_text(runLog['Short Description'][m])
    # axes.titlesize = 'x-large'
    fig.suptitle(runLog['Short Description'][m], fontsize=18)
    axes.set_xlabel('Easting [m]')
    axes.set_ylabel('Northing [m]')

    plt.show()

    fig.savefig(
        '//srv-mad3.baird.com/Projects/13632.101 South Shore Breakwater/06_Models/04_MIKE21_HD/Images/Production/' +
        runLog['Short Description'][m] + '_RevC_Current.png', bbox_inches='tight', dpi=300)

# %% Plot Bathymetry

fig, axes = plt.subplots(figsize=(8, 8))
# Convert to feet and ignore missing
bed_plot = (bed_list[m] * 3.28084) - 578.007
bed_nan = np.isnan(bed_plot)

axDHI = dfs_list[m].plot(bed_plot[~bed_nan], plot_type='contourf', show_mesh=False, cmap='magma_r', ax=axes, levels=9,
                         vmin=-80, vmax=0, label='Bed Elevation (ft)',
                         elements=dfs_list[m].element_ids[~bed_nan])

axes.set_xlim(left=427800, right=431000)
axes.set_ylim(bottom=4758000, top=4762500)

ctx.add_basemap(axes, source=mapbox, crs='EPSG:32616')

fig.suptitle('Bed Elevation', fontsize=18)
axes.set_xlabel('Easting [m]')
axes.set_ylabel('Northing [m]')

plt.show()

fig.savefig(
    '//srv-mad3/Projects/13632.101 South Shore Breakwater/10_Reports&Pres/13632.101.R3 Ph I BOD/Support files/' +
    'Bed_Elevation.png', bbox_inches='tight', dpi=300)

# %% Read time series
MIKEds_list  = [None] * (max(modelPlot) + 1)
MIKEdsT_list = [None] * (max(modelPlot) + 1)

for m in modelPlot:
    dfsIN = Dfs0(pl.Path(runLog['Run Location'][m], str(runLog['Number'][m]) + '_' +
                         runLog['Run Name'][m] + '.sw - Result Files', 'BreakPts.dfs0').as_posix())
    dfsIN_read = dfsIN.read()
    MIKEds_list[m] = dfsIN_read.to_dataframe()

    dfsTIN = Dfs0(pl.Path(runLog['Run Location'][m], str(runLog['Number'][m]) + '_' +
                          runLog['Run Name'][m] + '.sw - Result Files', 'TransectPTS.dfs0').as_posix())
    dfsTIN_read = dfsTIN.read()
    MIKEdsT_list[m] = dfsTIN_read.to_dataframe()

    # Cleanup unnecessary variables
    del dfsIN_read
    del dfsIN
    del dfsTIN_read
    del dfsTIN

# %% Read in Toe and Crest Shapefiles
breakwaterPTS = gp.read_file("//srv-mad3.baird.com/Projects/"
                             "13632.101 South Shore Breakwater/08_CADD/Outgoing/"
                             "20211211_Toe Extents (to Alexander)/"
                             "20211211_Toe_Extents_NAD83_WISStatePlaneSZn_USFt_Lines_OffshoreClipSimple_10m_vertexUTM.shp")
breakCrest = pd.read_csv(
    '//srv-mad3.baird.com/Projects/13632.101 South Shore Breakwater/08_CADD/Outgoing/20211214_Crest Points (to Alexander)/20211214_Crest_Points_m.csv',
    header=None, names=['x', 'y', 'z'])
breakCrest_gdf = gp.GeoDataFrame(breakCrest, crs='EPSG:32154',
                                 geometry=gp.points_from_xy(breakCrest.x, breakCrest.y))
breakCrest_gdf.to_crs('EPSG:32616', inplace=True)

breakwaterPTS_Crest = breakwaterPTS.sjoin_nearest(breakCrest_gdf)
breakwaterPTS_Crest.rename(columns={'x': 'breakCrest_x', 'y': 'breakCrest_y', 'z': 'Crest'}, inplace=True)

# %% Merge with data
breakPointsOut = [None] * (max(modelPlot) + 1)
breakTimesOut  = [None] * (max(modelPlot) + 1)

for m in modelPlot:
    breakwaterPTS_times = None

    for t in range(0, MIKEds_list[m].shape[0]):
        breakwaterPTS_merge = None
        breakwaterPTS_merge = breakwaterPTS_Crest

        for i in range(0, MIKEds_list[m].shape[1]):
            paramName = MIKEds_list[m].columns.values[i][MIKEds_list[m].columns.values[i].find('"', 5, -1) + 3:]
            if paramName not in breakwaterPTS_merge.columns:
                breakwaterPTS_merge[paramName] = np.full([breakwaterPTS_merge.shape[0], 1], np.nan)
            tmpFID = int(MIKEds_list[m].columns.values[i][1:MIKEds_list[m].columns.values[i].find('"', 1)])
            tmpIND = int(MIKEds_list[m].columns.values[i][5:MIKEds_list[m].columns.values[i].find('"', 5)])

            breakwaterPTS_merge.loc[((breakwaterPTS_merge.FID == tmpFID) &
                                     (breakwaterPTS_merge.vertex_ind == tmpIND)),
                                    paramName] = MIKEds_list[m].iloc[t, i]

        breakwaterPTS_merge['Time'] = MIKEds_list[m].index[t]

        breakwaterPTS_times = pd.concat([breakwaterPTS_times, breakwaterPTS_merge], ignore_index=True)

    breakTimesOut[m]  = breakwaterPTS_times

# %% Read in Breakwater Transect Points Sh
breakwaterT     = pd.read_csv("C:/Users/arey/files/Projects/SouthShore/Bathy/BreakTransectPTS_Names.csv", sep='\t')

breakwaterT_PTS = gp.GeoDataFrame(breakwaterT, geometry=gp.points_from_xy(breakwaterT.X, breakwaterT.Y),
                                  crs='EPSG:32616')

# %% Merge with data
breakPoints_TOut = [None] * (max(modelPlot) + 1)
breakTimes_TOut  = [None] * (max(modelPlot) + 1)

for m in modelPlot:
    breakwaterPTS_times = None

    for t in range(0, MIKEdsT_list[m].shape[0]):
        breakwaterPTS_merge = None
        breakwaterPTS_merge = breakwaterT_PTS

        for i in range(0, MIKEdsT_list[m].shape[1]):
            paramName = MIKEdsT_list[m].columns.values[i][MIKEdsT_list[m].columns.values[i].find(':', 5, -1) + 2:]
            if paramName not in breakwaterPTS_merge.columns:
                breakwaterPTS_merge[paramName] = np.full([breakwaterPTS_merge.shape[0], 1], np.nan)
            tmpName = MIKEdsT_list[m].columns.values[i][0:MIKEdsT_list[m].columns.values[i].find(':', 1)]

            breakwaterPTS_merge.loc[(breakwaterPTS_merge.Name == tmpName),
                                    paramName] = MIKEdsT_list[m].iloc[t, i]

        breakwaterPTS_merge['Time'] = MIKEdsT_list[m].index[t]

        breakwaterPTS_times = pd.concat([breakwaterPTS_times, breakwaterPTS_merge], ignore_index=True)

    breakTimes_TOut[m]  = breakwaterPTS_times


#%% Format and save

for m in modelPlot:
    saveTmp         = breakTimesOut[m].copy()
    saveTmp['X']    = saveTmp.geometry.x
    saveTmp['Y']    = saveTmp.geometry.y
    saveTmp.drop(['vertex_par', 'vertex_p_1', 'angle', 'geometry', 'index_right',
                  'breakCrest_x', 'breakCrest_y', 'Length'], axis=1, inplace=True)
    saveTmp.sort_values(by=['FID', 'vertex_ind'], inplace=True, ignore_index=True)

    # Reorder columns
    colNames    = saveTmp.columns.values
    saveTmp     = saveTmp[['X', 'Y', *colNames[0:-2]]]
    saveTmpT    = saveTmp.transpose(copy=True)

    # Transect points
    saveTmp2         = breakTimes_TOut[m].copy()
    saveTmp2.drop(['geometry'], axis=1, inplace=True)
    saveTmp2T       = saveTmp2.transpose(copy=True)


    saveTmpT.to_csv('//srv-mad3.baird.com/Projects/13632.101 South Shore Breakwater/06_Models/02_Mike21SW/Results/' +
                    'Toe' + runLog['Short Description'][m] + '.csv')

    saveTmp2T.to_csv('//srv-mad3.baird.com/Projects/13632.101 South Shore Breakwater/06_Models/02_Mike21SW/Results/' +
                    'Transect' + runLog['Short Description'][m] + '.csv')