#%% Example Penobscot Plotting Script
# Alexander Rey, May 18, 2023

#%% Import modules
import xarray as xr
import matplotlib.pyplot as plt
import contextily as ctx
import numpy as np
import dfm_tools as dfmt
import pandas as pd
import pathlib as pl

#%% Read in Model Log
pth = pl.Path("N:/Penobscot/Modelling/12345.678.W.AJMR.Rev0_PenobscotModelLog.xlsx")

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

dataPath = "FlowFM_map.nc"


#%% Read in Delft3D Data
i = 3

# Get hist file path
file_nc_hist = pl.Path(runLog['Run Location'][i],
                       'FlowFM', 'output', 'FlowFM_his.nc')
file_nc_hist = "N:/Penobscot/Modelling/Runs/04_SalinityRun.dsproj_data/FlowFM/output/FlowFM_his.nc"
# Open hist file dataset
data_xr = xr.open_mfdataset(file_nc_hist, preprocess=dfmt.preprocess_hisnc)

# Get stations
stations_pd = data_xr['stations'].to_dataframe()

#%% Extract and plot water level

# Extract water level
waterlevel_xr = data_xr['waterlevel'].sel(stations=stations_pd.index[1])

# Plot water level
fig, ax = plt.subplots(figsize=(12, 6))
waterlevel_xr.plot(ax=ax, label='Water Level')

plt.show()

#%% Extract and plot water level map
file_nc_map = "N:/Penobscot/Modelling/Runs/04_SalinityRun.dsproj_data/FlowFM/output/FlowFM_map.nc"

# Read in map file
d3dfm_DataArray = dfmt.open_partitioned_dataset(file_nc_map)

# Get Var info
vars_pd = dfmt.get_ncvarproperties(d3dfm_DataArray)

# Get map of water level
waterlevel_map_xr = d3dfm_DataArray.mesh2d_s1[3, :].compute()


#%% Plot water level map using delft3d map plotting function and contextily
fig, ax = plt.subplots(figsize=(12, 6))

pc = waterlevel_map_xr.ugrid.plot(ax=ax, edgecolors='face', cmap='jet')

# Define crs
crs = 'EPSG:32619'
ctx.add_basemap(ax=ax, source=ctx.providers.Esri.WorldImagery, crs=crs, attribution=False)

plt.show()