#%% 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, Polygon
import alphashape

#%% 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  = [19, 22]

# stormTime =  [datetime.datetime(2005, 8, 7, 0, 0, 0)]

ugrid_all           = [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)
wl                  = [None] * (max(modelPlot)+1)
wd                  = [None] * (max(modelPlot)+1)
pol_shp_list        = [None] * (max(modelPlot)+1)
alphaPolyList       = [None] * (max(modelPlot)+1)

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

    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)

    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_ucx', timestep=tStep)
    uy_mean[i]          = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucy', timestep=tStep)
    uy_mean[i]          = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucy', timestep=tStep)
    wl[i]               = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_s1', timestep=tStep)
    wd[i]               = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_waterdepth', timestep=tStep)


    # Setup polygons for shapefile export
    pol_shp_list[i] = []

    #partly from dfm_tools.ugrid.polygon_intersect()
    for iP, pol_data in enumerate(ugrid_all[i].verts):
        pol_data_nonan = pol_data[~np.isnan(pol_data).all(axis=1)]
        pol_shp = Polygon(pol_data_nonan)
        pol_shp_list[i].append(pol_shp)

    print('Polygon Complete')

    # Find convex hull of pcolor data
    # Convert grid to multipoint

    # Create a list of grid points where the water depth is greater than zero
    wd_mask = (wd[i][:] > 0)
    alphaPoints = list(map(tuple, np.column_stack((facex[i][wd_mask.tolist()],
                                                   facey[i][wd_mask.tolist()])).data))

    # Find the concave hull
    alphaPoly        = alphashape.alphashape(alphaPoints, 0.02)
    alphaPolyList[i] = alphaPoly


#%% Plot using DFM functions
modelPlot = [19, 22]

for i in modelPlot:
    fig, axes = plt.subplots(nrows=1, ncols=1, subplot_kw={'projection': ccrs.epsg(32615)}, figsize=(14, 10))
    fig.patch.set_facecolor('white')

    plt.tight_layout()

    plotCount = 1

    vel_mask = (wd[i][:] > 0)

    # pc = plot_netmapdata(ugrid_all[i].verts[vel_mask.tolist()],
    #                      values=np.sqrt(ux_mean[i][vel_mask]**2 +\
    #                                     uy_mean[i][vel_mask]**2),
    #                      ax=axes, linewidth=0.5, cmap="viridis")

    pc = plot_netmapdata(ugrid_all[i].verts[vel_mask.tolist()],
                         values=wd[i][vel_mask],
                         ax=axes, linewidth=0.5, cmap="viridis")

    axes.set_xlim(494649, 513647)
    axes.set_ylim(5514653, 5525256)

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

    # Add model bounds line
    alphaPoly = alphaPolyList[i]
    # for p in range(0, len(alphaPoly.geoms)):
    #     axes.plot(*alphaPoly.geoms[p].exterior.xy, color='k', linewidth=1)
    axes.plot(*alphaPoly.exterior.xy, color='k', linewidth=1)

    plt.show()

    # # Save as shapefile
    dataShape = gp.GeoDataFrame()
    dataShape['geometry'] = pol_shp_list[i]
    dataShape.set_crs(epsg=32615)
    dataShape['WaterLevel'] = wl[i].data.flatten()
    dataShape['WaterDepth'] = wd[i].data.flatten()
    dataShape.to_file('C:/Users/arey/files/Projects/Grassy Narrows/Modelling/Flooding/' + runLog['Run Name'][i] + '.shp')

    # dataBounds = gp.GeoDataFrame
    # dataBounds['geometry'] = alphaPoly.geoms

    features = [i for i in range(len(alphaPoly))]
    dataBounds = gp.GeoDataFrame(index=[0], geometry=[alphaPoly], crs='EPSG:32615')
    dataBounds.to_file('C:/Users/arey/files/Projects/Grassy Narrows/Modelling/Flooding/Bounds_' + runLog['Run Name'][i] + '.shp')


# fig.savefig(
#     'C:/Users/arey/files/Projects/LSU/Modelling/Figures/' +
#     'StormVelDisWide.png', bbox_inches='tight', dpi=400)