#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import pandas as pd
import numba


# In[26]:


@numba.jit(nopython=True, parallel=False)
def flowmodel(A, Q, Adown, Qup, zb, B, dx, dt, g, manning):
    imax = len(A)
    Anew, Qnew = np.zeros(imax), np.zeros(imax)
    
# continuous equation
    for i in numba.prange(1, imax-1) : Anew[i] = A[i] - dt * ( Q[i] - Q[i-1] ) / dx
        
    Anew[0], Anew[-1] = Anew[1], Adown
    
# moumentum equation
    for i in numba.prange(1,imax-1): 
        ip, ic, im = (i+1, i, i-1) 
        Cr1 = 0.5*( Q[ic]/A[ic] + Q[ip]/A[ip] )*dt/dx
        Cr2 = 0.5*( Q[ic]/A[ic] + Q[im]/A[im] )*dt/dx
        dHdx1 = ( Anew[ip]/B[ip] + zb[ip] - Anew[ic]/B[ic] - zb[ic] ) / dx
        dHdx2 = ( Anew[ic]/B[ic] + zb[ic] - Anew[im]/B[im] - zb[im] ) / dx
        dHdx = (1.0 - Cr1) * dHdx1 + Cr2 * dHdx2
        
        Qnew[ic] = Q[ic] - dt * ( Q[ic]**2/A[ic] - Q[im]**2/A[im] ) / dx  \
                   - dt * g * Anew[ic] * dHdx \
                   - dt * g * A[ic] * manning**2 * Q[ic]**2 / B[ic]**2 / ( A[ic]/B[ic] )**(10.0/3.0)
                
#     Qnew[0], Qnew[-1] = Qup, Qnew[-2]
    Qnew[0] = Qup
    
# check downstream boundary Q 
    i = -1
    ic, im = (i, i-1) 
    dHdx = ( Anew[ic]/B[ic] + zb[ic] - Anew[im]/B[im] - zb[im] ) / dx
    
    Qnew[ic] = Q[ic] - dt * ( Q[ic]**2/A[ic] - Q[im]**2/A[im] ) / dx  \
               - dt * g * Anew[ic] * dHdx \
               - dt * g * A[ic] * manning**2 * Q[ic]**2 / B[ic]**2 / ( A[ic]/B[ic] )**(10.0/3.0)
            
    return Anew, Qnew 


# In[27]:


get_ipython().run_cell_magic('time', '', 'length = 5000.0\ndx = 100.0\nimax = int(length/dx) + 1\ndt = 10.0\ntotalTime = 5.0*3600.0\nhini = 1.0\nmanning  = 0.03\nib = 1.0/1000.0\ng = 9.8\n\n# Initial & Boundary condition\nB = np.full(imax, 1.0, dtype=float)\nA = hini*B\nQ = ib**0.5*(hini)**(5.0/3.0)/manning*B #normal flow\nzb = np.zeros(imax)\nfor i in range(1,imax):\n    zb[i] = zb[i-1] + ib*dx\nzb = zb[::-1]\n\nWLdown = 2.0\nfor i, (Ap, zbp) in enumerate(zip(A, zb)):\n    if( Ap/B[i] + zbp < WLdown) : A[i] = (WLdown - zbp)*B[i]\n\nQup = Q[0]\nAdown = A[-1]\n\nfor _ in range(int(totalTime/dt)):\n    A, Q = flowmodel(A, Q, Adown, Qup, zb, B, dx, dt, g, manning)\n')


# In[28]:


import holoviews as hv
hv.extension('bokeh')


# In[42]:


get_ipython().run_cell_magic('output', "size=100 filename='tmp' fig='html'", '(hv.Curve(A+zb/B, "x", "E.L.")*hv.Curve(zb)).relabel(\'Elevation\') + hv.Curve(Q,"x","Q").redim(Q={\'range\':(1,1.1)}).relabel(\'Discharge\')\n')