# -*- coding: utf-8 -*-
#==========================================================#
#Last Updated - June 2018
#@author: z5025317 Valentin Heimhuber
#code for creating climate prioritization plots for NARCLIM variables. 
#Inputs: Uses CSV files that contain all 12 NARCLIM model runs time series for 1 grid cell created with: P1_NARCliM_NC_to_CSV_CCRC_SS.py

#==========================================================#
#Load packages
#==========================================================#
import numpy as np
import os
import pandas as pd
import glob
import matplotlib
import matplotlib.pyplot as plt
from datetime import datetime
from datetime import timedelta
from matplotlib.backends.backend_pdf import PdfPages
from ggplot import *
matplotlib.style.use('ggplot')

#==========================================================#
#Input parameters
#==========================================================#
#set scenario codes and beginning and end years and corresponding scenario code
fs=['Hwq003', 'Hwq005']
scenario_real_names = ['Present', 'Far_Future']

variable = 'elev' #depth or vel

startyear=1999
endyear=1999
#year=range(startyear, endyear+1)

#set directory path for output files
output_directory = 'H:/WRL_Projects/Hunter_CC_Modeling/Module_6/03_Results/Output/Postprocessed/Figures'

#set input directories and data
#postprocessed RMA results
RMA_data_csv = 'H:/WRL_Projects/Hunter_CC_Modeling/Module_6/03_Results/Output/Postprocessed/Compound_data/Scn__Hwq003_Hwq005_WQ_sal_1999_1999compound.csv'
#tidal boundary data
tides_csv='C:/Users/z5025317/OneDrive - UNSW/Hunter_CC_Modeling/07_Modelling/01_Input/Tide Generator/Tidal Simulation Python Script (BMM)/Tides_1990_2010.csv'
#CSV file of mesh nodes and corresponding chainages
nodes_csv = 'H:/WRL_Projects/Hunter_CC_Modeling/Module_6/03_Results/Chainages/Hunter_nodes.csv' 
#River flow boundary
Flow_csv = 'H:/WRL_Projects/Hunter_CC_Modeling/Module_6/01_Input/BCGeneration/Scenarios/Calibration/Calibration_Greta.csv' 
#==========================================================#

#==========================================================#
#set plot parameters
ALPHA_figs = 1
font = {'family' : 'sans-serif',
        'weight' : 'normal',
        'size'   : 14}
matplotlib.rc('font', **font)
#==========================================================#


#========================================================#
#automated part of the code doing the data extraction
#==========================================================

#read csv file with extracted RMA data
RMA.df = pd.read_csv(RMA_data_csv, parse_dates=True, index_col=0)

#read csv file with nodes to extract data from
node = pd.read_csv(nodes_csv)['Hunter'].values
chainages = pd.read_csv(nodes_csv)['x_km'].values

#Load boundary flow data
Flow_df = pd.read_csv(Flow_csv, parse_dates=True, index_col=0)['Q[ML/d]']
Flow_df.columns = ['Q[ML/d]']
Flow_df = Flow_df[datetime.strptime(str(startyear) + ' 07 13', '%Y %m %d').date():datetime.strptime(str(2005) + ' 07 19', '%Y %m %d').date()]
flowannual = Flow_df.resample('A').mean()
Flow_df.resample('A').mean().plot()

#Load tide boundary data
tides_df = pd.read_csv(tides_csv, parse_dates=True, index_col=0)
tides_df.columns = ['tide']


#append to summary data frame
if f in ['HCC010', 'HCC042']:
    scname = 'Present'
if f == 'HCC011':
    scname = 'Near_Future'
    df.index =  df.index  -  (datetime.strptime('2025 01 01', '%Y %m %d').date() - datetime.strptime('1995 01 01', '%Y %m %d').date())
if f == 'HCC012':
    scname = 'Far_Future'
    df.index =  df.index  -  (datetime.strptime('2065 01 01', '%Y %m %d').date() - datetime.strptime('1995 01 01', '%Y %m %d').date())
if f in ['HCC040', 'HCC041']:
    scname = 'Present_Resto'
df.columns = df.columns+'_'+ scname

#cut down the summary df to common years
Summary_df = Summary_df[datetime.strptime(str(startyear) + ' 07 13', '%Y %m %d').date():datetime.strptime(str(endyear) + ' 07 19', '%Y %m %d').date()]
Summary_df.tail()
#Summary_df.columns = chainages
#type(Summary_df.index)

#cut down the summary df to common years
Summary_df = Summary_df[datetime.strptime(str(startyear) + ' 01 05', '%Y %m %d').date():datetime.strptime(str(endyear) + ' 12 31', '%Y %m %d').date()]
Summary_df.tail()


variables = ['Sal']

Present_df = Summary_df.filter(regex='Present')
Present_df.columns

Far_future_df = Summary_df.filter(regex='Far')
Far_future_df.columns



#==========================================================#
#Tidal analysis
#==========================================================#
#read in tide data
Plot_scenario = 'Flood'
ylims = [-1.1,8]
tides_df = tides_df[datetime.strptime(str(startyear) + ' 07 13', '%Y %m %d').date():datetime.strptime(str(endyear) + ' 07 19', '%Y %m %d').date()]

central_tides_df = tides_df[datetime.strptime(str(startyear) + ' 07 16', '%Y %m %d').date():datetime.strptime(str(endyear) + ' 07 16', '%Y %m %d').date()]
maxtide_index = central_tides_df['tide'].idxmax()
tides_df['tide'].max()
end_maxtide_index =  maxtide_index + timedelta(days=0.4)
beg_maxtide_index =  maxtide_index - timedelta(days=0.2)
short_tides_df = tides_df[beg_maxtide_index: end_maxtide_index]
short_tides_df.plot(cmap='viridis',  legend=True)
mintide_index = short_tides_df['tide'].idxmin()

#plot the minimum and maximum elevation along the estuary occuring within a single tidal cycle
Short_tide_elev_df = Summary_df[min(short_tides_df.index):max(short_tides_df.index)]
max_min_1cycle_elev_0slr = pd.concat([Short_tide_elev_df.filter(regex='Present').max(), Short_tide_elev_df.filter(regex='Present').min()], axis=1, join='inner')  
max_min_1cycle_elev_09slr = pd.concat([Short_tide_elev_df.filter(regex='Far').max(), Short_tide_elev_df.filter(regex='Far').min()], axis=1, join='inner')
max_min_1cycle_elev_0slr.index = chainages
max_min_1cycle_elev_09slr.index = chainages
max_min_1cycle_elev = pd.concat([max_min_1cycle_elev_0slr, max_min_1cycle_elev_09slr], axis=1, join='outer')
max_min_1cycle_elev.columns = ['MaxElev 0 SLR','MinElev 0 SLR','MaxElev 0.9 SLR','MinElev 0.9 SLR' ]
max_min_1cycle_elev.plot(cmap='viridis',  legend=True)

#plot the elevation along the estuary at exactly the time of high and low tide
max_min_extime_elev_0slr = pd.concat([Summary_df.filter(regex='Present').loc[maxtide_index], Summary_df.filter(regex='Present').loc[mintide_index]], axis=1, join='outer')
max_min_extime_elev_09slr = pd.concat([Summary_df.filter(regex='Far').loc[maxtide_index], Summary_df.filter(regex='Far').loc[mintide_index]], axis=1, join='outer')
max_min_extime_elev_0slr.index = chainages
max_min_extime_elev_09slr.index = chainages
max_min_extime_elev = pd.concat([max_min_extime_elev_0slr, max_min_extime_elev_09slr], axis=1, join='outer')
max_min_extime_elev.columns = ['HigT Elev 0SLR','LowT Elev 0SLR','HigT Elev 0.9 SLR','LowT Elev 0.9 SLR' ]
max_min_extime_elev.plot(cmap='viridis',  legend=True)

png_out_name = png_output_directory + variable + '_high_low_tide_analysis_' + Plot_scenario + '.pdf' 
fig = plt.figure(figsize=(17,20))
ax=plt.subplot(3,1,1) 
plt.title('Max and min water surface elevation along the estuary over one tidal cycle') 
ax.set_prop_cycle(color=['slateblue', 'slateblue', 'tomato', 'tomato'], linestyle=['--', '-', '--', '-'], lw=[2,2,2,2])
max_min_1cycle_elev.plot( legend=False, ax=ax)
plt.axhline(y=0, color='slateblue', linestyle='--', lw=1, alpha=0.5)  
plt.axhline(y=0.9, color='tomato', linestyle='--', lw=1, alpha=0.5) 
plt.ylim(ylims)
plt.ylabel("Water level [m AHD]")
plt.xlabel("Distance from ocean boundary [km]")
# Put a legend to the right of the current axis
#lgd = ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))  
ax.legend()                    
#plot tides
ax=plt.subplot(3,1,2)
plt.title('Max and min water surface elevation along the estuary at time of high & low tide') 
ax.set_prop_cycle(color=['slateblue', 'slateblue', 'tomato', 'tomato'], linestyle=['--', '-', '--', '-'], lw=[2,2,2,2])
max_min_extime_elev.plot(legend=False, ax=ax)
plt.ylabel("Water level [m AHD]")
plt.ylim(ylims)
plt.xlabel("Distance from ocean boundary [km]")
#lgd = ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))   
ax.legend()    
#plot tides
ax=plt.subplot(3,1,3)
plt.ylim(ylims)
plt.title('Tide at open ocean boundary') 
short_tides_df.plot(cmap='viridis',  legend=False, ax=ax)
ax.axvline(maxtide_index, color='black', linestyle='--')
ax.axvline(mintide_index, color='black', linestyle='-')
plt.xlabel("Time")
plt.ylabel("Water level [m AHD]")
#lgd = ax.legend(loc='center left', bbox_to_anchor=(1, 0.5)) 
ax.legend()     
fig.tight_layout()
fig.savefig(png_out_name)
plt.close()




#==========================================================#
#animate 2 tidal cycles baseline scenario:
#==========================================================#
Six_day_tide_elev_df = Summary_df.filter(regex='Present')[maxtide_index - timedelta(days=2): maxtide_index + timedelta(days=3)]
Six_day_tide_elev_df.columns = chainages
Six_day_tides_df = tides_df[maxtide_index - timedelta(days=2): maxtide_index + timedelta(days=3)]
Six_day_tides_df.plot(cmap='viridis',  legend=True)

Six_day_flow_df = Flow_df[maxtide_index - timedelta(days=2): maxtide_index + timedelta(days=3)]
Six_day_flow_df = pd.concat([Six_day_tides_df,  Six_day_flow_df], axis=1, join='outer')['Q[ML/d]'].interpolate()
Six_day_flow_df.plot(cmap='viridis',  legend=True)

    
ii=1
for index in Six_day_tide_elev_df.index:
    #index = Six_day_tide_elev_df.index[1]
    png_out_path =  output_directory + '/TidalCycle_' + variable + '_' + Plot_scenario +  '3/'
    if not os.path.exists(png_out_path):
        os.makedirs(png_out_path)
    #png_out_name = png_out_path + str(index) + 'longitudinalXS.png' 
    png_out_name = png_out_path + str(ii) + variable +  '_longitudinalXS.png' 
    fig = plt.figure(figsize=(17,20))
    ax=plt.subplot(3,1,1) 
    plt.title('Water surface elevation along the estuary') 
    Six_day_tide_elev_df.loc[index].plot(color='blue',  legend=False)
    #plt.axhline(y=0, color='r', linestyle='--', lw=1, alpha=0.5) 
    plt.axhline(tides_df.loc[index].values, color='black',lw=0.5, linestyle='--')
    plt.ylim(ylims)
    plt.ylabel("Water level [m AHD]")
    plt.xlabel("Distance from ocean boundary [km]")
    #plot tides
    ax=plt.subplot(3,1,2)
    plt.ylim(-1.1,1.3)
    plt.title('Tide at open ocean boundary') 
    Six_day_tides_df.plot(cmap='viridis',  legend=False, ax=ax)
    ax.axvline(index, color='green', linestyle='--')
    plt.xlabel("Time")
    plt.ylabel("Water level [m AHD]")
    
    ax=plt.subplot(3,1,3)
    #plt.ylim(ylims)
    plt.title('Tide at open ocean boundary') 
    Six_day_flow_df.plot(color='black',  legend=False, ax=ax)
    ax.axvline(index, color='green', linestyle='--')
    plt.xlabel("Time")
    plt.ylabel("Hunter Inflow [ML/day]")

    fig.tight_layout()
    fig.savefig(png_out_name)
    plt.close()
    ii=ii+1
        
        
#==========================================================#
#animate 2 tidal cycles baseline + SLR scenario:
#==========================================================#
Six_day_tide_elev_df_0slr = Summary_df.filter(regex='Present')[maxtide_index - timedelta(days=1): maxtide_index + timedelta(days=1)]
Six_day_tide_elev_df_0slr.columns = chainages
Six_day_tide_elev_df_09slr = Summary_df.filter(regex='Far')[maxtide_index - timedelta(days=1): maxtide_index + timedelta(days=1)]
Six_day_tide_elev_df_09slr.columns = chainages
Six_day_tides_df = tides_df[maxtide_index - timedelta(days=1): maxtide_index + timedelta(days=1.5)]
Six_day_tides_df.plot(cmap='viridis',  legend=True)

ii=0
for index in Six_day_tide_elev_df_0slr.index:
    ii=ii+1
    #index = Six_day_tide_elev_df.index[1]
    png_out_path =  output_directory + '/TidalCycle_SLR_' + variable + '_' + Plot_scenario +   '/'
    if not os.path.exists(png_out_path):
        os.makedirs(png_out_path)
    #png_out_name = png_out_path + str(index) + 'longitudinalXS.png' 
    png_out_name = png_out_path + str(ii) + variable +  '_longitudinalXS.png' 
    max_min_extime_elev_both = pd.concat([Six_day_tide_elev_df_0slr.loc[index],  Six_day_tide_elev_df_09slr.loc[index]], axis=1, join='outer')
    max_min_extime_elev_both.columns = ['Present Day','0.9m SLR']
    
    fig = plt.figure(figsize=(17,12))
    ax=plt.subplot(2,1,1) 
    plt.ylim(ylims)
    plt.title('Water surface elevation along the estuary') 
    ax.set_prop_cycle(color=['slateblue','tomato'], linestyle=['-',  '-'], lw=[2,2])
    max_min_extime_elev_both.plot(ax=ax)
    ax.legend(loc=1)
    plt.axhline(y=0, color='r', linestyle='--', lw=1, alpha=0.5)  
    plt.ylim(-1.1,2.5)
    plt.ylabel("Water level [m AHD]")
    plt.xlabel("Distance from ocean boundary [km]")
    
    #plot tides
    ax=plt.subplot(2,1,2)
    plt.ylim(-1.1,1.5)
    plt.title('Tide at open ocean boundary') 
    Six_day_tides_df.plot(cmap='viridis',  legend=False, ax=ax)
    ax.axvline(index, color='green', linestyle='--')
    plt.xlabel("Time")
    plt.ylabel("Water level [m AHD]")
    fig.tight_layout()
    fig.savefig(png_out_name)
    plt.close()
 
        
































for NODE in node:
    NODE = str(NODE)
    #=========#
    #comparison time series plot at node
    #=========#
    #out name
    png_out_file_name = comparison_code + '_' +  variable  + '_' + NODE + '_time_series.png'
    png_out_path = png_output_directory + '/' + png_out_file_name
    #
    plotin_df = Summary_df.filter(regex=NODE)
    #prepare data frame for plot
    fig = plt.figure(figsize=(14,18))
    ax=plt.subplot(4,1,1) 
    #start and end dat of plots
    plt.title(variable  + '_' + NODE + ' - time series')
    start_date = datetime(int('1996'), int('02'), int('01'))
    end_date = datetime(int('1996'), int('02'), int('5'))
    plotin_df1=plotin_df[start_date:end_date]
    ymin = min(plotin_df1.min()) - 0.01 *min(plotin_df1.min())
    ymax = max(plotin_df1.max()) + 0.01 *max(plotin_df1.max())  
    plotin_df1.plot(ylim=(ymin,ymax) ,legend=False, ax=ax) 
    #plt.legend( loc=1)    
    plt.xticks([])
    
    ax=plt.subplot(4,1,2) 
    start_date = datetime(int('1996'), int('02'), int('01'))
    end_date = datetime(int('1996'), int('02'), int('3'))
    plotin_df2=plotin_df[start_date:end_date]
    ymin = min(plotin_df2.min()) - 0.1 *min(plotin_df2.min())
    ymax = max(plotin_df2.max()) + 0.1 *max(plotin_df2.max()) 
    plotin_df2.plot(ylim=(ymin,ymax),legend=False, ax=ax)
    #plt.legend( loc=1)
    ax.xaxis.grid(False)
    
    ax=plt.subplot(4,1,3)
    plt.title(variable  + '_' + NODE + ' - monthly maximum')
    Monthly_max = plotin_df.resample('M').max()
    Monthly_max.plot(ax=ax,legend=False)
    plt.legend( loc=1)
    ax.xaxis.grid(False)
    
    ax=plt.subplot(4,1,4)
    plt.title(variable  + '_' + NODE + ' - monthly mean')
    Monthly_mean = plotin_df.resample('M').mean()
    Monthly_mean.plot(ax=ax)
    #plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
    #   ncol=2, mode="expand", borderaxespad=0.)
    #plt.legend( loc=1)
    ax.xaxis.grid(False)

    #ax.patch.set_alpha(ALPHA_figs)
    fig.patch.set_alpha(ALPHA_figs)
    fig.tight_layout()
    fig.savefig(png_out_path)
    plt.close()
    #=========#