import numpy as np
import scipy.integrate as itg
#initial conditions; length unit = 2GM
r0 = 4. 
r1 = 13.5
t_ini = 0. #this can be arbitrary number

#define the r function
def r_approx(x):  #approximate r(x)
    C = np.e-1.-np.sqrt(2.)-2./(5.+2**(78./25.))
    lam = (25*(5.+2**(78./25.))**2*(2*np.e-np.sqrt(2.)-4.)-3*2.**(128./25.))/25/(2.-(5.+2.**(78./25.))*(np.e-1.-np.sqrt(2.)))**2
    return np.log(1.+ 2.*C*(x+1.)/(2.+ C*np.log(1.+lam*x*(x+1.))) + np.sqrt(2.+2*x)+(2.+2*x)/(5.+8.*(2.+2*x)**(3./25.)))


def r_exact(x):  #r(x)
    """for input x >= -1 find w>0 such that (w-1)e^w = x"""
    ##check error
    if x < -1:
        print("Error: Invalid input for r_exact")
        return None
    if x < -0.999:
        u = np.sqrt(2.*(1.+x))
        return u*(1.+u*(-1./3. + u*(11./72.+u*(-43./540.+u*769./17280.)))) 
    w = r_approx(x)
    lastw = w
    w = (x*np.exp(-w)-w+1.0)/w+w    
    while(abs(w-lastw)>1.e-13):
        lastw = w
        w = (x*np.exp(-w)-w+1.0)/w+w    
    return w

eq1_C = np.sqrt(1. - 1./r0)  #(tau+mu)d(tau-mu)/ds - (tau-mu) d(tau+mu)/ds
qiuku_tau_plus_mu = np.sqrt(abs(r1-1.))*np.exp((r1 + t_ini)/2.)
mymu0 = np.sqrt(r0-1.)*np.exp(r0/2.)*np.cosh(t_ini/2.)
mytau0 = np.sqrt(r0-1.)*np.exp(r0/2.)*np.sinh(t_ini/2.)
s_end = np.pi/2.*np.sqrt(r0**3)

def my_eom(z, s):  #evolution equations for z = [tau+mu, tau-mu]
    r = r_exact(z[0]*z[1])
    rp = np.exp(-r)/r
    a = z[0]
    b = eq1_C/rp
    c = 0.25*z[1]/rp
    dz1ds = (-b-np.sqrt(abs(b**2-4.*a*c)))/2./a
    dz0ds = -0.25/rp/dz1ds
    return [dz0ds, dz1ds]


n=2000
s = np.linspace(0., s_end*0.9999, n)
z = itg.odeint(my_eom, [mymu0+mytau0, mymu0-mytau0], s)
if(z[n-1, 0] < qiuku_tau_plus_mu):
    print("Oh no, I will miss the qiuku!")
else:    
    for i in range(n-1):
        if(z[i, 0] < qiuku_tau_plus_mu and z[i+1, 0] >= qiuku_tau_plus_mu):
            r = (r_exact(z[i, 0]*z[i, 1])*np.log(z[i+1, 0]/qiuku_tau_plus_mu)+r_exact(z[i+1,0]*z[i+1, 1])*np.log(qiuku_tau_plus_mu/z[i,0]))/np.log(z[i+1,0]/z[i, 0]) #interpolate to get r
            print("I get the qiuku at r/(2GM) = ", r)