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

import dfm_tools as dfmt
from dfm_tools.get_nc import get_netdata, get_ncmodeldata, plot_netmapdata
from dfm_tools.get_nc_helpers import get_ncvarproperties, get_stationid_fromstationlist, get_varnamefromattrs


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

import rioxarray
import xarray as xr

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


# 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)
mag_mean            = [None] * (max(modelPlot)+1)
wl                  = [None] * (max(modelPlot)+1)
wd                  = [None] * (max(modelPlot)+1)
shear               = [None] * (max(modelPlot)+1)
pol_shp_list        = [None] * (max(modelPlot)+1)
alphaPolyList       = [None] * (max(modelPlot)+1)
X2                  = [None] * (max(modelPlot)+1)
Y2                  = [None] * (max(modelPlot)+1)
mag2                = [None] * (max(modelPlot)+1)
shear2              = [None] * (max(modelPlot)+1)

for i in modelPlot:
    file_nc_map = os.path.join(runLog['Run Location'][i],
        '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)
    shear[i]            = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_taus', timestep=tStep)

    mag_mean[i]         = (ux_mean[i]**2 + uy_mean[i]**2)**0.5

    # 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[0].tolist()],
                                                   facey[i][wd_mask[0].tolist()])).data))

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


    # Create mask of regular grid inside concave hull
    # plottingMask = [alphaPoly.contains(Point(i[0], i[1]))
    #                    for i in np.array([facex[i].flatten(), facey[i].flatten()]).T]

    # Convert to regular grid for interpolation
    # X2[i], Y2[i], mag2[i] = scatter_to_regulargrid(xcoords=facex[i][plottingMask], ycoords=facey[i][plottingMask],
    #                                       ncellx=2000, ncelly=4000, values=mag_mean[i][0, plottingMask],
    #                                       maskland_dist=50, method='nearest')

    X2[i], Y2[i], mag2[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                          ncellx=4000, ncelly=3000, values=mag_mean[i][0, :],
                                          maskland_dist=50, method='nearest')
    X2[i], Y2[i], shear2[i] = scatter_to_regulargrid(xcoords=facex[i], ycoords=facey[i],
                                          ncellx=4000, ncelly=3000, values=shear[i][0, :],
                                          maskland_dist=50, method='nearest')

#%% Save as geotiff
modelPlot  = [18, 35]

sedIDX = 0

for i in modelPlot:
    tiffPlot = mag2[i]

    # Blank zero sed
    tiffPlot[tiffPlot == 0] = np.nan

    # Blank outside of hull
    # sedPlot[~regularMask2[i]] = np.nan

    # Convert to xarray dataset
    tiffPlot_xr = xr.DataArray(data=tiffPlot.T,
                           dims=["x", "y"],
                           coords=dict(
                               x=(["x"], X2[i][0, :]),
                               y=(["y"], Y2[i][:, 0])))

    # Transpose for geotiff writing
    tiffPlot_xr = tiffPlot_xr.transpose('y', 'x')
    tiffPlot_xr.rio.write_crs("epsg:32615", inplace=True)
    tiffPlot_xr.rio.to_raster('C:/Users/arey/files/Projects/Grassy Narrows/Slides/' +
         runLog['Run Title'][i] + '_mag.tif')


    tiffPlot = shear2[i]

    # Blank zero sed
    tiffPlot[tiffPlot == 0] = np.nan

    # Blank outside of hull
    # sedPlot[~regularMask2[i]] = np.nan

    # Convert to xarray dataset
    tiffPlot_xr = xr.DataArray(data=tiffPlot.T,
                           dims=["x", "y"],
                           coords=dict(
                               x=(["x"], X2[i][0, :]),
                               y=(["y"], Y2[i][:, 0])))

    # Transpose for geotiff writing
    tiffPlot_xr = tiffPlot_xr.transpose('y', 'x')
    tiffPlot_xr.rio.write_crs("epsg:32615", inplace=True)
    tiffPlot_xr.rio.to_raster('C:/Users/arey/files/Projects/Grassy Narrows/Slides/' +
         runLog['Run Title'][i] + '_shear.tif')

#%% Plot using DFM functions
modelPlot = [24]

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)


