import numpy as np
from scipy.integrate  import quad

def sing(x):  #(sqrt(x)-sin(sqrt(x))) / sqrt(x)**3
    if(abs(x)<0.08):
        return (1.-x/20.*(1.-x/42.*(1.-x/72.*(1.-x/110.*(1.-x/156.*(1.-x/210.))))))/6.
    elif(x > 0.):
        u = np.sqrt(x)
        return (u-np.sin(u))/u**3
    else:
        u = np.sqrt(-x)
        return (np.sinh(u) - u)/u**3

def sinf(x):
    if(abs(x)<1.e-2):
        return 1.-x/6.*(1.-x/20.*(1.-x/42.*(1.-x/72.*(1.-x/110.*(1.-x/156.)))))
    elif(x > 0.):
        u = np.sqrt(x)
        return np.sin(u)/u
    else:
        u = np.sqrt(-x)
        return np.sinh(u)/u

def tanf(x):  # arctan(sqrt(x))/ sqrt(x)
    if(abs(x)<1.e-2 ):
        return 1.-x*(1./3.-x*(1./5.-x*(1./7.-x*(1./9.-x*(1./11.-x*(1./13.-x*(1./15.-x/17.)))))))
    elif(x > 0.):
        u = np.sqrt(x)
        return np.arctan(u)/u
    else:
        u = np.sqrt(-x)
        if(u >= 1.):
            print("Warning: tanf overflow, returning 1.e99")
            return 1.e99
        else:
            return np.arctanh(u)/u



def time_from_theta(theta, e):
    if(e < 1.): #closed orbit, arbitrary theta
        period = 2.*np.pi/(1.-e**2)**1.5
        n_period = np.floor(theta/(np.pi*2.) + 0.5)        
        base_time = period*n_period
        theta = theta - (np.pi*2.)*n_period
        if(theta/np.pi >= 0.999999999):
            return base_time + period/2. - (np.pi-theta)/(1.+e*np.cos(theta))**2
        elif(theta/np.pi <= -0.999999999):
            return base_time - period/2.+ (theta + np.pi)/(1.+e*np.cos(theta))**2
    else:
        rinv = 1+e*np.cos(theta)        
        if(rinv <= 0.):
            print("Warning: unphysical input of theta.")
            print("For e = ", e)
            print("You must have |theta| < ", np.arccos(-1./e))                        
            exit()
        elif(rinv < 1.e-11):  #in this case the precisioni is actually limited by the round-off error, we cannot be very accurate unless we go beyond double-precsion float
            dtheta = -theta*1.e-11
            k1 = dtheta/rinv**2
            k2 = dtheta/(1.+e*np.cos(theta+k1/2.))**2
            k3 = dtheta/(1.+e*np.cos(theta+k2/2.))**2
            k4 = dtheta/(1.+e*np.cos(theta+k3))**2                
            base_time = base_time -  (k1+k4+2.*(k2+k3))/6.
            theta = theta + dtheta
        else:
            base_time = 0.
    u = np.tan(theta/2.)
    v = (1.-e)/(1.+e)*u**2
    Qtv = tanf(v)
    Erat = 2.*Qtv/(1.+e)*u  #E/(1-e^2)^{3/2}
    Esq = 4.*Qtv**2*v  # E^2
    return base_time + Erat*(sing(Esq)*Erat**2 + sinf(Esq)/(1.+e)) #(E-e sin(E))/(1-e^2)^{3/2}


def theta_from_time(t, e):
    if(e < 1.): #closed orbit, count periods
        period = 2.*np.pi/(1.-e**2)**1.5
        n_period = np.floor(t/period + 0.5)
        base_theta = np.pi*2.*n_period
        t = t - period*n_period
        if(t/period >= 0.4999999999999):
            return base_theta + np.pi
        elif(t/period <= -0.4999999999999):
            return base_theta - np.pi
    else:
        base_theta = 0.
    if(e < 0.99): 
        t  = t*(1.-e**2)**1.5
        F = max(min(t+e*np.sin(t), np.pi*0.9999), -np.pi*0.9999)
        corr = (t- F + e*np.sin(F))/(1.-e*np.cos(F))
        while(abs(corr) > 1.e-12*abs(F)):
            F = F + corr
            corr = (t- F + e*np.sin(F))/(1.-e*np.cos(F))
        F = F + corr
        theta = 2.*np.arctan(np.sqrt((1.+e)/(1.-e))*np.tan(F/2.))
    else:
        Qtv = 1.                
        lam = np.sqrt(2./(1.+e))
        t3 = 3.*t/lam**3
        s = (np.sqrt(1.+t3**2)+t3)**(1./3.)
        Erat = (s - 1./s)*lam
        u = Erat*(1.+e)/2./Qtv
        v = (1.-e)/(1.+e)*u**2
        if(v < -0.5):
            t  = t*(e**2-1)**1.5
            F = np.arcsinh(t/e)
            corr = (t - e*np.sinh(F) + F)/(e*np.cosh(F)-1.)
            while(abs(corr) > 1.e-12*abs(F)):
                F = F + corr
                corr = (t - e*np.sinh(F) + F)/(e*np.cosh(F)-1.)
            F = F + corr
            theta = 2.*np.arctan(np.sqrt((e+1.)/(e-1.))*np.tanh(F/2.))
        else:
            for i in range(10):
                Qtv = tanf(v)
                Esq = 4.*Qtv**2*v  # E^2
                lam = np.sqrt( sinf(Esq)/3./(1.+e)/sing(Esq) )        
                t3 = t/2./lam**3/sing(Esq)
                s = (np.sqrt(1.+t3**2)+t3)**(1./3.)
                Erat = (s - 1./s)*lam
                u = Erat*(1.+e)/Qtv/2.
                v = (1.-e)/(1.+e)*u**2            
            theta = 2.*np.arctan(u)
    return base_theta + theta

e = float(input("Enter eccentricity ="))
t = float(input("Enter time (in unit of sqrt(p^3/GM)) = "))
theta = theta_from_time(t, e)
print("theta = ", theta)
print("mapping back: time = ",  time_from_theta(theta, e))