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

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

import pandas as pd
import pathlib as pl
import numpy as np
import geopandas as gp
import datetime as datetime
from pyproj import CRS, Transformer

import matplotlib as mpl
import matplotlib.pyplot as plt
import time
import matplotlib.animation as animation
import matplotlib.cm as cm
from matplotlib.colors import Normalize
import contextily as ctx
import matplotlib.font_manager as fm

import os
import sys
sys.path.append('C:/Users/arey/Repo/Python_Baird/Mustique')
sys.path.append('C:/Users/arey/Repo/Python_Baird/Mustique/adcptool')
from adcploader import *


# %% Read Model Log

pth = pl.Path("//srv-ott3.baird.com/", "Projects", "13509.101 NCC New Edinburgh Ottawa",
              "06_Models", "13509.101.W.AJMR.Rev0_NCC_Boathouse_ModelLog.xlsx")

runLog = pd.read_excel(pth.as_posix(), "ModelLog")

# %% Read in ADCP Transect Points
importFiles = os.listdir('C:/Users/arey/files/Projects/OttawaRiver/MergedSurvey/')
adcpDat  = []
adcpAvg  = []
adcpName = []
# Initialize import variables
for i in range(0, len(importFiles)):
    if 'ASC.TXT' in importFiles[i]:
        OBS_IN = RawProfileObj('C:/Users/arey/files/Projects/OttawaRiver/MergedSurvey/' + importFiles[i])

        # Convert to UTM
        transformer = Transformer.from_crs(4326, 26918)
        for j in range(0, len(OBS_IN.ensembles)):
            [OBS_IN.ensembles[j].UTMx, OBS_IN.ensembles[j].UTMy] =\
                transformer.transform(OBS_IN.ensembles[j].lat, OBS_IN.ensembles[j].lon)

        # Set start and end points for the transect
        startingpoint = dict(start=Vector(OBS_IN.ensembles[0].UTMx, OBS_IN.ensembles[0].UTMy, 0),
                             end=Vector(OBS_IN.ensembles[-1].UTMx, OBS_IN.ensembles[-1].UTMy, 0))

        # Process
        processing_settings = dict(proj_method=3)
        p0 = ProcessedProfileObj(OBS_IN, processing_settings, startingpoint)

        # Save
        adcpDat.append(p0)
        adcpName.append(importFiles[i])
        adcpAvg.append(get_averaged_profile(p0, cfg={'order': 21}))



        # Plot
        # plot_profile_2d(thin_out(p0, {'keep_ensemble': 10}))


# %% Plot ADCP points on a map
mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckex9vtri0o6619p55sl5qiyv/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
# mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckemxgtk51fgp19nybfmdcb1e/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'

fig, axes = plt.subplots(figsize=(8, 8))
colormap = cm.plasma

# Set thinning factor
thin_factor = 10

# Set ADCP transect to plot
for adcp_to_plot in range(0, len(adcpDat)):
    # position of ensembles
    x_all = [e.position.x for e in adcpAvg[adcp_to_plot].ensembles]
    y_all = [e.position.y for e in adcpAvg[adcp_to_plot].ensembles]
    x = x_all[::thin_factor]
    y = y_all[::thin_factor]

    # velocity components. expects uv_rot to be False!!
    vx_all = [e.velocity.x if not e.void else 0 for e in adcpAvg[adcp_to_plot].ensembles]
    vy_all = [e.velocity.y if not e.void else 0 for e in adcpAvg[adcp_to_plot].ensembles]
    vx = vx_all[::thin_factor]
    vy = vy_all[::thin_factor]

    # average length for the quiverkey
    avglen = np.sqrt(np.array(vx) ** 2 + np.array(vy) ** 2).mean()
    displen = np.around(avglen, int(np.ceil(np.abs(np.log10(avglen)))))

    # scale the vectors
    vscale =  1

    axes.plot(x, y, marker='x', color='white')
    axes.scatter(x[0], y[0], marker='o')
    axes.axis('equal')
    axes.grid()

    norm = Normalize(vmin=0, vmax=2)
    colors = np.sqrt(np.array(vx) ** 2 + np.array(vy) ** 2)
    # norm.autoscale(colors)

    q = axes.quiver(x, y, vx * vscale, vy * vscale, width=0.003,
                    color=colormap(norm(colors)))
    # axes.quiverkey(q, 0.2, 0.05, displen * vscale, 'velocity: {} [m/s]'.format(displen))

# axes.set_xlim(left=446747.7, right=447246.7)
# axes.set_ylim(bottom=5033841.5, top=5034305.8)
axes.set_xlim(left=445959.7, right=447219.7)
axes.set_ylim(bottom=5033522.5, top=5034850.8)

ctx.add_basemap(axes, source=mapbox, crs='EPSG:26918')
axes.set_xlabel('Easting [m]')
axes.set_ylabel('Northing [m]')

plt.show()


# %% Read Model Results at point
modelPlot = [5]

ds_list = [None] * (max(modelPlot) + 1)
dsT = [None] * (15)

dfs2d_list = [None] * (max(modelPlot) + 1)
dfs3d_list = [None] * (max(modelPlot) + 1)

z_list = [None] * (max(modelPlot) + 1)
z_profile_list = [None] * (max(modelPlot) + 1)

elm_list = [None] * (max(modelPlot) + 1)
elm_df_list = [None] * (max(modelPlot) + 1)

elm2d_list = [None] * (max(modelPlot) + 1)
elm2d_df_list = [None] * (max(modelPlot) + 1)

# Identify points to read
x_all = []
y_all = []
ADCPidxCount = 0
ADCP_transect_idx = []
for adcp in adcpAvg:
    for e in adcp.ensembles:
        x_all.append(e.position.x)
        y_all.append(e.position.y)

        ADCPidxCount = ADCPidxCount + 1

    ADCP_transect_idx.append(ADCPidxCount)

readPoints = np.array((x_all, y_all)).T

# Drop NaNs
readPoints = readPoints[~np.isnan(readPoints).any(axis=1)]


for m in modelPlot:

    dfsIN = Dfsu(pl.Path(runLog['Run Location'][m], 'NCC_2Dout.dfsu').as_posix())
    dfs2d_list[m] = dfsIN

    # # Read Map
    ds_list[m] = dfs2d_list[m].read()

    # Read in transects
    # dsT[4] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T4.dfsu').as_posix())
    # dsT[5] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T5.dfsu').as_posix())
    # dsT[6] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T6.dfsu').as_posix())
    # dsT[7] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T7.dfsu').as_posix())
    # dsT[8] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T8.dfsu').as_posix(), time=90)
    # dsT[9] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T9.dfsu').as_posix())
    # dsT[10] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T10.dfsu').as_posix())
    # dsT[11] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T11.dfsu').as_posix())
    # dsT[13] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T13.dfsu').as_posix())
    # dsT[14] = mikeio.read(pl.Path(runLog['Run Location'][m], 'Run01_T14.dfsu').as_posix())

    # ## Read specific points in 2D
    # # Find nearest elements
    # elem_ids = dfs2d_list[m].find_nearest_elements(readPoints[:, 0], readPoints[:, 1])
    # # Read in data from nearest elements
    # elm2d_list[m] = dfs2d_list[m].read(elements=elem_ids)
    #
    # # Convert to Pandas DataFrame
    # elm2d_df_list[m] = [None] * readPoints.shape[0]
    # for p in range(0, readPoints.shape[0]):
    #     elm2d_df_list[m][p] = elm2d_list[m].isel(idx=p).to_dataframe()


    # ## Read specific points 3D
    # # Setup MIKE object
    # dfsIN = Dfsu(pl.Path(runLog['Run Location'][m], 'Run01_3DOut.dfsu').as_posix())
    # dfs3d_list[m] = dfsIN
    #
    # # Find nearest elements
    # elem_ids = dfs3d_list[m].find_nearest_profile_elements(readPoints[:, 0], readPoints[:, 1])
    # # Read from elements- flatten 2d element array (xy and z) to 1d for reading
    # elm_list[m] = dfs3d_list[m].read(elements=elem_ids.flatten().astype(int))
    #
    # elm_df_list[m] = [None] * readPoints.shape[0]
    # z_profile_list[m] = [None] * readPoints.shape[0]
    # # Convert to Pandas DataFrame
    # # Loop through points, selecting corresponding depths
    # for p in range(0, readPoints.shape[0]):
    #     z_profile_list[m][p] = dfs3d_list[m].element_coordinates[elem_ids[p, :].astype(int), 2]
    #
    #     # Convert to Pandas and add zidx
    #     for z in range(0, z_profile_list[m][p].shape[0]):
    #         df_tmp = elm_list[m].isel(idx=p * 5 + z).to_dataframe()
    #         df_tmp['Z'] = z_profile_list[m][p][z]
    #         df_tmp['Z_IDX'] = z
    #         if z == 0:
    #             elm_df_list[m][p] = df_tmp
    #         else:
    #             elm_df_list[m][p] = elm_df_list[m][p].append(df_tmp)
    #
    #         del df_tmp




# %% Map Plotting
# mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckex9vtri0o6619p55sl5qiyv/tiles/256/{z}/{x}/{y}@2x?access_token=pk.eyJ1IjoiYWxleGFuZGVyMDA0MiIsImEiOiJjazVmdG4zbncwMHY4M2VrcThwZGUzZDFhIn0.w6oDHoo1eCeRlSBpwzwVtw'
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 = [5]
# Set thinning factor
thin_factor = 10

for m in modelPlot:
    vmax = 1
    vmin = 0

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

    # Plot Velocity

    # Filter Out Land Values (no velocity)
    dhiFilt = ds_list[m]['Current speed'].sel(time="2018-01-01 02:30:00").values > 0
    dhiFiltIDX = [i for i, x in enumerate(dhiFilt) if x]

    # Plot Velocity at last time step for elements greater than zero
    axDHI = ds_list[m]['Current speed'].sel(time="2018-01-01 02:30:00").isel(element=dhiFiltIDX).plot.contourf(
                            show_mesh=False, cmap='inferno', ax=axes,
                             vmin=vmin, vmax=vmax, label='Current Speed (m/s)')


    # Add ADCP data
    for adcp_to_plot in range(0, len(adcpDat)):
        # position of ensembles
        x_all = [e.position.x for e in adcpAvg[adcp_to_plot].ensembles]
        y_all = [e.position.y for e in adcpAvg[adcp_to_plot].ensembles]
        x = x_all[::thin_factor]
        y = y_all[::thin_factor]

        # velocity components. expects uv_rot to be False!!
        vx_all = [e.velocity.x if not e.void else 0 for e in adcpAvg[adcp_to_plot].ensembles]
        vy_all = [e.velocity.y if not e.void else 0 for e in adcpAvg[adcp_to_plot].ensembles]
        vx = vx_all[::thin_factor]
        vy = vy_all[::thin_factor]

        # average length for the quiverkey
        avglen = np.sqrt(np.array(vx) ** 2 + np.array(vy) ** 2).mean()
        displen = np.around(avglen, int(np.ceil(np.abs(np.log10(avglen)))))

        # scale the vectors
        vscale = 1

        # axes.plot(x, y, marker='x', color='white')
        # axes.scatter(x[0], y[0], marker='o')
        axes.axis('equal')
        axes.grid()

        norm = Normalize(vmin=vmin, vmax=vmax)
        colors = np.sqrt(np.array(vx) ** 2 + np.array(vy) ** 2)
        # norm.autoscale(colors)

        # q = axes.quiver(x, y, vx * vscale, vy * vscale, width=0.003,
        #                 color=colormap(norm(colors)))
        axes.scatter(x, y, 25, colors, cmap='inferno', edgecolor='black', vmin=vmin, vmax=vmax)

        # Add transect number
        axes.text(x[0], y[0], str(adcp_to_plot), color='w', fontsize=20)

        # Print transect start and end points
        print('Transect ' + str(adcp_to_plot) +
              ' Start' + str(x[0]) + ', ' + str(y[0]) + ' to ' +
              str(x[-1]) + ', ' + str(y[-1]))


    # axes.set_xlim(left=445959.7, right=447219.7)
    # axes.set_ylim(bottom=5033522.5, top=5034850.8)
    # axes.set_xlim(left=445959.7, right=457000)
    # axes.set_ylim(bottom=5033522.5, top=5037000)
    axes.set_xlim(left=446747.7, right=447246.7)
    axes.set_ylim(bottom=5033841.5, top=5034305.8)

    ctx.add_basemap(axes, source=mapbox, crs='EPSG:26918')
    # 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()

# %% Transect Plotting
# Plot
# plot_profile_3d(adcpAvg[4])

for a in [4, 5, 6, 7, 8, 9, 10, 11, 13, 14]:

    fig, axes = plt.subplots(nrows=2, figsize=(8, 4), sharex=True)

    plot_profile_3d(adcpAvg[a], vmin=0, vmax=0.75,  axes=axes[0], cfg=dict(
         title='Observed Transect ' + str(a)), ymin=-7, ymax=0)

    DHI_tx_ax = dsT[a]['Current speed'].plot(ax=axes[1], label='Current Speed (m/s)', title='Model Transect ' + str(a),
                                 cmin=0, cmax=0.75, add_colorbar=False, cmap='turbo')
    DHI_tx_ax.set_ylim(-7, 0)
    DHI_tx_ax.grid()
    # Add Colorbar
    colorbar_ax = fig.add_axes([0.91, 0.1, 0.02, 0.8])
    cmap = mpl.cm.get_cmap('turbo')
    norm = mpl.colors.Normalize(vmin=0, vmax=0.75)

    cb1 = mpl.colorbar.ColorbarBase(ax=colorbar_ax, cmap=cmap,
                                    norm=norm,
                                    orientation='vertical', label='Current Speed (m/s)')
    plt.show()
    fig.savefig('C:/Users/arey/files/Projects/OttawaRiver/PythonADCP/Merge_Transect_' + str(a) + '.png',
                 bbox_inches='tight')



# %% Animated map
plt.rcParams['animation.ffmpeg_path'] = 'C:/Users/arey/Local/ffmpeg-2022-02-14-git-59c647bcf3-full_build/bin/ffmpeg.exe'

mapbox = 'https://api.mapbox.com/styles/v1/alexander0042/ckex9vtri0o6619p55sl5qiyv/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)

# Set model to plot
modelPlot = [15, 16, 17]
m = 16

# m = 11
vmax = 34.0
vmin = 33.7
cmap = 'inferno'

metadata = dict(title='Movie Test', artist='Matplotlib',
                comment='Movie support')
writer = animation.FFMpegWriter(fps=2, metadata=metadata, codec='h264')
fig, axes = plt.subplots(figsize=(8, 8))

writer.setup(fig, 'C:/Users/arey/files/Projects/Mustique/' +
             runLog['Run'][m][0:-20] + 'Intake.mp4')

# Animation function
for i in np.arange(1, 384, 1): #384
    if m == 8 or m == 17:
        disPlot = 7
    elif m == 11 or m == 15:
        disPlot = 13
    elif m == 14 or m == 16:
        disPlot = 8

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

    # Plot salinity
    # Select salinity data at last time step
    plotDat = ds_list[m]['Salinity'][i, :]
    # Identify nan values
    plot_nan = np.isnan(plotDat)

    # Add additional nan values to show bed
    plot_nan = plot_nan | (plotDat<33.725)

    # Plot all selected values at a given later and not nan
    axDHI = dfs_list[m].plot(plotDat[(dfs_list[m].layer_ids == 0) & (~plot_nan)], plot_type='contourf',
                             show_mesh=False, cmap='inferno', ax=axes, add_colorbar=False,
                             vmin=vmin, vmax=vmax, label='Salinity (PSU)', levels=24,
                             elements=dfs_list[m].element_ids[(dfs_list[m].layer_ids == 0) & (~plot_nan)])

    axes.scatter(readPoints[disPlot, 0], readPoints[disPlot, 1], marker='^', color='r', s=20)
    # axes.set_xlim(left=696680.7, right=698252.7)
    # axes.set_ylim(bottom=1425547.5, top=1426681.8)
    axes.set_xlim(left=696508.3, right=698252.7)
    axes.set_ylim(bottom=1425267.0, top=1426681.8)

    ctx.add_basemap(axes, source=mapbox, crs='EPSG:32620')
    axes.set_xlabel('Easting [m]')
    axes.set_ylabel('Northing [m]')

    axes.title.set_text(str(ds_list[m].time[i]))

    if i == 1:
        fig.subplots_adjust(right=0.8)
        cmap = mpl.cm.inferno
        norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
        cax = plt.axes([0.82, 0.25, 0.04, 0.5])
        cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap,
                                        norm=norm,
                                        orientation='vertical')
        cb1.set_label('Salinity (PSU)')


    writer.grab_frame()
    # plt.show()
    print(i)

writer.finish()


# %% Plot Bathymetry

fig, axes = plt.subplots(figsize=(8, 8))
# Convert to feet and ignore missing
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:32620')

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


# %% Setup path along pipe transect
m = 1

df_time_rows = elm_df_list[m][0].Z.drop_duplicates()
df_time_rowsidx = np.where(df_time_rows.index.values[0] == elm_df_list[m][0].index.values)[0]

tSteps = range(100, 2000)


u_plotarray = np.zeros((2000, 5, 81))
for layer in range(0, 5):
    for tStep in tSteps:
        u_plotarray[tStep, layer, :] = [x.iloc[df_time_rowsidx[layer] + tStep, 0] for x in elm_df_list[m]]

v_plotarray = np.zeros((2000, 5, 81))
for layer in range(0, 5):
    for tStep in tSteps:
        v_plotarray[tStep, layer, :] = [x.iloc[df_time_rowsidx[layer] + tStep, 0] for x in elm_df_list[m]]

mag_plotarray = np.sqrt(u_plotarray ** 2 + v_plotarray ** 2)

max_plotarray = np.max(mag_plotarray, axis=0)

# Find layer depths along transect and interpolate along vertical profile
layer_depth     = np.zeros((5, 81))
max_plot_interp = np.zeros((60, 81))
interp_depths   = np.arange(-6, 0, 0.1)

for chain in range(14, 81):
    layer_depth[:, chain] = elm_df_list[m][chain].Z.unique()
    max_plot_interp[:, chain] = np.interp(interp_depths, layer_depth[:, chain], max_plotarray[:, chain])
    # Set points below bed as nan
    max_plot_interp[interp_depths<layer_depth[0, chain], chain] = np.nan

# %% Plot along transect

fig, axes = plt.subplots(figsize=(8, 8))
# plt.pcolor(np.arange(0, 67*5, 5), elm_df_list[m][0].Z.unique(), max_plotarray[:, 14:], cmap='jet')
# for chain in range(15, 81):
    # plt.pcolor([(chain-15) * 5, (chain-14) * 5], layer_depth[:, chain-1:chain+1], max_plotarray[:, chain-1:chain+1], cmap='jet')

axes.set_ylabel('Depth (m)')
axes.set_xlabel('Distance along discharge pipe (m). Existing Pipe at 170 m')

plt.pcolor(np.arange(0, 67*5, 5), interp_depths, max_plot_interp[:, 14:], cmap='jet')
plt.colorbar(label="Maximum Current Speed (m/s)")
plt.show()

pth = pl.Path("//srv-ott3.baird.com/", "Projects", "13539.101 L'Ansecoy Bay, Mustique", "06_Models", "01_MIKE3", "00_Figures",
              "Discharge Velocity.png")

fig.savefig(pth.as_posix(), bbox_inches='tight', dpi=300)
