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

import alphashape

import rioxarray
import xarray as xr

from shapely import geometry, ops

#%% Read in centerline shapefile
river_centerline = gp.read_file('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/03_Data/02_Physical/16_Waterline/Centerline_for_Modelling_UTMZ15.shp')

river_centerlineExploded = river_centerline.explode(ignore_index=True)
river_centerlineExploded.reset_index(inplace=True)

tempMulti = river_centerlineExploded.iloc[[5,0,1,2,3,4,6,7,9], 4]
# Put the sub-line coordinates into a list of sublists
outcoords = [list(i.coords) for i in tempMulti]
# Flatten the list of sublists and use it to make a new line
river_centerline_merge = geometry.LineString([i for sublist in outcoords for i in sublist])
river_centerline_merge_gpd = gp.GeoSeries(river_centerline_merge)

river_centerline_merge_gpd2 =\
    gp.GeoDataFrame(geometry=gp.points_from_xy(
        river_centerline_merge.xy[0], river_centerline_merge.xy[1], crs="EPSG:32615"))

# Add distance along centerline
river_centerline_merge_gpd2['DistanceFromPrevious'] = river_centerline_merge_gpd2.distance(river_centerline_merge_gpd2.shift(1))
river_centerline_merge_gpd2['RiverKM'] = river_centerline_merge_gpd2['DistanceFromPrevious'].cumsum()
river_centerline_merge_gpd2.iloc[0, 1] = 0
river_centerline_merge_gpd2.iloc[0, 2] = 0

#%% Save River Centerline as Shapefile
river_centerline_merge_gpd2.to_file('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/03_Data/02_Physical/16_Waterline/Centerline_for_Modelling_UTMZ15_RiverKM.shp')

#%% Read Model Log
pth = pl.Path("//srv-ott3.baird.com/", "Projects", "12828.101 English Wabigoon River", "06_Models", "00_Delft3D", "ModelRuns.xlsx")

runLog = pd.read_excel(pth.as_posix(), "Test runs", skiprows=1)

dataPath = "FlowFM_map.nc"

#%% Import using DFM functions
#modelPlot  = [35, 36, 37, 39, 40, 41, 42, 43, 44, 45]
modelPlot  = [18, 35]

dataIN_gp           = [None] * (max(modelPlot)+1)
dataOUT             = [None] * (max(modelPlot)+1)

for i in modelPlot:

    if (i == 39) or (i == 40) or (i == 41) or (i == 42) or (i == 43) or (i == 44) or (i == 45):
        Flume = True
    else:
        Flume = False

    # Find Map File in Output Folder
    outputFiles = os.listdir(os.path.join(runLog['Run Location'][i],
            'output'))
    for file in outputFiles:
        if file.endswith("map.nc"):
            mapFile = file

    file_nc_map = os.path.join(runLog['Run Location'][i],
        'output', mapFile)

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

    # Find nearest time step to desired time
    # tStep = []
    # for s in stormTime:
    #     abs_deltas_from_target_date = np.absolute(tSteps - s)
    #     tStep.append(np.argmin(abs_deltas_from_target_date))

    # Otherwise, define a timestep
    tStep = len(tSteps)-2

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


    # Define extraction variables
    # if (i == 37) or (i == 39):
    #     dataVars = ['s1', 'waterdepth', 'taus', 'ucmaga']
    # else:

    dataVars = ['s1', 'waterdepth', 'taus', 'ucmag']

    # Create Pandas Output Array
    dataIN_x = get_ncmodeldata(file_nc=file_nc_map,
                               varname='mesh2d_face_x',
                               silent=True)
    dataIN_y = get_ncmodeldata(file_nc=file_nc_map,
                               varname='mesh2d_face_y',
                               silent=True)

    dataIN_gp[i] = gp.GeoDataFrame(geometry=gp.points_from_xy(dataIN_x, dataIN_y, crs="EPSG:32615"))

    # Extract variables
    for vIDX, v in enumerate(dataVars):
        # Extract data from NetCDF file
        if len(vars_pd.loc['mesh2d_' + v, 'shape']) == 3:
            dataIN = get_ncmodeldata(file_nc=file_nc_map,
                                     varname='mesh2d_' + v, timestep=len(tSteps) - 1,
                                     silent=True, layer='all')
            dataIN = dataIN.mean(axis=2)
        else:
            dataIN = get_ncmodeldata(file_nc=file_nc_map,
                                     varname='mesh2d_' + v, timestep=len(tSteps) - 1,
                                     silent=True)  # , layer=0


        dataIN_gp[i][v] = np.array(dataIN[0][:])
        print(v)


    # Extract Bed Level
    dataIN = get_ncmodeldata(file_nc=file_nc_map,
                             varname='mesh2d_flowelem_bl',
                             silent=True)  # , layer=0

    dataIN_gp[i]['mesh2d_flowelem_bl'] = np.array(dataIN[:])
    print('mesh2d_flowelem_bl')


    # Create Dataframe along river centerline for plotting
    if Flume == True:
        dataOUT[i] = dataIN_gp[i]
        dataOUT[i]['RiverKM'] = dataIN_x
    else:
        dataOUT[i] = gp.sjoin_nearest(river_centerline_merge_gpd2, dataIN_gp[i], how='left', max_distance=100)

    dataOUT[i].set_index('RiverKM', inplace=True)

#%% Key Profile Sites
KeySites = dict()
KeySites['Dryden'] = 0
KeySites['Rugby Creek'] = 2650
KeySites['Eagle River'] = 44200
KeySites['Beaver Creek'] = 38600
KeySites['Clay Lake'] = 90000
KeySites['Gullwing River'] = 20100
KeySites['Cowamula'] = 57400
# KeySites['Buller Creek'] = 66000

# Structure Sites
StructureSites = dict()
StructureSites['Mutrie Lake'] = 54214
StructureSites['Highway 105'] = 64998
StructureSites['Quibell'] = 70000
StructureSites['Wainwright Dam'] = 6013

#%% Plot Data along centerline
fig, axes = plt.subplots(nrows=5, ncols=1, figsize=(9, 9), sharex=True)
fig.patch.set_facecolor('white')
fig.tight_layout(pad=3)
ax = axes.flat

modelPlot  = [35, 37, 40, 41, 42, 43, 44]
modelPlot  = [40, 41, 42, 43, 44]
modelPlot  = [35, 37, 40, 45]
modelPlot  = [18, 35]


for i in modelPlot:
    # S1
    ax[0].set_title('Water Level')
    ax[0].plot(dataOUT[i].index, dataOUT[i]['s1'], linewidth=3, label=runLog['Run Title'][i])
    ax[0].set_ylim([330, 345])
    # ax[0].set_xlim([0, 10000])
    ax[0].set_xlim([0, 90000])

    ax[0].legend(loc='upper right')
    ax[0].set_ylabel('Water Level [m]')

    # Shear
    ax[1].set_title('Shear Stress')
    ax[1].plot(dataOUT[i].index, dataOUT[i]['taus'], linewidth=3, label=runLog['Run Title'][i])
    ax[1].set_ylim([0, 2])

    ax[1].legend(loc='upper right')
    ax[1].set_ylabel('Shear Stress [N/m^2]')

    # Water Depth
    ax[2].set_title('Water Depth')
    ax[2].plot(dataOUT[i].index, dataOUT[i]['waterdepth'], linewidth=3, label=runLog['Run Title'][i])
    ax[2].set_ylim([0, 20])

    ax[2].legend(loc='upper right')
    ax[2].set_ylabel('Water Depth [m]')

    # Bed Level
    ax[3].set_title('Bed Level')
    ax[3].plot(dataOUT[i].index, dataOUT[i]['mesh2d_flowelem_bl'], linewidth=3, label=runLog['Run Title'][i])
    # ax[3].set_ylim([320, 340])

    ax[3].legend(loc='upper right')
    ax[3].set_ylabel('Bed Level [m]')

    # Velocity
    ax[4].set_title('Depth Averaged Velocity')
    ax[4].plot(dataOUT[i].index, dataOUT[i]['ucmag'], linewidth=3, label=runLog['Run Title'][i])
    ax[4].set_ylim([0, 2])

    ax[4].legend(loc='upper right')
    ax[4].set_ylabel('Velocity [m/s]')
    ax[4].set_xlabel('Distance Along River [m]')

plt.show()

# fig.savefig('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/06_Models/00_Delft3D/12_WAQ/FullDomainFigures/DomainHDCompare.png',
#              bbox_inches='tight', dpi=200)

fig.savefig('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/06_Models/00_Delft3D/07_PostProcessing/HD_Profile_Compare.png',
             bbox_inches='tight', dpi=200)

#%% Plot select data along centerline
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(9, 6), sharex=True)
ax = axes.flat

# fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(9, 4), sharex=True)
# ax[0] = axes

fig.patch.set_facecolor('white')
# Expand plot to allow space for legend
# fig.subplots_adjust(left=-2)


modelPlot  = [18, 35]

for i in modelPlot:

    # Velocity
    ax[0].set_title('Depth Averaged Velocity Along River Centerline')
    ax[0].plot(dataOUT[i].index / 1000, dataOUT[i]['ucmag'], linewidth=3, label=runLog['Run Title'][i], zorder=1)
    ax[0].set_ylim([0, 2])
    ax[0].set_xlim([0, 90])
    ax[0].set_ylabel('Velocity [m/s]')

    # Shear
    ax[1].set_title('Shear Stress Along River Centerline')
    ax[1].plot(dataOUT[i].index/1000, dataOUT[i]['taus'], linewidth=3, label=runLog['Run Title'][i], zorder=1)
    ax[1].set_ylim([0, 2])
    ax[1].set_xlim([0, 90])
    ax[1].set_ylabel('Shear Stress [N/$\mathregular{m^2}$]')

    ax[1].set_xlabel('Distance Along River [km]')

# Add in structure sites and key sites
for site in KeySites.keys():
    for axID in range(0, 2):
        if site == 'Dryden':
            ax[axID].scatter(KeySites[site]/1000, 1, marker='^', color='k', s=30, label='Inflow', zorder=2)
        else:
            ax[axID].scatter(KeySites[site]/1000, 1, marker='^', color='k', s=30, zorder=10)

        ax[axID].text(KeySites[site]/1000, 1.1, site, rotation=45)

for structureSite in StructureSites.keys():
    for axID in range(0, 2):
        if structureSite == 'Mutrie Lake':
            ax[axID].scatter(StructureSites[structureSite]/1000, 1, marker='s', color='k', label='Hydraulic Structure',
                             s=30, zorder=10)
        else:
            ax[axID].scatter(StructureSites[structureSite]/1000, 1, marker='s', color='k', s=30, zorder=10)


        ax[axID].text(StructureSites[structureSite] / 1000, 1.1, structureSite, rotation=45)

# Add remobilization threshold lines
ax[0].plot([0, 90000], [0.1, 0.1], linestyle='--', linewidth=1, label='Silt Remobilization Threshold', zorder=5)
ax[0].plot([0, 90000], [1.3, 1.3], linestyle='--', linewidth=1, label='Sand Remobilization Threshold', zorder=5)
ax[0].legend(bbox_to_anchor=(1.01, 0.5), loc="upper left")

fig.tight_layout(pad=1)

plt.show()

# fig.savefig('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/06_Models/00_Delft3D/07_PostProcessing/HD_Profile_Compare_DAV.png',
#              bbox_inches='tight', dpi=300)
fig.savefig('//srv-ott3.baird.com/Projects/12828.101 English Wabigoon River/06_Models/00_Delft3D/07_PostProcessing/HD_Profile_Compare_RevD.png',
             bbox_inches='tight', dpi=300)