#%% Plotting script for LSU D3D data

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

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



#%% 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  = [18]

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)

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

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

    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)

    X[i],Y[i],U[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                            ncellx=240, ncelly=340, values=ux_mean[i][0, :], maskland_dist=3)
    X[i],Y[i],V[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                            ncellx=240, ncelly=340, values=uy_mean[i][0, :], maskland_dist=3)


#%% Plot using DFM functions

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

# ax = axes.flat

fig.tight_layout(pad=3.0)
plotCount = 0

for i in modelPlot:
    pc = plot_netmapdata(ugrid_all[i].verts, values=np.sqrt(ux_mean[i][0,:]**2 + uy_mean[i][0,:]**2),
                         ax=axes, linewidth=0.5, cmap="viridis")

    axes.set_xlim(675524, 677027)
    axes.set_ylim(3365191, 3367182)

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

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

    pc.set_clim([0, 0.5])
    cbar = fig.colorbar(pc, ax=axes, shrink=0.95)
    cbar.set_label('Depth Averaged Velocity [m/s]')

    plotCount = plotCount + 1

plt.show()