#%% 
#Project: 11934.201 Newtown Creek TPP - Privileged and Confidential 
#Confidentiality Note: For internal discussion only. 
#Description: This code creates plots for comparing temperature, salinity between the EFDC model, Delft3DFM, and the measured data. 
#To run, ensure that the model log is saved locally, and change the path in the code. 

#%% # -*- coding: utf-8 -*-
"""
Created on Wed Feb  8 11:14:28 2023

@author: mrobinson
"""
import pandas as pd 
import numpy as np 
import matplotlib.pyplot as plt 
import scipy
from scipy.interpolate import interp1d
from scipy import interpolate

from numpy import trapz

from datetime import datetime, timedelta
from matplotlib import dates as mpl_dates
from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)

import os
import pandas as pd
import numpy as np
from sklearn.metrics import mean_squared_error
import geopandas as gp
import pytz

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.animation as animation

import dfm_tools as dfmt

import cartopy.crs as ccrs
import contextily as ctx

from shapely.geometry import Point, MultiPoint
from shapely.ops import nearest_points

import pathlib as pl
import xarray as xr

from datetime import timedelta
import datetime as datetime


#%% Load in Measured data

YSI_df = pd.read_csv('//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/03_Data/02_Physical/21_YSI Data/D1-D9_Combined/Combined_with_2015/All_Combine.csv')

# Filter out invalid data
YSI_df.loc[YSI_df['Temp'] < -30, 'Temp'] = np.nan
YSI_df.loc[YSI_df['Sal'] < 0, 'Sal'] = np.nan
YSI_df['DateTime'] = pd.to_datetime(YSI_df['Date Time Local'])
YSI_df.set_index('DateTime', inplace=True)

# Set index as DateTime
YSI_df.index = YSI_df.index.tz_localize(pytz.timezone('America/New_York'),
                                  ambiguous='NaT', nonexistent='NaT')

YSI_df.index = YSI_df.index.tz_convert(pytz.utc)

#%% Load in EFDC Output 
FILE=r"//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/06_Models/02_EFDC/06_EFDC_Outputs/2014-2015/EFDC_Compiled_2014_2015.xlsx"

EFDC_dat = pd.read_excel(FILE, None)

# Loop through stations and add to dataframe
Stations=['EB043', 'EK108', 'NC310', 'NC313', 'NC316', 'NC318']
for Stat in Stations:
    # Set index as DateTime
    EFDC_dat[Stat].set_index('date fix', inplace=True)
    EFDC_dat[Stat].index = pd.to_datetime(EFDC_dat[Stat].index).tz_localize(
                        pytz.timezone('UTC'))

#%% Load in Telemac Output

FILE_T=r"C:/Users/arey/Downloads/Results_Summarized_Telemac.xlsx"

TM_dat = pd.read_excel(FILE_T, None)

# Loop through stations and adjust times
Stations=['EB043', 'EK108', 'NC310', 'NC313', 'NC316', 'NC318']
for Stat in Stations:
    # Set index as DateTime
    TM_dat[Stat].set_index('Date_Time', inplace=True)
    TM_dat[Stat].index = pd.to_datetime(TM_dat[Stat].index).tz_localize(
                        pytz.timezone('UTC'))

#%%  Load in Water Level Data
    # Field Gauge-EAST
FFG_pd = pd.read_csv("//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/03_Data/02_Physical/00_ProcessingCode/NetCDF/Gauge/FieldFacility.csv")
FFG_pd['DateTime'] = pd.to_datetime(FFG_pd.Date_time)
FFG_pd.set_index(FFG_pd['DateTime'], inplace=True)
# Put in UTC from Eastern Time (with DST!)
FFG_pd = FFG_pd.tz_localize(
    pytz.timezone('US/Eastern'), ambiguous='infer').tz_convert(pytz.utc)

# National Grid Gauge-WEST
NGG1_pd = pd.read_csv("//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/03_Data/02_Physical/00_ProcessingCode/NetCDF/Gauge/NationalGrid1.csv")
NGG1_pd['DateTime'] = pd.to_datetime(NGG1_pd.Date_time)
NGG1_pd.set_index(NGG1_pd['DateTime'], inplace=True)
# Put in UTC from Eastern Time (with DST!)
NGG1_pd = NGG1_pd.tz_localize(
    pytz.timezone('US/Eastern'), ambiguous='NaT').tz_convert(pytz.utc)
# Drop ambiguous times
NGG1_pd = NGG1_pd.dropna()

NGG2_pd = pd.read_csv("//srv-ott3/Projects/11934.201 Newtown Creek TPP – Privileged and Confidential/03_Data/02_Physical/00_ProcessingCode/NetCDF/Gauge/NationalGrid2.csv")
NGG2_pd['DateTime'] = pd.to_datetime(NGG2_pd.datetime)
NGG2_pd.set_index(NGG2_pd['DateTime'], inplace=True)
# Put in UTC from Eastern Time (with DST!)
NGG2_pd = NGG2_pd.tz_localize(
    pytz.timezone('US/Eastern'), ambiguous='NaT').tz_convert(pytz.utc)
# Drop ambiguous times
NGG2_pd = NGG2_pd.dropna()

#%% Read in model log
runLog = pd.read_excel('C:/Users/arey/files/Projects/Newtown/Model Log NTC.xlsx', 'ModelLog')


#%% Load in Delft Results
dataPath = "FlowFM_map.nc"
histPath = "FlowFM_his.nc"

#read in time series
modelPlot = [105, 139, 141, 142, 143, 144]#, 106]

Delft_WL = [None] * (max(modelPlot)+1)
Delft_T_B = [None] * (max(modelPlot)+1)
Delft_S_B = [None] * (max(modelPlot)+1)
Delft_T_T = [None] * (max(modelPlot)+1)
Delft_S_T = [None] * (max(modelPlot)+1)


# Note, this uses the standard netcdf lib + xarray
for i in modelPlot:
    file_nc_hist = pl.Path(runLog['Run Location'][i],
                                 'FlowFM', 'dflowfm', '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()
    
    # Get observation stations
    #observations_pd=data_xr['stations'].to_dataframe()
    
    # Convert to Pandas
    Delft_WL[i] = data_xr.waterlevel.sel(stations=['NC310', 'NC313','NC316','NC318','EB043','EK108','West_WL']).to_pandas()

    #Temp
    Delft_T_B[i] = data_xr.temperature.sel(stations=['NC310', 'NC313','NC316','NC318','EB043','EK108'], laydim=0).to_pandas()
    Delft_T_T[i] = data_xr.temperature.sel(stations=['NC310', 'NC313','NC316','NC318','EB043','EK108'], laydim=9).to_pandas()
    
    #Salinity 
    Delft_S_B[i] = data_xr.salinity.sel(stations=['NC310', 'NC313','NC316','NC318','EB043','EK108'], laydim=0).to_pandas()
    Delft_S_T[i] = data_xr.salinity.sel(stations=['NC310', 'NC313','NC316','NC318','EB043','EK108'], laydim=9).to_pandas()
    
    # Put in UTC if needed
    if i == 104 or i == 105 or i >= 135:
        Delft_WL[i] = Delft_WL[i].tz_localize(
            pytz.timezone('EST')).tz_convert(pytz.utc)

        Delft_T_B[i] = Delft_T_B[i].tz_localize(pytz.timezone('EST')).tz_convert(pytz.utc)
        Delft_T_T[i] = Delft_T_T[i].tz_localize(pytz.timezone('EST')).tz_convert(pytz.utc)

        Delft_S_B[i] = Delft_S_B[i].tz_localize(pytz.timezone('EST')).tz_convert(pytz.utc)
        Delft_S_T[i] = Delft_S_T[i].tz_localize(pytz.timezone('EST')).tz_convert(pytz.utc)
    else:
        Delft_WL[i] = Delft_WL[i].tz_localize(
            pytz.timezone('UTC'))
 

#%% Plot Water Levels
i = 142

# Create Figure
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(18, 8), sharex=False, sharey=True)
fig.suptitle('Comparison of Predicted and Measured Water Level at NGG', fontsize=18)

# Plot Oberservations
NGG2_pd['water_surface_elevation'].multiply(0.3048).plot(label='FFG', color='k', linewidth=2)

# Plot Delft
Delft_WL[i]['West_WL'].plot(label='Delft', color='r', linewidth=2)

# Plot YSI
# YSI_df.loc[YSI_df['Station'] == 'NC318A', 'Depth'].add(-1.3).plot(
#     label='YSI', color='b', linewidth=2, ax=ax)

# Plot EFDC
# EFDC_dat['NC318A'].loc[:, 'WSE_m'].plot(label='EFDC', color='m', linewidth=2, ax=ax)

# Plot Telemac
# TM_dat['NC318'].loc[:, 'WSE_m'].plot(label='EFDC', color='m', linewidth=2, ax=ax)


# Set Time Bounds
# dateMin = pd.to_datetime(datetime.datetime(2014, 11, 1, 0, 0, 0, 0), utc=True)
# dateMax = pd.to_datetime(datetime.datetime(2014, 11, 13, 0, 0, 0, 0), utc=True)
dateMin = pd.to_datetime(datetime.datetime(2015, 7, 15, 0, 0, 0, 0), utc=True)
dateMax = pd.to_datetime(datetime.datetime(2015, 9, 1, 0, 0, 0, 0), utc=True)

# Set X Axis
ax.set_xlim(dateMin, dateMax)

# Set Y Axis
ax.set_ylim(-1.5, 2.5)

plt.show()

#%% Plot of Temperature AJMR
def resampleToObs(model, obs, obsMatchFreq):
    # Resample to observation time
    modelSampled = model.resample(obsMatchFreq).interpolate().reindex(
        obs.index)
    return modelSampled

def syncModelRMSE(model, obs):
    rmseOut = mean_squared_error(obs.loc[model.notna() & obs.notna()].dropna(
                ).values, model.loc[model.notna() & obs.notna()].dropna(
                ).values,squared=False)
    return rmseOut

degree_sign = u'\N{DEGREE SIGN}'


# Create Figure
fig, ax = plt.subplots(nrows=2, ncols=6, figsize=(16, 7), sharex=True, sharey=True)
# fig.suptitle('Comparison of Predicted and Measured Temperature at NC316', fontsize=18)

# Set Time Bounds for all axis since linked
dateMin = pd.to_datetime(datetime.datetime(2015, 7, 15, 0, 0, 0, 0), utc=True)
dateMax = pd.to_datetime(datetime.datetime(2015, 9, 1, 0, 0, 0, 0), utc=True)
# dateMin = pd.to_datetime(datetime.datetime(2015, 7, 28, 0, 0, 0, 0), utc=True)
# dateMax = pd.to_datetime(datetime.datetime(2015, 8, 7, 0, 0, 0, 0), utc=True)
# dateMin = pd.to_datetime(datetime.datetime(2014, 7, 1, 0, 0, 0, 0), utc=True)
# dateMax = pd.to_datetime(datetime.datetime(2015, 10, 1, 0, 0, 0, 0), utc=True)

# Set up Delft3D Variables
DelftInterp = [None] * (max(modelPlot)+1)
RMSE_D = [None] * (max(modelPlot)+1)


# Loop through stations and plot
for StatIDX, Stat in enumerate(Stations):
    ax[0, StatIDX].set_title(Stat)

    # Resample to 5 minute data, reindex to observed and calculate RMSD
    ObsDat = YSI_df.loc[YSI_df['Station'] == Stat + 'A', 'Temp']
    ObsDat = ObsDat.loc[(ObsDat.index > dateMin) & (ObsDat.index <= dateMax)]

    # Loop through Delft Runs
    for i in modelPlot:
        DelftInterp[i] = resampleToObs(Delft_T_T[i][Stat], ObsDat, '5T')
        RMSE_D[i] = syncModelRMSE(DelftInterp[i], ObsDat)

    EFDCInterp = resampleToObs(EFDC_dat[Stat].loc[:, 'TEM10'].resample('10T').nearest(), ObsDat, '5T')
    RMSE_E = syncModelRMSE(EFDCInterp, ObsDat)

    # TMInterp = resampleToObs(TM_dat[Stat].loc[:, 'TEM10'].resample('15T').nearest(), ObsDat, '5T')
    # RMSE_T = syncModelRMSE(TMInterp, ObsDat)

    # Plot Top Layer Temperature
    ObsDat.plot(label='YSI', color='k', linewidth=2, ax=ax[0, StatIDX])
    for i in modelPlot:
        ax[0, StatIDX].plot(Delft_WL[i].index, Delft_T_T[i][Stat], linestyle='-', label='Delft3D_' + str(i) + ':' + str(round(RMSE_D[i], 2)) + degree_sign + 'C')

    EFDC_dat[Stat].loc[:, 'TEM10'].plot(linewidth=1, ax=ax[0, StatIDX], label='EFDC: ' + str(round(RMSE_E, 2)) + degree_sign + 'C')
    # TM_dat[Stat].loc[:, 'TEM10'].plot(linewidth=1, ax=ax[0, StatIDX], label='Telemac: ' + str(round(RMSE_T, 2)) + degree_sign + 'C')

    ax[0, StatIDX].legend(loc='lower right')


    # Plot Bottom Layer Temperature
    # Resample to 5 minute data, reindex to observed and calculate RMSD
    ObsDat = YSI_df.loc[YSI_df['Station'] == Stat + 'C', 'Temp']
    ObsDat = ObsDat.loc[(ObsDat.index > dateMin) & (ObsDat.index <= dateMax)]

    # Loop through Delft Runs
    for i in modelPlot:
        DelftInterp[i] = resampleToObs(Delft_T_B[i][Stat], ObsDat, '5T')
        RMSE_D[i] = syncModelRMSE(DelftInterp[i], ObsDat)

    EFDCInterp = resampleToObs(EFDC_dat[Stat].loc[:, 'TEM1'].resample('10T').nearest(), ObsDat, '5T')
    RMSE_E = syncModelRMSE(EFDCInterp, ObsDat)

    # TMInterp = resampleToObs(TM_dat[Stat].loc[:, 'TEM1'].resample('15T').nearest(), ObsDat, '5T')
    # RMSE_T = syncModelRMSE(TMInterp, ObsDat)

    ObsDat.plot(label='YSI', color='k', linewidth=2, ax=ax[1, StatIDX])
    for i in modelPlot:
        ax[1, StatIDX].plot(Delft_WL[i].index, Delft_T_B[i][Stat], linestyle='-', label='Delft3D_' + str(i) + ':' + str(round(RMSE_D[i], 2)) + degree_sign + 'C')
    EFDC_dat[Stat].loc[:, 'TEM1'].plot(linewidth=1, ax=ax[1, StatIDX], label='EFDC: ' + str(round(RMSE_E, 2)) + degree_sign + 'C')
    # TM_dat[Stat].loc[:, 'TEM1'].plot(linewidth=1, ax=ax[1, StatIDX], label='Telemac: ' + str(round(RMSE_T, 2)) + degree_sign + 'C')

    # Add Legend and Set X Axis label
    ax[1, StatIDX].legend(loc='lower right')
    ax[1, StatIDX].set_xlabel('UTC Date')


# Set X Axis
ax[0, 0].set_xlim(dateMin, dateMax)
ax[0, 0].xaxis.set_major_formatter(
    mpl_dates.ConciseDateFormatter(ax[0, 0].xaxis.get_major_locator()))

# Set Y Axis
# ax[0, 0].set_ylim(18, 28)
ax[0, 0].set_ylim(-10, 30)
ax[0, 0].set_ylabel('Temperature Top (' + degree_sign + 'C)')
ax[1, 0].set_ylabel('Temperature Bottom (' + degree_sign + 'C)')

# Reduce spacing between subplots
plt.tight_layout()
plt.show()

# Save figure
# fig.savefig('C:/Users/arey/files/Projects/Newtown/TempSalFigs/Temperature_All_SalSens.png',
#             bbox_inches='tight', dpi=300)


#%% Plot Salinity AJMR
Stations=['NC310', 'NC313', 'NC316', 'NC318', 'EB043', 'EK108']

# Create Figure
fig, ax = plt.subplots(nrows=2, ncols=6, figsize=(16, 10), sharex=True, sharey=True)


# Set Time Bounds for all axis since linked
dateMin = pd.to_datetime(datetime.datetime(2015, 7, 15, 0, 0, 0, 0), utc=True)
dateMax = pd.to_datetime(datetime.datetime(2015, 9, 1, 0, 0, 0, 0), utc=True)
# dateMin = pd.to_datetime(datetime.datetime(2015, 7, 28, 0, 0, 0, 0), utc=True)
# dateMax = pd.to_datetime(datetime.datetime(2015, 8, 7, 0, 0, 0, 0), utc=True)

# Set up Delft3D Variables
DelftInterp = [None] * (max(modelPlot)+1)
RMSE_D = [None] * (max(modelPlot)+1)

# Loop through stations and plot
for StatIDX, Stat in enumerate(Stations):
    ax[0, StatIDX].set_title(Stat)

    # Resample to 5 minute data, reindex to observed and calculate RMSE
    ObsDat = YSI_df.loc[YSI_df['Station'] == Stat + 'A', 'Sal']
    ObsDat = ObsDat.loc[(ObsDat.index > dateMin) & (ObsDat.index <= dateMax)]

    # Loop through Delft Runs
    for i in modelPlot:
        DelftInterp[i] = resampleToObs(Delft_S_T[i][Stat], ObsDat, '5T')
        RMSE_D[i] = syncModelRMSE(DelftInterp[i], ObsDat)


    EFDCInterp = resampleToObs(EFDC_dat[Stat].loc[:, 'SAL10'].resample('10T').nearest(), ObsDat, '5T')
    RMSE_E = syncModelRMSE(EFDCInterp, ObsDat)
    # TMInterp = resampleToObs(TM_dat[Stat].loc[:, 'SAL10'].resample('15T').nearest(), ObsDat, '5T')
    # RMSE_T = syncModelRMSE(TMInterp, ObsDat)

    # Plot Top Layer Salinity
    ObsDat.plot(label='YSI', color='k', linewidth=2, ax=ax[0, StatIDX])
    for i in modelPlot:
        ax[0, StatIDX].plot(Delft_S_T[i].index, Delft_S_T[i][Stat], linestyle='-', label='Delft3D_' + str(i) + ':' + str(round(RMSE_D[i], 2)) + ' psu')

    EFDC_dat[Stat].loc[:, 'SAL10'].plot(linewidth=1, ax=ax[0, StatIDX], label='EFDC: ' + str(round(RMSE_E, 2)) + ' psu')
    # TM_dat[Stat].loc[:, 'SAL10'].plot(linewidth=1, ax=ax[0, StatIDX], label='Telemac: ' + str(round(RMSE_T, 2)) + ' psu')

    ax[0, StatIDX].legend(loc='lower right')

    # Resample to 5 minute data, reindex to observed and calculate RMSD
    ObsDat = YSI_df.loc[YSI_df['Station'] == Stat + 'C', 'Sal']
    ObsDat = ObsDat.loc[(ObsDat.index > dateMin) & (ObsDat.index <= dateMax)]

    for i in modelPlot:
        DelftInterp[i] = resampleToObs(Delft_S_B[i][Stat], ObsDat, '5T')
        RMSE_D[i] = syncModelRMSE(DelftInterp[i], ObsDat)

    EFDCInterp = resampleToObs(EFDC_dat[Stat].loc[:, 'SAL1'].resample('10T').nearest(), ObsDat, '5T')
    RMSE_E = syncModelRMSE(EFDCInterp, ObsDat)
    # TMInterp = resampleToObs(TM_dat[Stat].loc[:, 'SAL1'].resample('15T').nearest(), ObsDat, '5T')
    # RMSE_T = syncModelRMSE(TMInterp, ObsDat)

    # Plot Bottom Layer Salinity
    ObsDat.plot(label='YSI', color='k', linewidth=2, ax=ax[1, StatIDX])
    for i in modelPlot:
        ax[1, StatIDX].plot(Delft_S_B[i].index, Delft_S_B[i][Stat], linestyle='-', label='Delft3D_' + str(i) + ':' + str(round(RMSE_D[i], 2)) + ' psu')
    EFDC_dat[Stat].loc[:, 'SAL1'].plot(linewidth=1, ax=ax[1, StatIDX], label='EFDC: ' + str(round(RMSE_E, 2)) + ' psu')
    # TM_dat[Stat].loc[:, 'SAL1'].plot(linewidth=1, ax=ax[1, StatIDX], label='Telemac: ' + str(round(RMSE_T, 2)) + 'psu')

    # Add Legend and Set X Axis label
    ax[1, StatIDX].legend(loc='lower right')
    ax[1, StatIDX].set_xlabel('UTC Date')

# Set X Axis
ax[0, 0].set_xlim(dateMin, dateMax)
ax[0, 0].xaxis.set_major_formatter(
    mpl_dates.ConciseDateFormatter(ax[0, 0].xaxis.get_major_locator()))

# Set Y Axis
ax[0, 0].set_ylim(4, 27)
ax[0, 0].set_ylabel('Salinity Top (psu)')
ax[1, 0].set_ylabel('Salinity Bottom (psu)')

# Reduce spacing between subplots
plt.tight_layout()
plt.show()

# Save figure
# fig.savefig('C:/Users/arey/files/Projects/Newtown/TempSalFigs/Salinity_2015_Zoom_SensTests.png',
#             bbox_inches='tight', dpi=300)

