from scipy import sin, cos, sqrt, pi, dot
from gaussPivot import *
from parameter import *

def root_g(a,b,t):
    dx = 0.01
    def g(u):
        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 = -1.0
    b = 1.0
    #dx = 0.01
    u = root_g(a,b,t)
    xi = -0.5*u*u + A*cos(2.0*pi/T*(t+u)) #This is the displacement
    #position = 1.0 + xi
    return [xi, u]#[position, u]

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

def u_at_O(t):
    a = -1.01
    b = 1.01
    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.0+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


def prescribe(x,t):
    from Numeric import array, zeros, Float
    q_shore = shore(t)
    q = zeros(2,Float)
    if x<1.0-abs(q_shore[0]):
        q[0] = w_at_O(t)
        q[1] = u_at_O(t)
    else:
        q = q_shore
    def fun(q): #Here q=(w, 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):
        for i in range(30):                                                                                                                                    
            h = 1.0e-4      ##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] = x-1.0+0.01
                #print "1.0+q[0]-x=",1.0+q[0]-x
                print "SomethingWRONG"
                #return q
                #q[0] =0.25
            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]

#w,u=prescribe(0.0,0.0)