import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline













x = np.array([1,3,4,6,8,9])
y = np.array([1,4,2,6,5,7])
p = np.polyfit(x,y,5)
print p
x2 = np.linspace(0,10,100)
y2 = np.polyval(p,x2)
plt.plot(x,y,'bo')
plt.plot(x2,y2)
plt.ylim(-1,8)

x = np.array([1,3,4,6,8,9])
y = np.array([1,4,2,6,5,7])
p = np.polyfit(x,y,3)
print 'the values of p are:',p
x2 = np.linspace(0,10,100)
y2 = np.polyval(p,x2)
plt.plot(x,y,'bo')
plt.plot(x2,y2)











# Step 1
def Rh(h,w):
    A = (w * h) + (h * h)        # m^2, is the cross sectional area of the flow
    P = w + 2 * h * np.sqrt(2.)     # m, is the wetted perimeter of the channel
    rv = A / P
    return rv
Rh(2,4)

# Step 2
def Vavg(h,w,n,S):
    rv = Rh(h,w)**(2.0/3) * np.sqrt(S) / n
    return rv
Vavg(2,4,0.025,0.005)

# Step 3
def Qtot(h,w,n,S):
    A = (w+h) * h
    return A * Vavg(h,w,n,S)
Qtot(2,4,0.025,0.005)

# Step 4
h = np.linspace(1,4,100)
Q = Qtot(h,4,0.025,0.005)
plt.plot(h,Q,label='w=4')
Q = Qtot(h,6,0.025,0.005)
plt.plot(h,Q,label='w=6')
Q = Qtot(h,8,0.025,0.005)
plt.plot(h,Q,label='w=8')
plt.legend(loc='best')
plt.xlabel('h (m)')
plt.ylabel('Q (m$^3$/s)')

# Step 5
def Qdiff(h,w,n,S,Qd):
    Qt = Qtot(h,w,n,S)
    return Qd - Qt
from scipy.optimize import fsolve
h50 = fsolve(Qdiff,3,args=(4,0.025,0.005,50))
print 'design discharge is 50'
print 'computed water height is ',h50
print 'discharge computed with Qtot is ',Qtot(h50,4,0.025,0.005)

# Step 6
w = 6.0
n = 0.02
S = 0.004
h = np.zeros(5)
Qd = np.arange(20,101,20)
for i in range(5):
    h[i] = fsolve(Qdiff,1,args=(w,n,S,Qd[i]))
print 'Design discharge is ',Qd
print 'Corresponding water height is ',h
plt.plot(Qd,h,'b-o')
plt.xlabel('Design discharge (m$^3$/d)')
plt.ylabel('Water height (m)')
plt.title('w=6 m, n=0.02, S=0.004')

# Step 1
head = np.loadtxt('emptying_reservoir.csv',skiprows=54,usecols=[1],delimiter=";")

# Step 2
plt.subplot(211)
plt.plot(head)
plt.xlim(170,190)
plt.ylim(1115,1120)
plt.subplot(212)
plt.plot(head)
plt.xlim(325,335)
plt.ylim(1040,1080)
print 'Data from index 184 to 332'

# Step 3
headnew = head[184:332]
print len(headnew)
time = np.arange(0,len(headnew)*30,30)
plt.plot(time,headnew)
plt.xlabel('time (seconds')
plt.ylabel('head (cm)')

# Step 4
a = np.polyfit(time,headnew,2)
hfit = np.polyval(a,time)
plt.plot(time,headnew,'b',label='data')
plt.plot(time,hfit,'r',label='parabola')
plt.xlabel('time (seconds')
plt.ylabel('head (cm)')
plt.legend(loc='best')

# Step 5
def dhdt(t):
    return 2*a[0]*t + a[1]
dhdtnum = (headnew[1:] - headnew[:-1]) / 30.0
plt.plot(0.5*(time[1:]+time[:-1]), dhdtnum,label='numerical 1') 
dhdtnum2 = (headnew[2:] - headnew[:-2]) / 60.0
plt.plot(time[1:-1],dhdtnum2,'r',label='numerical 2')
dhdtpoly = dhdt(time)
plt.plot(time,dhdtpoly,'k',label='parabola')
plt.legend(loc='best')
plt.xlabel('time (seconds')
plt.ylabel('head (cm)')