import os
from scipy.special import jn
from scipy import sin, cos, sqrt, linspace, pi, zeros
from rootsearch import *
from bisect import *
from Numeric import Float
from numpy import zeros,dot
from gaussPivot import *
from config import g

def newtonRaphson2(f,q,tol=1.0e-15): ##1.0e-9 may be too large.
    for i in range(30):                                                                                                                                    
        h = 1.0e-15      ##1.0e-4 may be too large.
        n = len(q)
        jac = zeros((n,n),Float)
        if 1.0+q[0]-x<0.0:
            temp1 = 1.0+q[0]-x
            q[0] = q[0]-temp1
            q[1] = v
            return q
        f0 = f(q)
        for i in range(n):
            temp = q[i]
            q[i] = temp + h
            f1 = f(q)
            q[i] = temp
            jac[:,i] = (f1 - f0)/h
        if sqrt(dot(f0,f0)/len(q)) < tol: return q
        dq = gaussPivot(jac,-f0)
        q = q + dq
        if sqrt(dot(dq,dq)) < tol*max(max(abs(q)),1.0): return q
    print 'Too many iterations'


def j0(x):
    return jn(0.0, x)

def j1(x):
    return jn(1.0, x)

def j2(x):
    return jn(2.0, x)

def j3(x):
    return jn(3.0, x)

def jm1(x):
    return jn(-1.0, x)

def jm2(x):
    return jn(-2.0, x)

def bed(x):
    return x-1.0

def prescribe(x,t):
    q = zeros(2, Float)
    def fun(q): #Here q=(z, u)
        f = zeros(2,Float)
        f[0] = q[0] + 0.5*q[1]**2.0 - A*j0(4.0*pi/T*(1.0+q[0]-x)**0.5)*cos(2.0*pi/T*(t+q[1]))
        f[1] = q[1] + A*j1(4.0*pi/T*(1.0+q[0]-x)**0.5)*sin(2.0*pi/T*(t+q[1]))/(1+q[0]-x)**0.5
        return f
    def newtonRaphson2(f,q,tol=1.0e-15): ##1.0e-9 may be too large.
        for i in range(30):                                                                                                                                    
            h = 1.0e-15      ##1.0e-4 may be too large.
            n = len(q)
            jac = zeros((n,n),Float)
            if 1.0+q[0]-x<0.0:
                temp1 = 1.0+q[0]-x
                q[0] = q[0]-temp1
                q[1] = v
                return q
            f0 = f(q)
            for i in range(n):
                temp = q[i]
                q[i] = temp + h
                f1 = f(q)
                q[i] = temp
                jac[:,i] = (f1 - f0)/h
            if sqrt(dot(f0,f0)/len(q)) < tol: return q
            dq = gaussPivot(jac,-f0)
            q = q + dq
            if sqrt(dot(dq,dq)) < tol*max(max(abs(q)),1.0): return q
        print 'Too many iterations'
    q = newtonRaphson2(fun,q)
    return q[0], q[1]


def root_g(a,b,t):
    dx = 0.01
    def g(u):
        #It was u + 8.0*A*pi/T*sin(2.0*pi/T*(t+u)). See equation (10) in Johns. Use L'Hospital rule.
        #Note that there is misprint in equation (10) in Johns.
        return u + 2.0*A*pi/T*sin(2.0*pi/T*(t+u))
    while 1:
        x1,x2 = rootsearch(g,a,b,dx)
        if x1 != None:
            a = x2
            root = bisect(g,x1,x2,1)
        else:
            break
    return root


def shore(t):
    a = -0.2#-1.0
    b = 0.2#1.0
    #dx = 0.01
    u = root_g(a,b,t)
    xi = -0.5*u*u + A*cos(2.0*pi/T*(t+u))
    position = 1.0 + xi
    return position, u


def w_at_O(t):
    return eps*cos(2.0*pi*t/T)

def u_at_O(t):
    a = -1.01#-1.0
    b = 1.01#1.0
    dx = 0.01
    w = w_at_O(t)
    def fun(u):
        return u + A*j1(4.0*pi/T*(1.0+w)**0.5)*sin(2.0*pi/T*(t+u))/(1+w)**0.5
    while 1:
        x1,x2 = rootsearch(fun,a,b,dx)
        if x1 != None:
            a = x2
            root = bisect(fun,x1,x2,1)
        else:
            break
    return root


##==========================================================================##
#DIMENSIONAL PARAMETERS
L = 5e4                 # Length of channel (m)
h_0 = 5e2               # Height at origin when the water is still
#N = 100#400                 # Number of computational cells
#cell_len = 1.1*L/N      # Origin = 0.0
Tp = 15.0*60.0                           # Period of oscillation
a = 1.0                                 # Amplitude at origin
#X_dimensionless = linspace(0.0, 1.1*L, N)             # Discretized spatial domain
##=========================================================================##
#DIMENSIONLESS PARAMETERS
eps = a/h_0
T = Tp*sqrt(g*h_0)/L
A = eps/j0(4.0*pi/T)
Time = linspace(0.0,T,1000)
#X = X_dimensionless/L
#Z = bed(X)
#N_X = len(X)
N_T = len(Time)


Stage = zeros(N_T, Float)    
Veloc = zeros(N_T, Float)
for i in range(N_T):
    t=Time[i]
    zet, vel = prescribe(0.0,t)
    Stage[i] = zet
    Veloc[i] = vel


Stage_johns = zeros(N_T, Float)    
Veloc_johns = zeros(N_T, Float)
for i in range(N_T):
    t=Time[i]
    Stage_johns[i] = w_at_O(t)
    Veloc_johns[i] = u_at_O(t)


num=len(Stage)
error_w=(1.0/num)*sum(abs(Stage-Stage_johns))*h_0
error_u=(1.0/num)*sum(abs(Veloc-Veloc_johns))*sqrt(g*h_0)
print "error_w=", error_w
print "error_u=", error_u

"""

from pylab import clf,plot,title,xlabel,ylabel,legend,savefig,show,hold,subplot

hold(False)
clf()
plot1 = subplot(111)
plot(Time/T,Stage*h_0,'b-', Time/T,Stage_johns*h_0,'k--')
xlabel('t/T')
ylabel('Stage')
plot1.set_xlim([0.000,0.030])
plot1.set_ylim([0.980,1.005]) #([-9.0e-3,9.0e-3])
legend(('C-G', 'Johns'),
       'lower left', shadow=False)

plot2 = subplot(212)
plot(Time/T,Veloc*sqrt(g*h_0),'b-', Time/T,Veloc_johns*sqrt(g*h_0),'k--')
xlabel('t/T')
ylabel('Velocity')
#plot1.set_xlim([0.0,1.1])
#plot2.set_ylim([-0.05,0.05]) #([-1.0e-12,1.0e-12])
legend(('C-G', 'Johns'),
       'upper right', shadow=False)


filename = "discrepancy-closer"
#filename += str(i)
filename += ".eps"
savefig(filename)
#show()

#plot(Time,Vel_at_O)
#show()
"""