import numpy as np
import matplotlib.pyplot as plt
from salishsea_tools import evaltools as et, viz_tools
import cmocean as cmo
import datetime as dt
import netCDF4 as nc
import matplotlib.dates as mdates
import matplotlib as mpl
import matplotlib.gridspec as gridspec
from matplotlib.colors import LogNorm
import cmocean
import pandas as pd
import pickle
import os
import gsw

%matplotlib inline

# path to model files:
PATH= '/results/SalishSea/nowcast-green.201812/'

# start and end dates for analysis:
start_date = dt.datetime(2015,1,1)
end_date = dt.datetime(2016,1,1)

# number of days per model file:
flen=1

# dictionary mapping desired model variables to the file types where they are found
filemap={'nitrate':'ptrc_T','silicon':'ptrc_T','ammonium':'ptrc_T','diatoms':'ptrc_T','ciliates':'ptrc_T',
         'flagellates':'ptrc_T','vosaline':'grid_T','votemper':'grid_T'}

# dictionary mapping model file types to their time resolution in hours (1 is hourly files, 24 is daily)
fdict={'ptrc_T':1,'grid_T':1}

# results format
# -- nowcast: files like 01jan15/SalishSea_1h_20150101_20150101_ptrc_T.nc
# -- long: files like SalishSea_1h_20150206_20150804_ptrc_T_20150427-20150506.nc, all in one directory
namfmt='nowcast'

def interpCTDvar(sta,yr,yd,ztarget,ctdvar):
    ctdlocs=(dfCTD.Station==sta)&(dfCTD.Year==yr)&(dfCTD.YD==yd)
    if np.sum(ctdlocs)==0:
        print(f'Warning: Station {sta}, Year {yr}, year day {yd} not found in dfCTD')
        return np.nan
    else:
        val=np.interp(ztarget,dfCTD.loc[ctdlocs,['Z']].values.flatten(),
                  dfCTD.loc[ctdlocs,[ctdvar]].values.flatten())
        return val

# load new datetime info
dfTime=pd.read_excel('/ocean/eolson/MEOPAR/obs/WADE/WDE_Data/OlsonSuchyAllen_UBC_PDR_P003790-010721.xlsx',
                    engine='openpyxl',sheet_name='EventDateTime')
# drop duplicate rows, taking first time
dfTime.drop_duplicates(subset=['FlightDate','SiteCode'],keep='first',inplace=True)
dfTime['dtPac']=[dt.datetime.combine(idate, itime) for idate, itime \
         in zip(dfTime['FlightDate'],dfTime['TimeDown \n(Local - PST or PDT)'])]
dfTime['dtUTC']=[et.pac_to_utc(ii) for ii in dfTime['dtPac']]
# join to date/time
dfTime['dtUTC']=[et.pac_to_utc(dt.datetime.combine(idate,itime)) for idate,itime in \
                zip(dfTime['FlightDate'],dfTime['TimeDown \n(Local - PST or PDT)'])]
dfTime2=dfTime.loc[:,['FlightDate','SiteCode','dtUTC']]

dfs=dict()
for year in range(2015,2019):
    PATH= '/results2/SalishSea/nowcast-green.201905/'
    datadir='/ocean/eolson/MEOPAR/obs/WADE/ptools_data/ecology'
    dfSta=pickle.load(open(os.path.join(datadir,'sta_df.p'),'rb'))
    dfCTD0=pickle.load(open(os.path.join(datadir,f'Casts_{str(year)}.p'),'rb'))
    dfs0=pd.merge(left=dfSta,right=dfCTD0,how='right',
             left_on='Station',right_on='Station')
    dfs0.rename(columns={'Latitude':'Lat','Longitude':'Lon'},inplace=True)
    # Calculate Absolute (Reference) Salinity (g/kg) and Conservative Temperature (deg C) from 
    # Salinity (psu) and Temperature (deg C):
    dfs0['Z']=-1*dfs0['Z']
    press=gsw.p_from_z(-1*dfs0['Z'],dfs0['Lat'])
    dfs0['SA']=gsw.SA_from_SP(dfs0['Salinity'],press,
                               dfs0['Lon'],dfs0['Lat'])
    dfs0['CT']=gsw.CT_from_t(dfs0['SA'],dfs0['Temperature'],press)
    # merge time
    dfs[year]=pd.merge(left=dfs0,right=dfTime2,how='left',
                      left_on=['Date','Station'],right_on=['FlightDate','SiteCode'])
    # if no time available, guess midday
    dfs[year]['dtUTC']=[iiD+dt.timedelta(hours=20) if pd.isnull(iiU) else iiU \
                    for iiU,iiD in zip(dfs[year]['dtUTC'],dfs[year]['Date'])]
    #dfs[year]['Year']=[ii.year for ii in dfs[year]['dtUTC']]
    #dfs[year]['YD']=et.datetimeToYD(dfs[year]['dtUTC'])

dfAll.to_csv('/data/eolson/results/MEOPAR/oldDBs/WADECTD.csv',index=False)

#Chl:N ratio used later in plots: grab from appropriate namelist
mod_chl_N=et.getChlNRatio(basedir=PATH,nam_fmt=namfmt)

# create dictionary of dataframe views by year
datyr=dict()
yy=np.array([ii.year for ii in data['dtUTC']])
yys=np.unique(yy)
for yr in yys:
    datyr[yr]=data.loc[yy==yr]