import os
import Gnuplot
from math import sqrt, pi
from shallow_water_1d import *
from Numeric import allclose, array, zeros, ones, Float, take, sqrt
from config import g, epsilon


#h_0 = 0.48
#u_0 = 0.625
#h_1 = 0.48
#u_1 = 0.625

#h_0 = 0.48
#u_0 = 0.625
#h_1 = 0.106
#u_1 = 2.829228

h_0 = 0.5
u_0 = 0.6
h_1 = 0.5
u_1 = 0.6
g = 9.81

#h_c = (q**2/g)**(1/3)
#u_c = sqrt(g*h_c)

q = 0.3
g = 9.81

b = 4.0
xmax = 25.0
z_bmax = 0.2
Fr_0 = u_0/sqrt(g*h_0)

def stage(x):
    y,u,h = analytical_sol(x)
    return y
    

def elevation(x):
    b = 4.0
    xmax = 25.0
    z_bmax = 0.2
    z_b = zeros(len(x),Float)
    for i in range(len(x)):
        if (x[i] >= xmax-b) & (x[i] <= xmax+b):
            z_b[i] = z_bmax*(1.0-(x[i]-xmax)**2/b**2)
        else:
            z_b[i] = 0.0 
    return z_b
    #return 0.0

def xmomentum(x):
    return u_0*h_0


def bisection(func,xR,xL,H):
    while ((xR - xL) > epsilon):
  
        #Calculate midpoint of domain
        xM = xL + (xR - xL) / 2.0
  
        #Find f(xM)
        if ((f(xL,H) * f(xM,H)) > 0):
            #Throw away left half
            xL = xM
        else:
            #Throw away right half
            xR = xM
    return xR

def f(D,H):
    return D**3+D**2*(H-1.0-Fr_0**2/2.0)+Fr_0**2/2.0

def fprime(D,H):
    Fr_0 = u_0/sqrt(g*h_0)
    return 3*D**2+2*D*(H-1-Fr_0**2/2.0)

def analytical_sol(x):
    y = zeros(len(x),Float)
    u = zeros(len(x),Float)
    height = zeros(len(x),Float)
    for i in range(len(x)):
        if (x[i] >= xmax-b) & (x[i] <= xmax+b):
            zb = z_bmax*(1.0-(x[i]-xmax)**2/b**2)
        else:
            zb = 0.0 

        H = zb/h_0
        y1 = bisection(f,1.0,0.2/0.5,H)
        height[i] = y1*h_0
        y[i] = height[i]+zb
        u[i] = q/(height[i])
        

    return y,u,height

Fr_0 = u_0/sqrt(g*h_0)
zb=0.2
H = zb/h_0
D = bisection(f,1,0.2/0.5,H)

L = 50.0
N = 50 #800
cell_len = L/N
points = zeros(N+1,Float)
for j in range(N+1):
    points[j] = j*cell_len
boundary = { (0,0): 'left',(N-1,1): 'right'}
domain = Domain(points,boundary)
D1 = Dirichlet_boundary([h_0,u_0*h_0])
D2 = Dirichlet_boundary([h_1,u_1*h_1])
domain.set_boundary({'left':D1,'right':D2})
domain.set_quantity('elevation',elevation)
domain.set_quantity('stage',stage)
domain.set_quantity('xmomentum',xmomentum)
X = domain.vertices
C = domain.centroids

Stage = domain.quantities['stage']
Xmom = domain.quantities['xmomentum']

StageQ = Stage.vertex_values
StageC = Stage.centroid_values
XmomQ = Xmom.vertex_values
XmomC = Xmom.centroid_values
ElevationQ = domain.quantities['elevation'].vertex_values

stage_bdry = Stage.boundary_values
xmom_bdry =  Xmom.boundary_values

w,u,h = analytical_sol(X.flat)
wc, uc, hc = analytical_sol(C)

from pylab import plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot,ylim,xlim, rc
hold(False)
rc('text', usetex=True)
h_error = zeros(1,Float)
uh_error = zeros(1,Float)
k = 0
yieldstep = 25.0
finaltime = 25.0
domain.limiter = "vanleer"
#domain.limiter = "steve_minmod"
domain.default_order = 2
domain.default_time_order = 2#check order is not working
for t in domain.evolve(yieldstep = yieldstep, finaltime = finaltime):
    domain.write_time()
print 'Test 1'    
plot1 = subplot(211)
print 'Test 2'
plot(X,w,X,StageQ,X,ElevationQ)
print 'Test 3'
#xlabel('Position')
#ylabel('Stage (m)')
#legend(('Analytical Solution', 'Numerical Solution', 'Channel Bed'),
#       'upper right', shadow=True)
plot1.set_ylim([-0.1,1.0])
print 'Test 4'
plot2 = subplot(212)
print 'Test 5'
plot(X,u*h,X,XmomQ)#/(StageQ-ElevationQ))
print 'Test 6'
#xlabel('x (m)')
#ylabel(r'X-momentum ($m^2/s$)')
h_error[k] = 1.0/(N)*sum(abs(wc-StageC))
uh_error[k] = 1.0/(N)*sum(abs(uc*hc-XmomC))
print "h_error %.10f" %(h_error[k])
print "uh_error %.10f"% (uh_error[k])
print "h_max %.10f"%max(wc-StageC)
print "uh max  %.10f"%max(uc*hc-XmomC)
filename = "steady_flow_"
filename += domain.limiter
filename += str(400)
filename += ".ps"
#savefig(filename)
show()