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

import rioxarray
import xarray as xr
#%% 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 = range(66, 72)

# stormTime =  [datetime.datetime(2005, 8,  29, 14, 0, 0),
#               datetime.datetime(2005, 9,  24, 14, 0, 0),
#               datetime.datetime(2005, 12, 28, 14, 0, 0)]
# stormTime =  [datetime.datetime(2005, 1, 14, 18, 0, 0),
#               datetime.datetime(2005, 2,  3, 18, 0, 0),
#               datetime.datetime(2005, 2,  10, 18, 0, 0)]
#stormTime =  [datetime.datetime(2005, 12, 31, 18, 0, 0)]
#stormTime =  [datetime.datetime(2005, 8, 7, 0, 0, 0)]


ugrid_all           = [None] * (max(modelPlot)+1)
X2                  = [None] * (max(modelPlot)+1)
Y2                  = [None] * (max(modelPlot)+1)
sed2                = [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)
ux                  = [None] * (max(modelPlot)+1)
uy                  = [None] * (max(modelPlot)+1)
sed                 = [None] * (max(modelPlot)+1)
mag                 = [None] * (max(modelPlot)+1)
bed                 = [None] * (max(modelPlot)+1)
bed2                = [None] * (max(modelPlot)+1)
mag2                = [None] * (max(modelPlot)+1)
wd                  = [None] * (max(modelPlot)+1)
regularMask2        = [None] * (max(modelPlot)+1)
meshMask            = [None] * (max(modelPlot)+1)
alphaPolyList       = [None] * (max(modelPlot)+1)

sedIDX = 0
tPlot = 0

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

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

    # Find nearest time step to desired time
    # stormTime = [runLog['End Date'][i]]
    # tStep = []
    # for s in stormTime:
    #     abs_deltas_from_target_date = np.absolute(tSteps - s)
    #     tStep.append(np.argmin(abs_deltas_from_target_date))
    tStep = len(tSteps) - 1


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

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

    sed[i]              = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_bodsed',
                                          timestep=tStep, station='all')
    bed[i]              = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_flowelem_bl')
    wd[i]               = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_waterdepth',
                                          timestep=tStep)
    mag[i]              = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucmag',
                                          timestep=tStep, layer=4)

    # Create a list of grid points
    alphaPoints = list(map(tuple, np.column_stack((facex[i][wd[i][0, :] != 0],
                                                   facey[i][wd[i][0, :] != 0])).data))

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

    plottingMask = (facex[i] > 676470) & (facex[i] < 677021) & \
                   (facey[i] > 3365939) & (facey[i] < 3366617)

    # Convert to regular grid for interpolation
    X2[i], Y2[i], sed2[i] = scatter_to_regulargrid(xcoords=facex[i][plottingMask], ycoords=facey[i][plottingMask],
                                          ncellx=500, ncelly=500, values=sed[i][tPlot, plottingMask, sedIDX],
                                          maskland_dist=25, method='linear')

    X2[i], Y2[i], bed2[i] = scatter_to_regulargrid(xcoords=facex[i][plottingMask], ycoords=facey[i][plottingMask],
                                          ncellx=750, ncelly=750, values=bed[i][plottingMask],
                                          maskland_dist=25, method='cubic')

    X2[i], Y2[i], mag2[i] = scatter_to_regulargrid(xcoords=facex[i][plottingMask], ycoords=facey[i][plottingMask],
                                          ncellx=750, ncelly=750, values=mag[i][0, plottingMask, 0],
                                          maskland_dist=25, method='cubic')


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

    # Reshape back to regular grid for mask
    regularMask2[i] = np.array(regularMask2[i]).reshape(X2[i].shape)


# %% Setup forebay polygon
# forebayPoly = gp.read_file('C:/Users/arey/files/Projects/LSU/ForebayPoly.shp')
forebayPoly = gp.read_file('C:/Users/arey/files/Projects/LSU/ForebayPolyExpand.shp')

# Setup mask. Can only do this once if grids are the same
i = 45
forebayPolyMask = [forebayPoly.iloc[0, 1].contains(Point(i[0], i[1]))
                   for i in np.array([facex[i], facey[i]]).T]

#%% Plot Sediment  DFM functions
modelPlot = range(66, 72)

sedIDX = 0

for tPlot in range(0, 1):
    for i in modelPlot:
        # sedPlot = bed2[i]
        sedPlot = mag2[i]

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

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

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

        # Transpose for geotiff writing
        sed2_xr = sed2_xr.transpose('y', 'x')
        sed2_xr.rio.write_crs("epsg:26915", inplace=True)
        sed2_xr.rio.to_raster('C:/Users/arey/files/Projects/LSU/Modelling/GeoTiffMag/Mag_' +
             runLog['Run Name'][i] + '.tif')
