#%% Plotting script for Penobscot 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 as mpl
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


import dataretrieval.nwis as nwis
import noaa_coops as noaa
#%% Read Model Log
pth = pl.Path("C:/Users/arey/files/Projects/Penobscot/Modelling/12345.678.W.AJMR.Rev0_ModelLog.xlsx")

runLog = pd.read_excel(pth.as_posix(), "ModelLog")

dataPath = "FlowFM_map.nc"

#%% Import using DFM functions
modelPlot = range(3, 4)

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)
sal                 = [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)
sal2                = [None] * (max(modelPlot)+1)
shear2              = [None] * (max(modelPlot)+1)

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], 'dflowfm', 'output', 'FlowFM_map.nc')

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

    # Find nearest time step to desired time
    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_ucxa', timestep=tStep)
    uy_mean[i]          = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_ucya', 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)
    sal[i]              = get_ncmodeldata(file_nc=file_nc_map, varname='mesh2d_sa1', timestep=tStep, layer='all')

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


    # 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.0001)
    alphaPolyList[i] = alphaPoly

    plottingMask = (facex[i] > 510000) & (facex[i] < 521984) & \
                   (facey[i] > 4914195) & (facey[i] < 4960418)

    # Convert to regular grid for interpolation
    X2[i], Y2[i], mag2[i] = scatter_to_regulargrid(xcoords=facex[i][plottingMask], ycoords=facey[i][plottingMask],
                                          ncellx=150, ncelly=470, values=mag_mean[i][0, plottingMask],
                                          maskland_dist=400, method='linear')

    # Convert to regular grid for interpolation
    X2[i], Y2[i], sal2[i] = scatter_to_regulargrid(xcoords=facex[i][plottingMask], ycoords=facey[i][plottingMask],
                                          ncellx=150, ncelly=470, values=sal[i][0, plottingMask, 0],
                                          maskland_dist=400, method='linear')


#%% Import hist using DFM functions

modelPlot = range(3, 4)

modelHistWL   = [None] * (max(modelPlot)+1)
modelHistTime  = [None] * (max(modelPlot)+1)
# tSteps        = [None] * (max(modelPlot)+1)

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

    vars_pd, dims_pd = get_ncvardimlist(file_nc=file_nc_hist)

    # Get station names
    histStations = get_hisstationlist(file_nc=file_nc_hist, varname='waterlevel')

    # Import
    modelHistTime[i] = get_timesfromnc(file_nc=file_nc_hist, varname='time')
    modelHistWL[i]   = get_ncmodeldata(file_nc=file_nc_hist, varname='waterlevel',
                                       timestep='all', station='all')

# %% Import Validation sites
## Validation sites
# Winterport, Maine:
# Gage from 2017-09-01 to 2017-12-01
site = '443810068502201'
Winterport_df = nwis.get_record(sites=site, service='iv', start='2017-09-01', end='2017-12-01',
                                parameterCd='00065')
Winterport_df['pd_Datetime'] = pd.to_datetime(Winterport_df.index, format='%Y-%m-%dT%H:%M:%S.%f%z')
Winterport_df['pd_Datetime'] = Winterport_df['pd_Datetime'].apply(lambda x: x.tz_convert('UTC'))
Winterport_df.set_index(pd.DatetimeIndex(Winterport_df['pd_Datetime']), inplace=True)


# Bucksport, Maine:
# Gage from 2017-09-01 to 2017-12-01
site = '443409068471801'
Bucksport_df = nwis.get_record(sites=site, service='iv', start='2017-09-01', end='2017-12-01',
                               parameterCd='00065')
Bucksport_df['pd_Datetime'] = pd.to_datetime(Bucksport_df.index, format='%Y-%m-%dT%H:%M:%S.%f%z')
Bucksport_df['pd_Datetime'] = Bucksport_df['pd_Datetime'].apply(lambda x: x.tz_convert('UTC'))
Bucksport_df.set_index(pd.DatetimeIndex(Bucksport_df['pd_Datetime']), inplace=True)

# FortPoint, Maine:
# Gage from 2017-09-01 to 2017-12-01
site = '442810068480101'
FortPoint_df = nwis.get_record(sites=site, service='iv', start='2017-09-01', end='2017-12-01',
                               parameterCd='00065')
FortPoint_df['pd_Datetime'] = pd.to_datetime(FortPoint_df.index, format='%Y-%m-%dT%H:%M:%S.%f%z')
FortPoint_df['pd_Datetime'] = FortPoint_df['pd_Datetime'].apply(lambda x: x.tz_convert('UTC'))
FortPoint_df.set_index(pd.DatetimeIndex(FortPoint_df['pd_Datetime']), inplace=True)

# Import NOAA CO-OPS data at Bar Harbor
barHarbor = noaa.Station(8413320)

barHarbor_df = barHarbor.get_data(
     begin_date="20170901",
     end_date="20171201",
     product="water_level",
     units="metric",
     datum="MSL",
     time_zone="gmt")

barHarbor_df['DateTime'] = barHarbor_df.index
barHarbor_df['DateTime'] = barHarbor_df['DateTime'].dt.tz_localize('UTC')

#%% Plot Validation
i = modelPlot
i = 3


fig, axes = plt.subplots(nrows=3, ncols=1, figsize=(10, 6), sharex=True, sharey=True)

axes[0].plot(modelHistTime[i][:], modelHistWL[i][:, 0]*3.28084, label='Model')
FortPoint_df.loc['2017-9-1':'2017-9-7', '00065_feet above navd88'].plot(ax=axes[0], label='Observation')
barHarbor_df.loc['2017-9-1':'2017-9-7', 'water_level'].multiply(3.28084).plot(ax=axes[0], label='Bar Harbour')
axes[0].legend()


axes[1].plot(modelHistTime[i][:], modelHistWL[i][:, 1]*3.28084, label='Model')
Bucksport_df.loc['2017-9-1':'2017-9-7', '00065_feet above navd88'].plot(ax=axes[1], label='Observation')
axes[1].legend()


axes[2].plot(modelHistTime[i][:], modelHistWL[i][:, 2]*3.28084, label='Model')
Winterport_df.loc['2017-9-1':'2017-9-7', '00065_feet above navd88'].plot(ax=axes[2], label='Observation')
axes[2].legend()

axes[2].set_xlabel('Date')
axes[1].set_ylabel('Water Level [ftMSL]')

plt.show()

# fig.savefig('C:/Users/arey/files/Projects/Penobscot/Modelling/WL Validation.png',
#             bbox_inches='tight', dpi=300)

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

for i in modelPlot:
    fig, axes = plt.subplots(nrows=1, ncols=1, subplot_kw={'projection': ccrs.epsg(32619)}, figsize=(10, 14))
    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=sal[i][vel_mask, 0],
                         ax=axes, linewidth=0.5, cmap="viridis")

    cb1 = fig.colorbar(pc, ax=axes)

    axes.set_xlim(510101, 521984)
    axes.set_ylim(4925241, 4960418)


    # 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:32619')

    # Add model bounds line
    # alphaPoly = alphaPolyList[i]
    # axes.plot(*alphaPoly.exterior.xy, color='k', linewidth=1)

    plt.show()


#%% Plot Sediment  DFM functions
modelPlot = range(0, 1)

sedIDX = 0

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

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

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

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

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