import numpy
from matplotlib import pyplot
%matplotlib notebook

def horizon_RHS(H, theta, z_singularities):
    """
    The RHS function for the apparent horizon problem.
    
    Parameters
    ----------
    
    H : array
        vector [h, dh/dtheta]
    theta : double
        angle
    z_singularities : array
        Location of the singularities on the z axis; non-negative
    
    Returns
    -------
    
    dHdtheta : array
        RHS
    """
    
    assert(numpy.all(numpy.array(z_singularities) >= 0.0)), "Location of singularities cannot be negative"
    
    h = H[0]
    dh = H[1]
    
    psi = 1.0
    dpsi_dr = 0.0
    dpsi_dtheta = 0.0
    for z in z_singularities:
        distance = numpy.sqrt((h*numpy.sin(theta))**2 + (h*numpy.cos(theta) - z)**2)
        psi += 0.5/distance
        dpsi_dr -= 0.5*(h-z*numpy.cos(theta))/distance**3
        dpsi_dtheta -= 0.5**h*z*numpy.sin(theta)/distance**3
        # Apply reflection symmetry
        if z > 0.0:
            distance = numpy.sqrt((h*numpy.sin(theta))**2 + (h*numpy.cos(theta) + z)**2)
            psi += 0.5/distance
            dpsi_dr -= 0.5*(h+z*numpy.cos(theta))/distance**3
            dpsi_dtheta += 0.5**h*z*numpy.sin(theta)/distance**3
            

    C2 = 1.0 / (1.0 + (dh / h)**2)
    # Impose that the term involving cot(theta) vanishes on axis.
    if (abs(theta) < 1e-16) or (abs(theta - numpy.pi) < 1e-16):
        cot_theta_dh_C2 = 0.0
    else:
        cot_theta_dh_C2 = dh / (numpy.tan(theta) * C2)
        
    dHdtheta = numpy.zeros_like(H)
    dHdtheta[0] = dh
    dHdtheta[1] = 2.0*h - cot_theta_dh_C2 + 4.0*h**2/(psi*C2)*(dpsi_dr - dpsi_dtheta*dh/h**2) + 3.0*dh**2/h
    
    return dHdtheta

"""
Explanation: Apparent horizons
We're now going to use finite differences to find a black hole apparent horizon.
The spacetime we're going to look at is simplified:

$3+1$ split (we're looking at one slice, so one instant in "time");
axisymmetric (so we can consider only two dimensions in space, using $r, \theta$);
"bitant" or "reflection" symmetric (so we only consider $\theta \in [0, \pi/2]$);
all singularities have bare mass $1$;
time-symmetric (the extrinsic curvature vanishes).

We then compute the expansion of outgoing null geodesics, and look for where this vanishes. The surface with radius $h(\theta)$ where this occurs is the apparent horizon. With our assumptions, $h$ obeys the boundary value problem
$$
  \frac{d^2 h}{d \theta^2} = 2 h + \frac{3}{h} \left( \frac{d h}{d \theta} \right)^2 + f \left( \theta, h, \frac{d h}{d \theta} \right), \qquad \frac{d h}{d \theta} ( \theta = 0 ) = 0 = \frac{d h}{d \theta} ( \theta = \pi/2 ).
$$
The function $f$ encodes the spacetime effects due to the singularities. 
To solve this problem we convert to first order form. Introduce the vector 
$$
  {\bf H} = \begin{pmatrix} h \ \frac{d h}{d \theta} \end{pmatrix}.
$$
Then we have the problem
$$
  \frac{d}{d \theta} {\bf H} = {\bf F}({\bf H}, \theta) = \begin{pmatrix} H_2 \ 2 H_1 + \frac{3}{H_1} H_2^2 + f(\theta, {\bf H}) \end{pmatrix}, \qquad H_2(\theta = 0) = 0 = H_2(\theta = \pi/2).
$$
We'll give the entire right-hand-side as code:
End of explanation
"""


def euler_step(Hi, theta_i, dtheta, z_singularity):
    """
    Euler's method - one step
    """
    
    return Hi + dtheta*horizon_RHS(Hi, theta_i, z_singularity)

Ntheta = 100
z_singularity = [0.0]
theta = numpy.linspace(0.0, numpy.pi/2.0, Ntheta)
dtheta = theta[1] - theta[0]
H = numpy.zeros((2, Ntheta))
H[:, 0] = [0.5, 0.0]

for i in range(Ntheta-1):
    H[:, i+1] = euler_step(H[:, i], theta[i], dtheta, z_singularity)

pyplot.figure()
pyplot.polar(theta, H[0,:])
pyplot.show()

"""
Explanation: We now need to solve the boundary value problem. We'll do this using shooting.
Shooting
Initial Value Problems
If we knew the initial radius of the horizon, $h(\theta = 0) = H_1(\theta = 0) = h_0$, we could solve the initial value problem
$$
  \frac{d}{d \theta} {\bf H} = {\bf F}({\bf H}, \theta) = \begin{pmatrix} H_2 \ 2 H_1 + \frac{3}{H_1} H_2^2 + f(\theta, {\bf H}) \end{pmatrix}, \qquad {\bf H}(\theta = 0) = \begin{pmatrix} h_0 \ 0 \end{pmatrix}.
$$
For example, the simple Schwarzschild black hole will have $h_0 = 1/2$, in this slicing.
To solve the initial value problem we can re-use our finite differencing algorithms. For example, we evaluate the initial value problem equation at $\theta_i$ using forward differencing, to get
\begin{align}
  \left. \frac{d}{d \theta} {\bf H} \right|_{\theta = \theta_i} & \approx \frac{1}{\Delta \theta} \left( {\bf H}^{(i+1)} - {\bf H}^{(i)} \right) \
  & = {\bf F}({\bf H}^{(i)}, \theta_i),
\end{align}
where we have denoted ${\bf H}(\theta_i) \equiv {\bf H}^{(i)}$. We then re-arrange this to get Euler's method
$$
  {\bf H}^{(i+1)} = {\bf H}^{(i)} + \Delta \theta \, {\bf F}({\bf H}^{(i)}, \theta_i).
$$
We can use this to solve for the Schwarzschild case:
End of explanation
"""


def rk2_step(Hi, theta_i, dtheta, z_singularity):
    """
    RK2 method - one step
    """
    
    Hp = Hi + dtheta * horizon_RHS(Hi, theta_i, z_singularity)
    
    return 0.5*(Hi + Hp + dtheta*horizon_RHS(Hp, theta_i+dtheta, z_singularity))

H = numpy.zeros((2, Ntheta))
H[:, 0] = [0.5, 0.0]

for i in range(Ntheta-1):
    H[:, i+1] = rk2_step(H[:, i], theta[i], dtheta, z_singularity)

pyplot.figure()
pyplot.polar(theta, H[0,:])
pyplot.show()

"""
Explanation: We see that this has worked nicely. However, Euler's method is very inaccurate on more complex problems, as it's only first order convergent. We would like to use a higher order method.
Runge-Kutta methods
When looking at central differencing earlier we used information from both sides of the point where we took the derivative. This gives higher accuracy, but isn't helpful in the initial value case, where we don't have half the information.
Instead, we use many Euler steps combined. Each one gives an approximation to "future" data, which can be used to approximate the derivative at more locations.
For example, the Euler step above starts from ${\bf H}^{(i)}$ and computes ${\bf F}^{(i)}$ to approximate ${\bf H}^{(i+1)}$. We can use this approximation to give us ${\bf F}^{(i+1)}$.
Now, a more accurate solution would be
$$
  {\bf H}^{(i+1)} = {\bf H}^{(i)} + \int_{\theta_i}^{\theta_{i+1}} \text{d} \theta \, {\bf F}({\bf H}, \theta).
$$
In Euler's method we are effectively representing the value of the integral by the value of the integrand at the start, multiplied by the width $\Delta \theta$. We could now approximate it by the average value of the integrand, $({\bf F}^{(i)} + {\bf F}^{(i+1)})/2$, multiplied by the width $\Delta \theta$. This gives the algorithm
\begin{align}
  {\bf H}^{(p)} &= {\bf H}^{(i)} + \Delta \theta \, {\bf F}({\bf H}^{(i)}, \theta_i), \
  {\bf H}^{(i+1)} &= {\bf H}^{(i)} + \frac{\Delta \theta}{2} \left( {\bf F}({\bf H}^{(i)}, \theta_i) + {\bf F}({\bf H}^{(p)}, \theta_{i+1}) \right) \
  &= \frac{1}{2} \left( {\bf H}^{(i)} + {\bf H}^{(p)} + \Delta \theta \, {\bf F}({\bf H}^{(p)}, \theta_{i+1}) \right).
\end{align}
The final re-arrangement ensures we do not have to store or re-compute ${\bf F}^{(i)}$. This is one of the Runge-Kutta methods. This version is second order accurate, and a big improvement over Euler's method.
End of explanation
"""


initial_guesses = numpy.linspace(0.4, 0.6, 10)
solutions = []
z_singularity = [0.0]

for h0 in initial_guesses:
    H = numpy.zeros((2,Ntheta))
    H[:, 0] = [h0, 0.0]
    for i in range(Ntheta-1):
        H[:, i+1] = rk2_step(H[:, i], theta[i], dtheta, z_singularity)
    solutions.append(H[0,:])
    
pyplot.figure()
for r in solutions:
    pyplot.polar(theta, r)
pyplot.show()

"""
Explanation: Root finding
We still can't find a horizon unless we know the initial radius. However, there is a way around this. Let's see what happens if we compute in the Schwarzschild case, using the wrong initial data:
End of explanation
"""


initial_guesses = numpy.linspace(0.4, 0.6, 100)
dhdtheta_end = numpy.zeros_like(initial_guesses)
z_singularity = [0.0]

for guess, h0 in enumerate(initial_guesses):
    H = numpy.zeros((2,Ntheta))
    H[:, 0] = [h0, 0.0]
    for i in range(Ntheta-1):
        H[:, i+1] = rk2_step(H[:, i], theta[i], dtheta, z_singularity)
    dhdtheta_end[guess] = H[1, -1]
    
pyplot.figure()
pyplot.plot(initial_guesses, dhdtheta_end)
pyplot.xlabel(r"$h_0$")
pyplot.ylabel(r"$dh/d\theta(\pi/2)$")
pyplot.show()

"""
Explanation: We see that the the surfaces that start with the radius too small are curving back in; their derivative is negative. The surfaces with radius too large are diverging; their derivative is positive. We know that the true solution has vanishing derivative.
Let's explicitly plot the derivative at the endpoint.
End of explanation
"""


def residual(h0, z_singularities):
    """
    The residual function for the shooting method.
    """
    
    H = numpy.zeros((2, Ntheta))
    H[:, 0] = [h0, 0.0]
    for i in range(Ntheta-1):
        H[:, i+1] = rk2_step(H[:, i], theta[i], dtheta, z_singularity)
        
    return H[1, -1]

"""
Explanation: We see that the derivative vanishes precisely where the horizon should be, exactly as expected.
This also gives us a way of solving for the apparent horizon. We want to solve the equation
$$
  R(h_0) = 0.
$$
The function $R$ is given by $R(h_0) = H_2(\pi/2 ; h_0)$. In other words, we

compute the solution ${\bf H}$ given the initial guess $h_0$ for the unknown initial radius $H_1(0)$;
from the solution for ${\bf H}$ at $\theta = \pi/2$, set $R(h_0) = H_2$.

We can code this residual function as
End of explanation
"""


def secant(R, x0, x1, args, tolerance = 1e-10):
    """
    Secant method
    """
    
    x = x1
    x_p = x0
    residual = abs(R(x, args))
    while residual > tolerance:
        x_new = x - (R(x, args) * (x - x_p)) / (R(x, args) - R(x_p, args))
        x_p = x
        x = x_new
        residual = abs(R(x, args))
    
    return x

"""
Explanation: Finally, we need to find the root of this equation.
Secant method
Problems where we are given an algebraic, nonlinear function ${\bf R}$ and asked to find ${\bf x}$ such that ${\bf R}({\bf x}) = {\bf 0}$ are nonlinear root-finding problems. Many standard solution methods are based on Newton's algorithm:

Guess the root to be ${\bf x}^{(0)}$, and set $n=0$;
Compute the tangent planes to ${\bf R}$ at ${\bf x}^{(n)}$;
Find where these planes intersect zero, and set this to be ${\bf x}^{(n+1)}$;
If not converged to root, go to 2.

Computing the derivative for the tangent in step 2 is slow; instead we use finite differencing again.
In one dimension, Newton's method is
$$
  x^{(n+1)} = x^{(n)} - \frac{R(x^{(n)})}{R'(x^{(n)})}.
$$
Replacing the derivative with a finite difference gives
$$
  x^{(n+1)} = x^{(n)} - \frac{R(x^{(n)}) \left( x^{(n)} - x^{(n-1)} \right)}{R(x^{(n)}) - R(x^{(n-1)})}.
$$
This is the secant method. It's much easier to implement, but requires two initial guesses.
End of explanation
"""


h0 = secant(residual, 0.4, 0.6, z_singularity)
print("Computed initial radius is {}".format(h0))

"""
Explanation: We apply this to the Schwarzschild case:
End of explanation
"""


z_singularity = [0.75]

"""
Explanation: And from this we can compute the correct horizon.
What happens if we get the guess wildly wrong? In this simple case it will nearly always converge to the "right" answer, but in general a poor initial guess means the algorithm - or most root-finding algorithms! - won't converge.
Put it together
We can now compute the more interesting binary black hole case, where the singularities are at $z = \pm 0.75$. Using the symmetry, we need:
End of explanation
"""


initial_guesses = numpy.linspace(1.2, 1.4, 100)
dhdtheta_end = numpy.zeros_like(initial_guesses)
z_singularity = [0.75]
interval = [0.0, numpy.pi/2.0]

for guess, h0 in enumerate(initial_guesses):
    H = numpy.zeros((2,Ntheta))
    H[:, 0] = [h0, 0.0]
    for i in range(Ntheta-1):
        H[:, i+1] = rk2_step(H[:, i], theta[i], dtheta, z_singularity)
    dhdtheta_end[guess] = H[1, -1]
    
pyplot.figure()
pyplot.plot(initial_guesses, dhdtheta_end)
pyplot.xlabel(r"$h_0$")
pyplot.ylabel(r"$dh/d\theta(\pi/2)$")
pyplot.show()

"""
Explanation: We can now check what sorts of initial radius $h_0$ will be needed for the horizon:
End of explanation
"""


h0 = secant(residual, 1.26, 1.3, z_singularity)

"""
Explanation: We see the algorithms are having problems for small radii, but that it suggests that the correct answer is roughly $h_0 \in [1.26, 1.3]$. So we use root-finding:
End of explanation
"""


z_singularity = [0.75]
H0 = [h0, 0.0]
H = numpy.zeros((2,Ntheta))
H[:, 0] = [h0, 0.0]

for i in range(Ntheta-1):
    H[:, i+1] = rk2_step(H[:, i], theta[i], dtheta, z_singularity)

pyplot.figure()
pyplot.polar(theta, H[0,:])
pyplot.show()

"""
Explanation: And finally, we compute and plot the horizon surface.
End of explanation
"""


