# Starter code for week 6 illustrating how to build an AR model 
# and plot it

# %%
# Import the modules we will use
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
import datetime
#note you may need to do pip install for sklearn

# %%
# ** MODIFY **
# Set the file name and path to where you have stored the data
filename = 'streamflow_week6.txt'
filepath = os.path.join('../data', filename)
print(os.getcwd())
print(filepath)

# %%
#Read the data into a pandas dataframe
data=pd.read_table(filepath, sep = '\t', skiprows=30,
        names=['agency_cd', 'site_no', 'datetime', 'flow', 'code'],
        parse_dates=['datetime']
        )

# Expand the dates to year month day
data['year'] = pd.DatetimeIndex(data['datetime']).year
data['month'] = pd.DatetimeIndex(data['datetime']).month
data['day'] = pd.DatetimeIndex(data['datetime']).dayofweek
data['dayofweek'] = pd.DatetimeIndex(data['datetime']).dayofweek

# Aggregate flow values to weekly 
flow_weekly = data.resample("W", on='datetime').mean()

# %%
# Building an autoregressive model 
# You can learn more about the approach I'm following by walking 
# Through this tutorial
# https://realpython.com/linear-regression-in-python/

# Step 1: setup the arrays you will build your model on
# This is an autoregressive model so we will be building
# it based on the lagged timeseries

flow_weekly['flow_tm1'] = flow_weekly['flow'].shift(1)
flow_weekly['flow_tm2'] = flow_weekly['flow'].shift(2)

# %% Step 2 - pick what portion of the time series you want to use as training data
# here I'm grabbing the first 800 weeks 
# Note1 - dropping the first two weeks since they wont have lagged data
# to go with them  

train = flow_weekly[2:800][['flow', 'flow_tm1', 'flow_tm2']]
test = flow_weekly[800:][['flow', 'flow_tm1', 'flow_tm2']]

# %% Step 3: Fit a linear regression model using sklearn 
model = LinearRegression()
x=train['flow_tm1'].values.reshape(-1,1) #See the tutorial to understand the reshape step here 
y=train['flow'].values
model.fit(x,y)

#Look at the results
# r^2 values
r_sq = model.score(x, y)
print('coefficient of determination:', np.round(r_sq,2))

#print the intercept and the slope 
print('intercept:', np.round(model.intercept_, 2))
print('slope:', np.round(model.coef_, 2))

# %% Step 4 Make a prediction with your model 
# Predict the model response for a  given flow value
q_pred_train = model.predict(train['flow_tm1'].values.reshape(-1,1))
q_pred_test = model.predict(test['flow_tm1'].values.reshape(-1,1))

#altrenatievely you can calcualte this yourself like this: 
q_pred = model.intercept_ + model.coef_ * train['flow_tm1']

# you could also predict the q for just a single value like this
last_week_flow = 500
prediction = model.intercept_ + model.coef_ * last_week_flow


# %%
# Another example but this time using two time lags as inputs to the model 
model2 = LinearRegression()
x2=train[['flow_tm1','flow_tm2']]
model2.fit(x2,y)
r_sq = model2.score(x2, y)
print('coefficient of determination:', np.round(r_sq,2))
print('intercept:', np.round(model2.intercept_, 2))
print('slope:', np.round(model2.coef_, 2))

# generate preditions with the funciton
q_pred2_train = model2.predict(train[['flow_tm1', 'flow_tm2']])

# or by hand
q_pred2 = model2.intercept_   \
         + model2.coef_[0]* train['flow_tm1'] \
         +  model2.coef_[1]* train['flow_tm2'] 

# %% 
# Here are some examples of things you might want to plot to get you started:

# 1. Timeseries of observed flow values
# Note that date is the index for the dataframe so it will 
# automatically treat this as our x axis unless we tell it otherwise
fig, ax = plt.subplots()
ax.plot(flow_weekly['flow'], label='full')
ax.plot(train['flow'], 'r:', label='training')
ax.set(title="Observed Flow", xlabel="Date", 
        ylabel="Weekly Avg Flow [cfs]",
        yscale='log')
ax.legend()
# an example of saving your figure to a file
#fig.set_size_inches(5,3)
#fig.savefig("Observed_Flow.png")

#2. Time series of flow values with the x axis range limited
fig, ax = plt.subplots()
ax.plot(flow_weekly['flow'], label='full')
ax.plot(train['flow'], 'r:', label='training')
ax.set(title="Observed Flow", xlabel="Date", ylabel="Weekly Avg Flow [cfs]",
        yscale='log', xlim=[datetime.date(2000, 1, 26), datetime.date(2014, 2, 1)])
ax.legend()


# 3. Line  plot comparison of predicted and observed flows
fig, ax = plt.subplots()
ax.plot(train['flow'], color='grey', linewidth=2, label='observed')
ax.plot(train.index, q_pred_train, color='green', linestyle='--', 
        label='simulated')
ax.set(title="Observed Flow", xlabel="Date", ylabel="Weekly Avg Flow [cfs]",
        yscale='log')
ax.legend()

# 4. Scatter plot of t vs t-1 flow with log log axes
fig, ax = plt.subplots()
ax.scatter(train['flow_tm1'], train['flow'], marker='p',
              color='blueviolet', label='obs')
ax.set(xlabel='flow t-1', ylabel='flow t', yscale='log', xscale='log')
ax.plot(np.sort(train['flow_tm1']), np.sort(q_pred_train), label='AR model')
ax.legend()

# 5. Scatter plot of t vs t-1 flow with normal axes
fig, ax = plt.subplots()
ax.scatter(train['flow_tm1'], train['flow'], marker='p',
              color='blueviolet', label='observations')
ax.set(xlabel='flow t-1', ylabel='flow t')
ax.plot(np.sort(train['flow_tm1']), np.sort(q_pred_train), label='AR model')
ax.legend()

plt.show()

# %% Q1. AR model that you ended up building, including (1) what you are using as your prediction variables, 
# (2) the final equation for your model and (3) what you used as your testing and training periods.
# In your discussion please include graphical outputs that support why you made 
# the decisions you did with your model.
fl_wk = data.resample("W", on='datetime').mean()
# 1st step
fl_wk['flow_tm1'] = fl_wk['flow'].shift(1)
fl_wk['flow_tm2'] = fl_wk['flow'].shift(2)

# %% Plot historical data
fig, ax = plt.subplots()
ax.plot(fl_wk.index, fl_wk["flow"], color='cyan', linestyle='-', 
        label='Observed')
ax.set(title=" Flow since 1989", xlabel="Date", ylabel="Avg Flow [cfs]",#,
       yscale='log')
ax.legend()
fig.savefig("Historical.png")

# %%
# 2nd step
# Aggregate flow values to weekly 
lw = fl_wk["flow"].size
tr = fl_wk[lw-50:lw-20][['flow', 'flow_tm1', 'flow_tm2']] # weeks for training (2019 and 2020 data)
ts = fl_wk[lw-20:][['flow', 'flow_tm1', 'flow_tm2']] # weeks to test (last 5 months of 2020)

# 3rd step 
mw = LinearRegression()
xw=ts['flow_tm1'].values.reshape(-1,1) 
yw=ts['flow'].values
mw.fit(xw,yw)

#Look at the results
# r^2 values
rw_sq = mw.score(xw, yw)
print('coefficient of determination:', np.round(rw_sq,2))

#print the intercept and the slope (Final equation) y=mx+b
b = np.round(mw.intercept_, 2) #b
m = np.round(mw.coef_, 2) #m
print('Final equation is y = :',m[:1],'x + ',b)

# 4th step prediction with model 
# Predict the model response for a  given flow value
q_p_tr = mw.predict(tr['flow_tm1'].values.reshape(-1,1))
q_p_ts = mw.predict(ts['flow_tm1'].values.reshape(-1,1))

#altrenatievely you can calcualte this yourself like this: 
q_p = mw.intercept_ + mw.coef_ * tr['flow_tm1']

# %% Q2. Provide an analysis of your final model performance. 
# This should include at least one graph that shows the historical vs predicted streamflow and 
# some discussion of qualitatively how you think your model is good or bad.

#2. Time series of flow values with the x axis range limited

fig, ax = plt.subplots()
ax.plot(ts['flow'], color='blue',linestyle='-', linewidth=2, label='observed')
ax.plot(ts.index, q_p_ts, color='red', linestyle='--', 
        label='simulated')
ax.set(title="Observed Flow", xlabel="Date", ylabel="Weekly Avg Flow [cfs]")
       # yscale='log')
ax.legend()
fig.set_size_inches(5,3)
fig.savefig("Observed_simulated_Flow.png")

fig, ax = plt.subplots()
ax.scatter(ts['flow_tm1'], ts['flow'], marker='p',
              color='blueviolet', label='observations')
ax.set(xlabel='flow t-1', ylabel='flow t')
ax.plot(np.sort(ts['flow_tm1']), np.sort(q_p_ts), label='AR model')
ax.legend()

plt.show()

fig.set_size_inches(5,3)
fig.savefig("Scatter.png")

fig, ax = plt.subplots()
ax.plot(tr['flow'], color='blue',linestyle='-', linewidth=2, label='observed')
ax.plot(tr.index, q_p_tr, color='red', linestyle='--', 
        label='simulated')
ax.set(title="Observed Flow", xlabel="Date", ylabel="Weekly Avg Flow [cfs]")
       # yscale='log')
ax.legend()
fig.set_size_inches(5,3)
fig.savefig("Observed_simulated_Flow_hst.png")

# %% Q3. provide discussion on what you actually used for your forecast. 
# Did you use your AR model, why or why not? If not how did you generate your forecast this week?

# I think the model estimates somehow the behaviour that the observed data follows,
# however it seems to underestimate peak values. Also, it seems that there is 
# a lag in the simulated values. Thsi conclusion can be obtained from both graphs

# %% Weekly prediction

lw_t = flow_weekly['flow'][lw-20].min()
pr_lw_t = mw.intercept_ + mw.coef_ * lw_t
print('Test:',pr_lw_t.round(2))

lw_f = flow_weekly['flow'][lw-10].min()
pr_lw_f = mw.intercept_ + mw.coef_ * lw_f
print('Week 1 forecast using model:',pr_lw_f.round(2))

lw_f2= flow_weekly['flow'][lw-20].min()
pr_lw_f2 = mw.intercept_ + mw.coef_ * lw_f2
print('Week 2 forecast using model:',pr_lw_f2.round(2))

# %% Seasonal forecast

fw_ag = flow_weekly['flow'][lw-62:lw-59].mean()
pr_lt1 = mw.intercept_ + mw.coef_ * fw_ag
print('Wlt1 forecast using model:',pr_lt1.round(2))

fw_lt2 = flow_weekly['flow'][lw-61:lw-58].mean()
pr_lt2 = mw.intercept_ + mw.coef_ * fw_lt2
print('Wlt2 forecast using model:',pr_lt2.round(2))

fw_lt3 = flow_weekly['flow'][lw-60:lw-57].mean()
pr_lt3 = mw.intercept_ + mw.coef_ * fw_lt3
print('Wlt3 forecast using model:',pr_lt3.round(2))

fw_lt4 = flow_weekly['flow'][lw-59:lw-56].mean()
pr_lt4 = mw.intercept_ + mw.coef_ * fw_lt4
print('Wlt4 forecast using model:',pr_lt4.round(2))

fw_lt5 = flow_weekly['flow'][lw-58:lw-55].mean()
pr_lt5 = mw.intercept_ + mw.coef_ * fw_lt5
print('Wlt5 forecast using model:',pr_lt5.round(2))

fw_lt6 = flow_weekly['flow'][lw-57:lw-54].mean()
pr_lt6 = mw.intercept_ + mw.coef_ * fw_lt6
print('Wlt6 forecast using model:',pr_lt6.round(2))

fw_lt7 = flow_weekly['flow'][lw-56:lw-53].mean()
pr_lt7 = mw.intercept_ + mw.coef_ * fw_lt7
print('Wlt7 forecast using model:',pr_lt7.round(2))

fw_lt8 = flow_weekly['flow'][lw-55:lw-52].mean()
pr_lt8 = mw.intercept_ + mw.coef_ * fw_lt8
print('Wlt4 forecast using model:',pr_lt8.round(2))

fw_lt9 = flow_weekly['flow'][lw-55:lw-51].mean()
pr_lt9 = mw.intercept_ + mw.coef_ * fw_lt9
print('Wlt5 forecast using model:',pr_lt9.round(2))

fw_lt10 = flow_weekly['flow'][lw-54:lw-50].mean()
pr_lt10 = mw.intercept_ + mw.coef_ * fw_lt10
print('Wlt1 forecast using model:',pr_lt10.round(2))

fw_lt11 = flow_weekly['flow'][lw-53:lw-49].mean()
pr_lt11 = mw.intercept_ + mw.coef_ * fw_lt11
print('Wlt2 forecast using model:',pr_lt11.round(2))

fw_lt12 = flow_weekly['flow'][lw-52:lw-48].mean()
pr_lt12 = mw.intercept_ + mw.coef_ * fw_lt12
print('Wlt2 forecast using model:',pr_lt12.round(2))

fw_lt13 = flow_weekly['flow'][lw-51:lw-47].mean()
pr_lt13 = mw.intercept_ + mw.coef_ * fw_lt13
print('Wlt3 forecast using model:',pr_lt13.round(2))

fw_lt14 = flow_weekly['flow'][lw-50:lw-46].mean()
pr_lt14 = mw.intercept_ + mw.coef_ * fw_lt14
print('Wlt4 forecast using model:',pr_lt14.round(2))

fw_lt15 = flow_weekly['flow'][lw-49:lw-45].mean()
pr_lt15 = mw.intercept_ + mw.coef_ * fw_lt15
print('Wlt5 forecast using model:',pr_lt15.round(2))

fw_lt16 = flow_weekly['flow'][lw-48:lw-44].mean()
pr_lt16 = mw.intercept_ + mw.coef_ * fw_lt16
print('Wlt16 forecast using model:',pr_lt16.round(2))


# %%
