# Display graphics inline with the notebook
%matplotlib inline

# Standard Python modules
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import pandas as pd
import os
import datetime

# Module to enhance matplotlib plotting
import seaborn
seaborn.set()

# Modules to display images and data tables
from IPython.display import Image
from IPython.core.display import display

# Data Directory
dir = './data/'

# Styles
from IPython.core.display import HTML
HTML(open("styles/custom.css", "r").read())

h = np.array([335.0, 336.0, 336.5, 337.0, 337.5, 338.0, 339.0, 340.0])
v = np.array([112.67, 798.00, 1176.42, 1577.25, 2002.06, 2450.57, 3416.85, 4458.97])

plt.subplot(2,1,1)
plt.scatter(h,v)
plt.xlim(h.min(),h.max())
plt.ylim(0,plt.ylim()[1])
plt.title('Rainy Lake Stage-Volume Relationship')
plt.xlabel('Water Elevation (m)')
plt.ylabel('Volume (million cu. m)')

pf = sp.polyfit(h,v,2)

VRL= sp.poly1d(pf)
plt.hold(True)
plt.plot(h,VRL(h))
plt.hold(False)

ARL = sp.poly1d(np.array([2.0*pf[0],pf[1]]))
plt.subplot(2,1,2)
plt.plot(h,ARL(h))
plt.title('Rainy Lake Stage-Surface Area Relationship')
plt.xlabel('Elevation (m)')
plt.ylabel('Area (sq. km.)')

plt.tight_layout()

df = pd.DataFrame(zip(h,v,VRL(h),ARL(h)),columns=['h','V','Vhat','Ahat'])
print df.to_string(formatters={'V':'  {:.0f}'.format, 
                                'Vhat':'  {:.0f}'.format,
                                'Ahat':'  {:.0f}'.format})

h = np.array([337.0, 338.0, 338.5, 339.0, 339.5, 340.0,
     340.5, 341.0, 341.5, 342.0, 343.0])
v = np.array([65.33, 259.95, 364.20, 475.58, 592.46, 712.28,
     836.56, 966.17, 1099.79, 1239.68, 1540.75])

plt.subplot(2,1,1)
plt.scatter(h,v)
plt.xlim(h.min(),h.max())
plt.ylim(0,plt.ylim()[1])
plt.title('Namakan Lake Stage-Volume Relationship')
plt.xlabel('Water Elevation (m)')
plt.ylabel('Volume (million cu. m)')

pf = sp.polyfit(h,v,2)

VNL= sp.poly1d(pf)
plt.hold(True)
plt.plot(h,VNL(h))
plt.hold(False)

ANL = sp.poly1d(np.array([2.0*pf[0],pf[1]]))
plt.subplot(2,1,2)
plt.plot(h,ANL(h))
plt.title('Namakan Lake Stage-Surface Area Relationship')
plt.xlabel('Elevation (m)')
plt.ylabel('Area (sq. km.)')

plt.tight_layout()

df = pd.DataFrame(zip(h,v,VNL(h),ANL(h)),columns=['h','V','Vhat','Ahat'])
print df.to_string(formatters={'V':'  {:.0f}'.format, 
                                'Vhat':'  {:.0f}'.format,
                                'Ahat':'  {:.0f}'.format})

h = np.array([335.4, 336.0, 336.5, 336.75, 337.0, 337.25, 337.5, 
              337.75, 338.0, 338.5, 339, 339.5, 340.0])
d = np.array([0.0, 399., 425., 443., 589., 704., 792., 909.,
              1014., 1156., 1324., 1550., 1778.])

#plt.subplot(2,1,1)

plt.hold(True)
plt.scatter(d,h)
plt.plot(d,h)

plt.ylim(h.min(),h.max())
plt.xlim(0,plt.xlim()[1])
plt.title('Rainy Lake Stage-Discharge Relationship')
plt.ylabel('Water Elevation (m)')
plt.xlabel('Discharge (cu. per sec.)')


# Get historical flowrates and levels on RR and RL
RR = pd.read_pickle(dir+'RR.pkl')['1970':]
RL = pd.read_pickle(dir+'RL.pkl')['1970':]

Q = pd.concat([RR,RL],axis=1)
Q.columns = ['RR','RL']

A = Q['1970':'2000']

plt.scatter(A['RR'],A['RL'],marker='+',color='b')
B = Q['2000':]
plt.scatter(B['RR'],B['RL'],marker='+',color='r')

ax = plt.axis()
plt.plot([ax[0],ax[1]],[336.70,336.70],'y--')
plt.plot([ax[0],ax[1]],[337.20,337.20],'y--')
plt.plot([ax[0],ax[1]],[337.75,337.75],'m--')
plt.plot([ax[0],ax[1]],[337.90,337.90],'c--')

plt.xlabel('Upper Rainy River Flow [cubic meters/sec]')
plt.ylabel('Rainy Lake Level [meters]')
ax = plt.axis()
plt.axis([0,ax[1],ax[2],ax[3]])
plt.legend(['Winter Drought Line','Summer Drought Line','URC_max','All Gates Open','1970-2000','2000-2010'],loc="upper left")

plt.hold(False)
#plt.savefig('./images/RainyRiverDischarge.png')


Q = lambda x: np.interp(x,h,d)
x = np.linspace(336.,339.)
plt.plot(x,Q(x))