# -*- coding: utf-8 -*-

import enum
import sys

from . import curvature
from . import diff
from . import interp
from . import utils

import itertools as it
import numpy as np
import scipy.interpolate
import scipy.integrate
import scipy.optimize
from scipy.optimize import root_scalar

def calc_expansion(x, z, metric, curv, surf, direction = 1, diff_op = diff.fd4):
    """
    Calculate expansion of null geodesics on a sequence of surfaces [1]. The
    surfaces are specified as level surfaces of F(r, θ) = r - h(θ).

    [1] Alcubierre (2008): Introduction to 3+1 Numerical Relativity, chapter
        6.7, specifically equation (6.7.13).

    :param array_like x: 1D array of x coordinates.
    :param array_like z: 1D array of z-coordinates.
    :param array_like metric: 4D array of spatial metric values at the grid
                              formed by x and z. metric[i, j, k, l] is the ijth
                              component of the metric at the point (X=x[l],
                              Z=z[k]).
    :param array_like curv: values of the extrinsic curvature, otherwise same as
                            metric.
    :param callable surf: A callable that specifies the surfaces. Accepts an
                          array of θ and returns the array of correponding h.
    :param int direction: Values of 1/-1 specify that the expansion of outgoing
                          or ingoing geodesics is to be computed.
    :rtype: array_like, shape (z.shape[0], x.shape[0])
    :return: Expansion values at the grid formed from x and z.
    """
    dX = [x[1] - x[0], 0, z[1] - z[0]]

    X, Z  = np.meshgrid(x, z)
    R     = np.sqrt(X ** 2 + Z ** 2)
    Theta = np.where(R > 1e-8, np.arccos(Z / R), 0.0)

    metric_u = utils.matrix_invert(metric)

    Gamma = curvature.calc_christoffel(x, z, metric, diff_op)

    trK = np.einsum('ij...,ij...', metric_u, curv)

    F = R[:]
    for i in range(Theta.shape[0]):
        F[i] -= surf.eval(Theta[i])

    dF = np.empty((3,) + F.shape)
    dF[0] = diff_op(F, 1, dX[0])
    dF[1] = 0.0
    dF[2] = diff_op(F, 0, dX[2])

    s_l  = direction * dF[:]
    s_u  = np.einsum('ij...,j...->i...', metric_u, s_l)
    s_u /= np.sqrt(np.einsum('ij...,i...,j...', metric, s_u, s_u))

    ds_u = np.zeros((3,) + s_u.shape)
    for i in range(3):
        for j in range(3):
            if i == 1 or j == 1:
                continue
            diff_dir = 1 if (i == 0) else 0
            ds_u[i, j] = diff_op(s_u[j], diff_dir, dX[i])
    ds_u[1, 1] = np.where(np.abs(X) > 1e-8, s_u[0] / X, ds_u[0, 0])

    Div_s_u = np.einsum('ii...', ds_u) + np.einsum('iki...,k...', Gamma, s_u)

    H = Div_s_u - trK + np.einsum('ij...,i...,j...', curv, s_u, s_u)

    return H

def _tensor_transform_cart_spherical_d2(theta, r, x, z, tens):
    zero = np.zeros_like(theta)
    st   = np.sin(theta)
    ct   = np.cos(theta)
    jac  = np.array([[st,    r * ct,   zero ],
                     [zero,    zero, r * st ],
                     [ct,   -r * st,   zero ]])
    return np.einsum('ai...,bj...,ab...->ij...', jac, jac, tens)

def _tensor_transform_cart_spherical_u2(theta, r, x, z, tens):
    eps = 1e-10 # small number to replace NaNs with just huge numbers
    x = np.where(np.abs(x) > eps, x, eps)

    zero = np.zeros_like(theta)
    st = np.sin(theta)
    ct = np.cos(theta)
    r2 = r * r
    # ∂ spherical / ∂ cart
    jac_inv = np.array([[st,       zero,      ct ],
                        [z / r2,   zero, -x / r2 ],
                        [zero,   1. / x,     zero ]])
    return np.einsum('ai...,bj...,ij...->ab...', jac_inv, jac_inv, tens)

def _dtensor_transform_cart_spherical_d2(theta, r, x, z, tens, dtens):
    zero = np.zeros_like(theta)
    st = np.sin(theta)
    ct = np.cos(theta)
    jac = np.array([[st,    r * ct,   zero ],
                    [zero,    zero, r * st ],
                    [ct,   -r * st,   zero ]])
    djac = np.array([[ # ∂^2 x
                        [ zero, ct,    zero ], # / ∂r
                        [ ct, -r * st, zero ], # / ∂θ
                        [ zero, zero,    -x ], # / ∂φ
                     ],
                     [ # ∂^2 y
                         [ zero, zero,     st ], # / ∂r
                         [ zero, zero, r * ct ], # / ∂θ
                         [ st, r * ct, zero   ], # / ∂φ
                     ],
                     [ # ∂^2 z
                         [ zero,     -st, zero ], # / ∂r
                         [ -st,  -r * ct, zero ], # / ∂θ
                         [ zero,    zero, zero ], # / ∂φ
                     ]])

    return np.einsum('bij...,ck...,bc...->ijk...', djac, jac, tens) +\
           np.einsum('bj...,cik...,bc...->ijk...', jac, djac, tens) +\
           np.einsum('bj...,ck...,ai...,abc...->ijk...', jac, jac, jac, dtens)

def _christoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart):
    eps = 1e-10 # small number to replace NaNs with just huge numbers
    x = np.where(np.abs(x) > eps, x, eps)

    zero = np.zeros_like(theta)
    st = np.sin(theta)
    ct = np.cos(theta)
    r2 = r * r
    # ∂ cart / ∂ spherical
    jac = np.array([[st,    r * ct,   zero ],
                    [zero,    zero, r * st ],
                    [ct,   -r * st,   zero ]])
    # ∂ spherical / ∂ cart
    jac_inv = np.array([[st,       zero,      ct ],
                        [z / r2,   zero, -x / r2 ],
                        [zero,   1. / x,     zero ]])
    djac = np.array([[ # ∂^2 x
                        [ zero, ct,    zero ], # / ∂r
                        [ ct, -r * st, zero ], # / ∂θ
                        [ zero, zero,    -x ], # / ∂φ
                     ],
                     [ # ∂^2 y
                         [ zero, zero,     st ], # / ∂r
                         [ zero, zero, r * ct ], # / ∂θ
                         [ st, r * ct, zero   ], # / ∂φ
                     ],
                     [ # ∂^2 z
                         [ zero,     -st, zero ], # / ∂r
                         [ -st,  -r * ct, zero ], # / ∂θ
                         [ zero,    zero, zero ], # / ∂φ
                     ]])

    return np.einsum('ai...,jb...,kc...,ijk...->abc...', jac_inv, jac, jac, Gamma_cart) +\
           np.einsum('ak...,kbc...->abc...', jac_inv, djac)

def _dchristoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart, dGamma_cart):
    eps = 1e-10 # small number to replace NaNs with just huge numbers
    x = np.where(np.abs(x) > eps, x, eps)

    zero = np.zeros_like(theta)
    st = np.sin(theta)
    ct = np.cos(theta)

    r2 = r * r
    r3 = r2 * r
    r4 = r3 * r

    x2 = x * x
    x4 = x2 * x2

    z2 = z * z
    # ∂ cart / ∂ spherical
    jac = np.array([[st,    r * ct,   zero ],
                    [zero,    zero, r * st ],
                    [ct,   -r * st,   zero ]])
    # ∂ spherical / ∂ cart
    jac_inv = np.array([[st,       zero,      ct ],
                        [z / r2,   zero, -x / r2 ],
                        [zero,   1. / x,     zero ]])
    djac = np.array([[ # ∂^2 x
                        [ zero, ct,    zero ], # / ∂r
                        [ ct, -r * st, zero ], # / ∂θ
                        [ zero, zero,    -x ], # / ∂φ
                     ],
                     [ # ∂^2 y
                         [ zero, zero,     st ], # / ∂r
                         [ zero, zero, r * ct ], # / ∂θ
                         [ st, r * ct, zero   ], # / ∂φ
                     ],
                     [ # ∂^2 z
                         [ zero,     -st, zero ], # / ∂r
                         [ -st,  -r * ct, zero ], # / ∂θ
                         [ zero,    zero, zero ], # / ∂φ
                     ]])
    djac_inv = np.array([[ # ∂^2 r
                             [ z2 / r3,     zero,   -x * z / r3 ], # / ∂x
                             [ zero,      1. / r,          zero ], # / ∂y
                             [ -x * z / r3, zero,       x2 / r3 ], # / ∂z
                         ],
                         [ # ∂^2 θ
                             [     -2 * z * np.abs(x) / r4,                                       zero, np.sign(x) * (x2 - z2) / r4], # / ∂x
                             [                        zero, z * (x4 + x2 * z2) * np.abs(x) / (r4 * x4),                       zero ], # / ∂y
                             [ np.sign(x) * (x2 - z2) / r4,                                       zero,   2 * z * np.abs(x) / r4   ], # / ∂z
                         ],
                         [ # ∂^2 φ
                             [ zero,     -1 / x2, zero ], # / ∂x
                             [ -1 / x2,     zero, zero ], # / ∂y
                             [ zero,        zero, zero ], # / ∂z
                         ]])

    d2jac = np.array([[ # ∂^3 x
                         [ # / ∂r
                             [ zero,   zero,    zero ], # / ∂r
                             [ zero,    -st,    zero ], # / ∂θ
                             [ zero,   zero,     -st ], # / ∂φ
                         ],
                         [ # / ∂θ
                             [ zero,    -st,    zero ], # / ∂r
                             [  -st,    -ct,    zero ], # / ∂θ
                             [ zero,   zero, -r * ct ], # / ∂φ
                         ],
                         [ # / ∂φ
                             [ zero,   zero,     -st ], # / ∂r
                             [ zero,   zero, -r * ct ], # / ∂θ
                             [ -st, -r * ct,    zero ], # / ∂φ
                         ],
                      ],
                      [ # ∂^3 y
                         [ # / ∂r
                             [ zero,   zero,    zero ], # / ∂r
                             [ zero,   zero,      ct ], # / ∂θ
                             [ zero,     ct,    zero ], # / ∂φ
                         ],
                         [ # / ∂θ
                             [ zero,   zero,      ct ], # / ∂r
                             [ zero,   zero, -r * st ], # / ∂θ
                             [  ct, -r * st,    zero ], # / ∂φ
                         ],
                         [ # / ∂φ
                             [ zero,     ct,    zero ], # / ∂r
                             [  ct, -r * st,    zero ], # / ∂θ
                             [ zero,   zero, -r * st ], # / ∂φ
                         ],
                      ],
                      [ # ∂^3 z
                         [ # / ∂r
                             [ zero,   zero,    zero ], # / ∂r
                             [ zero,    -ct,    zero ], # / ∂θ
                             [ zero,   zero,    zero ], # / ∂φ
                         ],
                         [ # / ∂θ
                             [ zero,    -ct,    zero ], # / ∂r
                             [  -ct, r * st,    zero ], # / ∂θ
                             [ zero,   zero,    zero ], # / ∂φ
                         ],
                         [ # / ∂φ
                             [ zero,   zero,    zero ], # / ∂r
                             [ zero,   zero,    zero ], # / ∂θ
                             [ zero,   zero,    zero ], # / ∂φ
                         ],
                      ]])

    for i, j, k, l in it.product(range(3), repeat = 4):
        assert(np.all(d2jac[i, j, k, l] == d2jac[i, l, j, k]))
        assert(np.all(d2jac[i, j, k, l] == d2jac[i, k, l, j]))
        assert(np.all(d2jac[i, j, k, l] == d2jac[i, l, k, j]))

    # this is broken (probably borked index somewhere)
    #return np.einsum('ia...,bj...,ck...,dl...,dabc...->lijk...', jac_inv,  jac, jac,  jac, dGamma_cart) +\
    #       np.einsum('iad...,dl...,bj...,ck...,abc...->lijk...', djac_inv, jac, jac,  jac,  Gamma_cart) +\
    #       np.einsum('ia...,bjl...,ck...,abc...->lijk...',       jac_inv,      djac,  jac,  Gamma_cart) +\
    #       np.einsum('ia...,bj...,ckl...,abc...->lijk...',       jac_inv,       jac, djac,  Gamma_cart) +\
    #       np.einsum('ajkl...,ia...->lijk...', d2jac, jac_inv)                                          +\
    #       np.einsum('ajk...,iab...,bl...->lijk...', djac, djac_inv, jac_inv)

    ret = np.empty((3,) + Gamma_cart.shape)
    for a, b, c, d in it.product(range(3), repeat = 4):
        val = np.zeros_like(theta)
        for i, j, k, l in it.product(range(3), repeat = 4):
            val += jac[l, d] * jac_inv[a, i] * jac[j, b] * jac[k, c] * dGamma_cart[l, i, j, k]

        for i, j, k in it.product(range(3), repeat = 3):
            val += (np.einsum('i...,i...', djac_inv[a, i], jac[:, d]) * jac[j, b] * jac[k, c] +
                    jac_inv[a, i] * djac[j, b, d] * jac[k, c] +
                    jac_inv[a, i] * jac[j, b] * djac[k, c, d]) * Gamma_cart[i, j, k]

        for i, j in it.product(range(3), repeat = 2):
            val += djac[i, b, c] * djac_inv[a, i, j] * jac[j, d]

        for i in range(3):
            val += d2jac[i, b, c, d] * jac_inv[a, i]

        ret[d, a, b, c] = val

    return ret

class _MetricFunc:
    _X      = None
    _Z      = None
    _data   = None

    _interp = None
    
    _spline = None

    def __init__(self, X, Z, data, interp_order):
        self.__array_interface__ = data.__array_interface__

        self._X = X
        self._Z = Z
        self._data = data

        origin = (Z[0, 0], X[0, 0])
        step   = (Z[1, 0] - Z[0, 0], X[0, 1] - X[0, 0])

        self._interp = np.empty(data.shape[:-2], dtype = np.object)
        self._spline = np.empty(data.shape[:-2], dtype = np.object)

        for idx in np.ndindex(self._interp.shape):
            self._interp[idx] = interp.Interp2D_C(origin, step, data[idx], interp_order)
            #self._interp[idx] = lambda x, z, o = origin, s = step, d = data[idx]: interp.interp2d(o, s, d, (x, z), interp_order)

            self._spline[idx] = scipy.interpolate.RectBivariateSpline(Z[:, 0], X[0], data[idx])

    def interp(self, idx, x, z):
        #ret = np.empty_like(x)
        #for pt in np.ndindex(ret.shape):
        #    ret[pt] = self._interp[idx].interp(z[pt], x[pt])
        #return ret
        #ret = list(map(lambda pt: self._interp[idx].interp(z[pt], x[pt]), np.ndindex(x.shape)))
        #return np.array(ret)
        #return np.array([self._interp[idx].interp(z, x)])
        return self._interp[idx].interp(z, x)
        #return self._spline[idx](z, x, grid = False)
        #return self._interp[idx](z, x)

    @property
    def shape(self):
        return self._data.shape
    @property
    def strides(self):
        return self._data.strides
    @property
    def dtype(self):
        return self._data.dtype

    def reshape(self, newshape):
        return self._data.reshape(newshape)

    def __getitem__(self, key):
        return self._data[key]

@enum.unique
class ShootResult(enum.Enum):
    EXPIRED    = enum.auto()
    BOUNDARY   = enum.auto()
    AXIS       = enum.auto()
    AXIS_REPEL = enum.auto()

class ShootSolution:
    result = None
    sol    = None
    def __init__(self, result, sol):
        self.result = result
        self.sol    = sol

    def __call__(self, N):
        t = np.linspace(self.sol.t_min, self.sol.t_max, N)
        return self.sol(t)

class NoBracketFound(Exception):
    pass

class AHCalc(object):
    """
    Object encapsulating an axisymmetric apparent horizon calculation, intended
    to be used with the nonlin_solve_1d solver.
    """

    """
    A list of callables calculating the right-hand side of the axisymmetric AH
    equation and its variational derivatives.
    Intended to be passed to nonlin_solve_1d.
    """
    Fs = None

    _verbose        = None

    _diff_op        = None
    _diff2_op       = None
    _interp_order   = None

    _X              = None
    _Z              = None
    _metric_cart    = None
    _curv_cart      = None
    _metric_u_cart  = None
    _dmetric_cart   = None
    _dmetric_u_cart = None
    _d2metric_cart  = None
    _Gamma_cart     = None
    _dGamma_cart    = None
    _dcurv_cart     = None

    class _RHSFunc:
        _hc = None

        def __init__(self, hc):
            self._hc = hc

        def __call__(self, val = 0.0):
            rhs      = lambda x, y: self._hc._ah_rhs(x, y, val = val)
            rhs_var0 = lambda x, y: self._hc._var_diff(x, y, rhs, 0)
            rhs_var1 = lambda x, y: self._hc._var_diff(x, y, rhs, 1)

            return (rhs, rhs_var0, rhs_var1)

        def __getitem__(self, key):
            return self()[key]

        def __len__(self):
            return len(self())

        def __iter__(self):
            return iter(self())

    def __init__(self, X, Z, metric_cart, curv_cart, *,
                 interp_order = 6, verbose = False):
        self._verbose      = verbose
        self._diff_op      = diff.fd8
        self._diff2_op     = diff.fd28
        self._interp_order = interp_order

        self._X            = X
        self._Z            = Z
        self._metric_cart  = _MetricFunc(X, Z, metric_cart, self._interp_order)
        self._curv_cart    = _MetricFunc(X, Z, curv_cart,   self._interp_order)

        self.Fs = [self._ah_rhs, self._dF_r_fd, self._dF_R_fd]
        self.rhs           = self._RHSFunc(self)

    @property
    def X(self):
        return self._X
    @property
    def Z(self):
        return self._Z
    @property
    def metric_cart(self):
        return self._metric_cart
    @property
    def curv_cart(self):
        return self._curv_cart

    @property
    def metric_u_cart(self):
        if self._metric_u_cart is None:
            res = utils.matrix_invert(self.metric_cart)
            res = _MetricFunc(self.X, self.Z, np.ascontiguousarray(res),
                              self._interp_order)
            self._metric_u_cart = res
        return self._metric_u_cart
    @property
    def dmetric_cart(self):
        if self._dmetric_cart is None:
            dmetric = curvature._calc_dmetric(self.X, self.Z, self.metric_cart, self._diff_op)
            self._dmetric_cart = _MetricFunc(self.X, self.Z, dmetric, self._interp_order)
        return self._dmetric_cart
    @property
    def d2metric_cart(self):
        if self._d2metric_cart is None:
            d2metric = curvature._calc_d2metric(self.X, self.Z, self.metric_cart,
                                                self.dmetric_cart, self._diff_op, self._diff2_op)
            self._d2metric_cart = _MetricFunc(self.X, self.Z, d2metric, self._interp_order)
        return self._d2metric_cart
    @property
    def dmetric_u_cart(self):
        if self._dmetric_u_cart is None:
            dmetric_u = -np.einsum('ij...,km...,ljk...->lim...',
                                   self.metric_u_cart, self.metric_u_cart,
                                   self.dmetric_cart)
            self._dmetric_u_cart = _MetricFunc(self.X, self.Z, dmetric_u, self._interp_order)
        return self._dmetric_u_cart
    @property
    def Gamma_cart(self):
        if self._Gamma_cart is None:
            Gamma = curvature._calc_christoffel(self.metric_cart, self.metric_u_cart,
                                                self.dmetric_cart)
            self._Gamma_cart = _MetricFunc(self.X, self.Z, Gamma, self._interp_order)
        return self._Gamma_cart
    @property
    def dGamma_cart(self):
        if self._dGamma_cart is None:
            dGamma = curvature._calc_dchristoffel(self.metric_cart,  self.metric_u_cart,
                                                  self.dmetric_cart, self.dmetric_u_cart,
                                                  self.d2metric_cart)
            self._dGamma_cart = _MetricFunc(self.X, self.Z, dGamma, self._interp_order)
        return self._dGamma_cart
    @property
    def dcurv_cart(self):
        if self._dcurv_cart is None:
            dcurv = curvature._calc_dmetric(self.X, self.Z, self.curv_cart, self._diff_op)
            self._dcurv_cart = _MetricFunc(self.X, self.Z, dcurv, self._interp_order)
        return self._dcurv_cart

    def _metric_cart_interp(self, x, z):
        metric_cart_dst = np.empty((3, 3) + x.shape)
        for idx in it.product(range(3), repeat = 2):
            metric_cart_dst[idx] = self.metric_cart.interp(idx, x, z)
        return metric_cart_dst

    def _dmetric_cart_interp(self, x, z):
        dmetric_cart_dst = np.empty((3, 3, 3) + x.shape)
        for idx in it.product(range(3), repeat = 3):
            dmetric_cart_dst[idx] = self.dmetric_cart.interp(idx, x, z)
        return dmetric_cart_dst

    def _metric_u_cart_interp(self, x, z):
        metric_u_cart_dst = np.empty((3, 3) + x.shape)
        for idx in it.product(range(3), repeat = 2):
            metric_u_cart_dst[idx] = self.metric_u_cart.interp(idx, x, z)
        return metric_u_cart_dst

    def _dmetric_u_cart_interp(self, x, z):
        dmetric_u_cart_dst = np.empty((3, 3, 3) + x.shape)
        for idx in it.product(range(3), repeat = 3):
            dmetric_u_cart_dst[idx] = self.dmetric_u_cart.interp(idx, x, z)
        return dmetric_u_cart_dst

    def _curv_cart_interp(self, x, z):
        curv_cart_dst = np.empty((3, 3) + x.shape)
        for idx in it.product(range(3), repeat = 2):
            curv_cart_dst[idx] = self.curv_cart.interp(idx, x, z)
        return curv_cart_dst

    def _Gamma_cart_interp(self, x, z):
        Gamma_cart_dst = np.empty((3, 3, 3) + x.shape)
        for idx in it.product(range(3), repeat = 3):
            Gamma_cart_dst[idx] = self.Gamma_cart.interp(idx, x, z)
        return Gamma_cart_dst

    def _spherical_metric(self, theta, r, x, z):
        metric_cart_dst = np.empty((3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 2):
            metric_cart_dst[idx] = self.metric_cart.interp(idx, x, z)
        return _tensor_transform_cart_spherical_d2(theta, r, x, z, metric_cart_dst)

    def _spherical_dmetric(self, theta, r, x, z, metric, Gamma):
        # extract the derivatives of the metric from the Christoffel symbols
        Gamma_l = np.einsum('ij...,ikl...->jkl...', metric, Gamma)
        dmetric = np.empty_like(Gamma_l)
        for i, j, k in it.product(range(3), repeat = 3):
            dmetric[i, j, k] = Gamma_l[j, i, k] + Gamma_l[k, i, j]

        return dmetric

    def _spherical_metric_u(self, theta, r, x, z):
        metric_u_cart_dst = np.empty((3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 2):
            metric_u_cart_dst[idx] = self.metric_u_cart.interp(idx, x, z)
        return _tensor_transform_cart_spherical_u2(theta, r, x, z, metric_u_cart_dst)

    def _spherical_dmetric_u(self, theta, r, x, z, dmetric, metric_u):
        return -np.einsum('ik...,jl...,mkl...->mij...', metric_u, metric_u, dmetric)

    def _spherical_curv(self, theta, r, x, z):
        curv_cart_dst = np.empty((3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 2):
            curv_cart_dst[idx] = self.curv_cart.interp(idx, x, z)
        return _tensor_transform_cart_spherical_d2(theta, r, x, z, curv_cart_dst)

    def _spherical_dcurv(self, theta, r, x, z):
        curv_cart_dst  = np.empty(   (3, 3) + theta.shape)
        dcurv_cart_dst = np.empty((3, 3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 2):
            curv_cart_dst[idx] = self.curv_cart.interp(idx, x, z)
        for idx in it.product(range(3), repeat = 3):
            dcurv_cart_dst[idx] = self.dcurv_cart.interp(idx, x, z)
        return _dtensor_transform_cart_spherical_d2(theta, r, x, z, curv_cart_dst, dcurv_cart_dst)

    def _spherical_Gamma(self, theta, r, x, z):
        Gamma_cart_dst = np.empty((3, 3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 3):
            Gamma_cart_dst[idx] = self.Gamma_cart.interp(idx, x, z)
        return _christoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart_dst)

    def _spherical_dGamma(self, theta, r, x, z):
        Gamma_cart_dst  = np.empty(   (3, 3, 3) + theta.shape)
        dGamma_cart_dst = np.empty((3, 3, 3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 3):
            Gamma_cart_dst[idx] =  self.Gamma_cart.interp(idx, x, z)
        for idx in it.product(range(3), repeat = 4):
            dGamma_cart_dst[idx] = self.dGamma_cart.interp(idx, x, z)

        return _dchristoffel_transform_cart_spherical(theta, r, x, z, Gamma_cart_dst, dGamma_cart_dst)

    def _spherical_trK(self, theta, r, x, z):
        metric_u = np.empty((3, 3) + theta.shape)
        curv     = np.empty((3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 2):
            metric_u[idx] = self.metric_u_cart.interp(idx, x, z)
            curv[idx]     = self.curv_cart.interp(idx, x, z)

        return np.einsum('ij...,ij...', metric_u, curv)

    def _var_diff(self, x, y, func, idx, eps = 1e-8):
        y_var      = list(y)
        y_var[idx] = y[idx] + eps
        f_plus      = func(x, y_var)

        y_var[idx]  = y[idx] - eps
        f_minus      = func(x, y_var)

        return (f_plus - f_minus) / (2 * eps)

    def _ah_rhs(self, theta, H, val = 0.0):
        r, dr = H

        r2 = r * r
        st = np.sin(theta)
        ct = np.cos(theta)

        x = r * np.sin(theta)
        z = r * np.cos(theta)

        axis = np.abs(theta / np.pi - (theta / np.pi).astype(int)) < 1e-12

        metric   = self._spherical_metric   (theta, r, x, z)
        metric_u = self._spherical_metric_u(theta, r, x, z)
        curv     = self._spherical_curv    (theta, r, x, z)
        Gamma    = self._spherical_Gamma   (theta, r, x, z)
        trK      = self._spherical_trK     (theta, r, x, z)

        dmetric            = self._spherical_dmetric(theta, r, x, z, metric, Gamma)
        dmetric_u          = self._spherical_dmetric_u(theta, r, x, z, dmetric, metric_u)
        metric_u_cart_yy   = self.metric_u_cart.interp((1, 1),   x, z)
        dmetric_cart_yy_dz = self.dmetric_cart.interp((2, 1, 1), x, z)
        dmetric_cart_yy_dx = self.dmetric_cart.interp((0, 1, 1), x, z)

        u2 = metric_u[1, 1] * dr * dr - 2 * metric_u[0, 1] * dr + metric_u[0, 0]
        u  = np.sqrt(u2)
        u3 = u * u2

        du_F    = np.zeros((3,) + u.shape)
        du_F[0] = metric_u[0, 0] - metric_u[0, 1] * dr
        du_F[1] = metric_u[0, 1] - metric_u[1, 1] * dr

        # F_term = (∂^i F)(∂^j F)(Γ_{ij}^k ∂_k F + u K_{ij})
        F_term = np.zeros_like(u)
        for i, j in it.product(range(2), repeat = 2):
            F_term += du_F[i] * du_F[j] * (Gamma[0, i, j] - dr * Gamma[1, i, j] + u * curv[i, j])

        # X_term = γ^{ij} Γ_{ij}^k ∂_k F
        # this is the only term that is singular at θ = nπ
        X_term = np.zeros_like(u)
        for i, j in it.product(range(2), repeat = 2):
            X_term += metric_u[i, j] * (Gamma[0, i, j] - dr * Gamma[1, i, j])

        # γ^{φφ} Γ_{φφ}^r ∂_r F
        X_term_phi_reg_r  = metric_u[2, 2] * Gamma[0, 2, 2]
        X_term_phi_sing_r = -0.5 * (
                metric_u[0, 0] / r * (2. + r *
                    metric_u_cart_yy *  dmetric_cart_yy_dz) +
                dmetric_u[1, 0, 1] * 2
                )
        #X_term_phi_sing_r = -dmetric_u[1, 0, 1] - metric_u[0, 1] / metric_cart
        # γ^{φφ} Γ_{φφ}^θ
        # at the axis this term goes to γ^{θθ} * d2r, i.e. the rhs we are computing;
        # when multiplied by prefactors they all cancel out, so we get:
        #  rhs = <regular part> - rhs
        # in other words, at the axis we need to set X_term_phi_theta to 0 and
        # instead divide rhs by 2
        X_term_phi_reg_theta  = -metric_u[2, 2] * Gamma[1, 2, 2] * dr

        X_term += np.where(axis,
                           X_term_phi_sing_r + 0.0,
                           X_term_phi_reg_r + X_term_phi_reg_theta)

        denom = metric_u[0, 0] * metric_u[1, 1] - metric_u[0, 1] * metric_u[0, 1]
        ret   = -(u2 * X_term + u3 * trK - F_term + np.sqrt(2.0) * val * u3) / denom

        #ret  *= np.where(theta == 0.0, 0.5, 1.0)
        ret  *= np.where(axis, 0.5, 1.0)

        return ret

    def _ah_init_ux(self, z):
        gxx = self.metric_cart.interp((0, 0), np.array([0.0]), z)
        return 1. / np.sqrt(gxx)

    def _ah_init_uz(self, x):
        gzz = self.metric_cart.interp((2, 2), x, np.array([0.0]))
        return 1. / np.sqrt(gzz)

    def _ah_rhs_xz(self, theta, H, expansion_val = 0.0, clockwise = True,
                   normalize_u = True):
        x, z, dx, dz = H

        bound_x = self._X[0, -(self._interp_order // 2 + 1)]
        bound_z = self._Z[-(self._interp_order // 2 + 1), 0]
        if ((not np.all(np.isfinite(x))) or (not np.all(np.isfinite(x))) or
            np.any(x > bound_x) or np.any(x < -bound_x) or
            np.any(z > bound_z) or np.any(z < -bound_z)):
            return [dx, dz, np.array([0.0]), np.array([0.0])]

        # interpolate the metric components
        metric    = self._metric_cart_interp   (x, z)
        dmetric   = self._dmetric_cart_interp  (x, z)
        metric_u  = self._metric_u_cart_interp (x, z)
        dmetric_u = self._dmetric_u_cart_interp(x, z)
        Gamma     = self._Gamma_cart_interp    (x, z)
        curv      = self._curv_cart_interp     (x, z)

        detg     = utils.matrix_det(metric)

        x2       = x * x
        R2       = x2 * metric[1, 1]

        # tangent vector to the surface
        u = np.zeros((3,) + x.shape)
        u[0] = dx
        u[2] = dz

        if normalize_u:
            u2 = np.einsum('ij...,i...,j...', metric, u, u)
            u /= np.sqrt(u2)
            dx = u[0]
            dz = u[2]

        eta = np.zeros((3,) + x.shape)
        eta[1] = x

        sign = 1.0 if clockwise else -1.0

        s = np.zeros((3,) + x.shape)
        s[0] = -dz
        s[2] = dx
        s *= sign * np.sqrt(detg) * np.sqrt(metric_u[1, 1])

        su = np.einsum('ij...,j...->i...', metric_u, s)

        mu = np.empty_like(metric_u)
        for i, j in it.product(range(3), repeat = 2):
            mu[i, j] = metric_u[i, j] - su[i] * su[j]

        term1_reg  = (su[0] * dmetric[0, 1, 1] + su[2] * dmetric[2, 1, 1]) / (2.0 * metric[1, 1])
        term1_sing = np.where(np.abs(x) > 1e-12,
                              su[0] / x,
                              s[2] * dmetric_u[0, 0, 2])

        term2 = np.einsum('ij...,ij...', mu, curv)

        val = term1_reg + term1_sing - term2 - np.sqrt(2.0) * expansion_val

        d2x = su[0] * val - np.einsum('ij...,i...,j...', Gamma[0], u, u)
        d2z = su[2] * val - np.einsum('ij...,i...,j...', Gamma[2], u, u)

        d2z *= np.where(np.abs(x) > 1e-13, 1.0,
                1.0 / (1. + su[2] * sign * metric_u[0, 0] * np.sqrt(detg * metric_u[1, 1]) / dx))

        return np.array([dx, dz, d2x, d2z])

    def _dF_r_fd(self, theta, H):
        eps = 1e-8
        r, R = H

        rp = r + eps
        rm = r - eps
        Fp = self._ah_rhs(theta, (rp, R))
        Fm = self._ah_rhs(theta, (rm, R))
        return (Fp - Fm) / (2 * eps)

    def _dF_r_exact(self, theta, H):
        r, dr = H

        x = r * np.sin(theta)
        z = r * np.cos(theta)

        metric    = self._spherical_metric   (theta, r, x, z)
        metric_u  = self._spherical_metric_u (theta, r, x, z)
        curv      = self._spherical_curv     (theta, r, x, z)
        dcurv     = self._spherical_dcurv    (theta, r, x, z)
        Gamma     = self._spherical_Gamma    (theta, r, x, z)
        dmetric   = self._spherical_dmetric  (theta, r, x, z, metric, Gamma)
        dmetric_u = self._spherical_dmetric_u(theta, r, x, z, dmetric, metric_u)
        dGamma    = self._spherical_dGamma   (theta, r, x, z)

        u2     = metric_u[1, 1] * dr * dr - 2 * metric_u[0, 1] * dr + metric_u[0, 0]
        u      = np.sqrt(u2)
        var_u2 = dmetric_u[0, 1, 1] * dr * dr - 2 * dmetric_u[0, 0, 1] * dr + dmetric_u[0, 0, 0]
        var_u  = var_u2 / (2 * u)

        df_ij     = np.empty_like(metric_u)
        var_df_ij = np.empty_like(metric_u)
        for i, j in it.product(range(3), repeat = 2):
            df_ij[i, j]     =  ( metric_u   [0, i] -  metric_u   [1, i] * dr) * (metric_u[0, j] - metric_u[1, j] * dr)
            var_df_ij[i, j] = ((dmetric_u[0, 0, i] - dmetric_u[0, 1, i] * dr) * (metric_u[0, j] - metric_u[1, j] * dr) +
                               (dmetric_u[0, 0, j] - dmetric_u[0, 1, j] * dr) * (metric_u[0, i] - metric_u[1, i] * dr))

        term1     = u2 * metric_u                         - df_ij
        var_term1 = var_u2 * metric_u + u2 * dmetric_u[0] - var_df_ij
        term2     = Gamma[0] - dr * Gamma[1] + u * curv
        var_term2 = dGamma[0, 0] - dr * dGamma[0, 1] + dcurv[0] * u + curv * var_u

        denom     =  metric_u   [0, 0] * metric_u[1, 1]                                       -      metric_u   [0, 1] * metric_u[0, 1]
        var_denom = dmetric_u[0, 0, 0] * metric_u[1, 1] + metric_u[0, 0] * dmetric_u[0, 1, 1] - 2 * dmetric_u[0, 0, 1] * metric_u[0, 1]

        ret = -(np.einsum('ij...,ij...', var_term1, term2) + np.einsum('ij...,ij...', term1, var_term2) - np.einsum('ij...,ij...', term1, term2) * var_denom / denom) / denom
        return ret

    def _dF_R_fd(self, theta, H):
        eps = 1e-8
        r, R = H

        Rp = R + eps
        Rm = R - eps
        Fp = self._ah_rhs(theta, (r, Rp))
        Fm = self._ah_rhs(theta, (r, Rm))
        return (Fp - Fm) / (2 * eps)

    def _dF_R_exact(self, theta, H):
        r, dr = H

        x = r * np.sin(theta)
        z = r * np.cos(theta)

        metric_u = self._spherical_metric_u(theta, r, x, z)
        curv     = self._spherical_curv    (theta, r, x, z)
        Gamma    = self._spherical_Gamma   (theta, r, x, z)

        u2     = metric_u[1, 1] * dr * dr - 2 * metric_u[0, 1] * dr + metric_u[0, 0]
        u      = np.sqrt(u2)
        var_u2 = metric_u[1, 1] * dr * 2  - 2 * metric_u[0, 1]
        var_u  = var_u2 / (2 * u)

        df_ij     = np.empty_like(metric_u)
        var_df_ij = np.empty_like(metric_u)

        for i, j in it.product(range(3), repeat = 2):
            df_ij[i, j]     = (metric_u[0, i] - metric_u[1, i] * dr) * (metric_u[0, j] - metric_u[1, j] * dr)
            var_df_ij[i, j] = -metric_u[1, i] * (metric_u[0, j] - metric_u[1, j] * dr) - metric_u[1, j] * (metric_u[0, i] - metric_u[1, i] * dr)

        term1     = u2 * metric_u - df_ij
        var_term1 = var_u2 * metric_u - var_df_ij
        term2     = Gamma[0] - dr * Gamma[1] + u * curv
        var_term2 = -Gamma[1] + curv * var_u

        ret = -(np.einsum('ij...,ij...', var_term1, term2) + np.einsum('ij...,ij...', term1, var_term2)) / (metric_u[0, 0] * metric_u[1, 1] - metric_u[0, 1] * metric_u[0, 1])
        return ret

    def hor_area(self, hor, log2points = 10):
        theta = np.linspace(0, np.pi, (1 << log2points) + 1)
        r     = hor.eval(theta)
        dr    = hor.eval(theta, 1)

        st    = np.sin(theta)
        ct    = np.cos(theta)

        x     = r * st
        z     = r * ct

        r2    = r * r
        z2    = z * z

        metric_u = np.empty((3, 3) + theta.shape)
        for idx in it.product(range(3), repeat = 2):
            metric_u[idx] = self.metric_u_cart.interp(idx, x, z)

        metric_u_det = utils.matrix_det(metric_u)

        ds = np.empty((3,) + theta.shape)
        ds[0] = x / r - dr * x * z / (r2 * np.sqrt(r2 - z2))
        ds[1] = 0.0
        ds[2] = z / r + dr * np.sqrt(r2 - z2) / r2
        lm2 = np.einsum('ij...,i...,j...', metric_u, ds, ds)

        dA = np.sqrt(lm2 / metric_u_det) * r2 * st
        dA[0]  = 0.0
        dA[-1] = 0.0

        return scipy.integrate.romb(dA, theta[1] - theta[0]) * 2. * np.pi

    def hor_mass(self, hor, *args, **kwargs):
        A = self.hor_area(hor, *args, **kwargs)
        return np.sqrt(A / (16. * np.pi))

    def calc_expansion(self, surf, direction):
        """
        Calculate expansion of null geodesics on a sequence of surfaces [1]. The
        surfaces are specified as level surfaces of F(r, θ) = r - h(θ).

        [1] Alcubierre (2008): Introduction to 3+1 Numerical Relativity, chapter
            6.7, specifically equation (6.7.13).

        :param callable surf: A callable that specifies the surfaces. Accepts an
                              array of θ and returns the array of correponding h.
        :param int direction: Values of 1/-1 specify that the expansion of outgoing
                              or ingoing geodesics is to be computed.
        :rtype: array_like, shape (self.Z.shape[0], self.X.shape[0])
        :return: Expansion values at the grid formed from X and Z.
        """
        X, Z = self.X, self.Z
        dX = [X[0, 1] - X[0, 0], 0, Z[1, 0] - Z[0, 0]]

        R     = np.sqrt(X ** 2 + Z ** 2)
        Theta = np.where(R > 1e-12, np.arccos(Z / R), 0.0)

        trK = np.einsum('ij...,ij...', self.metric_u_cart, self.curv_cart)

        F = R[:]
        for i in range(Theta.shape[0]):
            F[i] -= surf.eval(Theta[i])

        dF = np.empty((3,) + F.shape)
        dF[0] = self._diff_op(F, 1, dX[0])
        dF[1] = 0.0
        dF[2] = self._diff_op(F, 0, dX[2])

        s_l  = direction * dF[:]
        s_u  = np.einsum('ij...,j...->i...', self.metric_u_cart, s_l)
        s_u /= np.sqrt(np.einsum('ij...,i...,j...', self.metric_cart, s_u, s_u))

        ds_u = np.zeros((3,) + s_u.shape)
        for i in range(3):
            for j in range(3):
                if i == 1 or j == 1:
                    continue
                diff_dir = 1 if (i == 0) else 0
                ds_u[i, j] = self._diff_op(s_u[j], diff_dir, dX[i])
        ds_u[1, 1] = np.where(np.abs(X) > 1e-8, s_u[0] / X, ds_u[0, 0])

        Div_s_u = np.einsum('ii...', ds_u) + np.einsum('iki...,k...', self.Gamma_cart, s_u)

        H = Div_s_u - trK + np.einsum('ij...,i...,j...', self.curv_cart, s_u, s_u)

        return H

    def shoot_expansion_surf(self, y0, *,
                             expansion_val = 0.0, clockwise = True,
                             atol = 1e-8, rtol = 1e-8,
                             bound_fact = 0.85, axis_tol = 1e-8,
                             lambda_bound = 1024.0,
                             terminate_axis_repel = False, axis_approach_fact = 5e-2):

        def rhs(t, y):
            ret = self._ah_rhs_xz(t, y, expansion_val = expansion_val,
                                   clockwise = clockwise)
            #print('%s -> %s' % (str(y), str(ret)))
            return ret

        x, z, dx, dz = (np.array([val]) for val in y0)

        metric = self._metric_cart_interp(x, z)
        u = np.zeros((3,) + x.shape)
        u[0] = dx
        u[2] = dz
        u2 = np.einsum('ij...,i...,j...', metric, u, u)
        u /= np.sqrt(u2)
        dx = u[0]
        dz = u[2]

        y0 = (x[0], z[0], dx[0], dz[0])

        solver = scipy.integrate.LSODA(rhs, y0 = y0, t0 = 0.0, t_bound = lambda_bound,
                rtol = rtol, atol = atol, vectorized = True)

        ts           = [solver.t]
        interpolants = []

        status = None
        y_old  = solver.y
        max_x  = 0.0
        while solver.t < solver.t_bound:
            msg = solver.step()
            if solver.status == 'failed':
                raise Exception('Integration failed at step %g: %s' % (solver.t, msg))
            elif solver.status == 'finished':
                status = ShootResult.EXPIRED
                break

            ts.append(solver.t)
            interpolants.append(solver.dense_output())

            if (abs(solver.y[0]) > self._X[0, -1] * bound_fact or
                abs(solver.y[1]) > self._Z[-1, 0] * bound_fact):
                status = ShootResult.BOUNDARY
                break
            if solver.y[0] < axis_tol and solver.y[2] < 0:
                status = ShootResult.AXIS
                break
            if (terminate_axis_repel and y_old[2] < 0 and solver.y[2] >= 0
                and (solver.y[3] > 0 or solver.y[0] < axis_approach_fact * max_x)):
                status = ShootResult.AXIS_REPEL
                break
            y_old = solver.y
            max_x = max(max_x, y_old[0])

        return ShootSolution(status, scipy.integrate.OdeSolution(ts, interpolants))

    def calc_expansion_surf(self, z, initial_guess = 0.0, tol = 1e-3):
        """
        For a given point z on the z-axis, compute the curve of constant
        expansion that passes through it, and the value of that expansion.
        """
        y0 = (0.0, z, 1.0, 0.0)

        ######################################################################
        # How this works:
        # - the z axis is repulsive (equation coefficients diverge to ∞ as x→0)
        #   for curves that are not perpendicular to the axis (z'≠0), i.e. only
        #   curves with z'=0 can hit the axis (in numerical approximation, z'
        #   has to be small enough)
        # - for a large enough value of expansion, the curve will loop back on
        #   itself, i.e. it will go from the z-axis to some maximum x value, then
        #   turn back towards the axis and then loop "upwards" (i.e. x'=0,z'>0)
        #   in the limit of expansion→∞ it collapses to a point
        # - decreasing the expansion, the curve stops looping upwards and loops
        #   downwards instead (i.e. x'=0,z'<0)
        # - between those two is the curve that hits the axis with exactly z'=0,
        #   which is our solution
        # - decreasing the expansion even further, the curve may hit the domain
        #   boundary without any loops
        #
        # So our algorithm looks for the zero of a function that is proportional
        # to some power (we pick 2, some other power might be better?) of the
        # final distance from the axis and is
        # - positive if the curve looped back on itself
        # - negative otherwise
        ######################################################################

        def expansion_weight(expansion_val):
            sys.stderr.write('shooting curve expansion=%g\n' % expansion_val)
            sol = self.shoot_expansion_surf(y0,
                    expansion_val = expansion_val, terminate_axis_repel = True)

            if sol.result in (ShootResult.EXPIRED, ShootResult.BOUNDARY):
                x = sol.sol(sol.sol.t_max)[0]
                sys.stderr.write('hit boundary; final x=%g\n' % x)
                return -(x ** 2)
            if sol.result == ShootResult.AXIS:
                sys.stderr.write('hit z-axis\n')
                return 0.0
            if sol.result == ShootResult.AXIS_REPEL:
                y   = sol.sol(sol.sol.t_max)
                ret = -(y[0] ** 2) * np.sign((y[1] - z) * y[3])
                sys.stderr.write('axis repel %s weight=%g\n' % (str(y), ret))
                return ret
            sys.stderr.write('Reached integration interval end\n')
            return np.nan

        # try the initial guess
        sys.stderr.write('looking for bracket\n')
        val0 = initial_guess
        w0   = expansion_weight(val0)

        # look for a bracket,
        # unless we just happened to hit zero on the first try,
        if w0 != 0.0:
            direction = -np.sign(w0)

            for offset in np.linspace(0.1, 3.0, 32):
                val1 = initial_guess + direction * (offset ** 3)
                w1   = expansion_weight(val1)
                if w0 * w1 <= 0.0:
                    break
                else:
                    # sign didn't change, update old side of the bracket
                    w0   = w1
                    val0 = val1

            if w1 * w0 > 0.0:
                raise ValueError('Could not bracket the target curve')

            bracket = [val0, val1]
            sys.stderr.write('got bracket: %s\n' % str(bracket))

            root = scipy.optimize.root_scalar(expansion_weight,
                    method = 'brentq', bracket = bracket, xtol = tol).root
        else:
            root = val0

        sys.stderr.write('found surface for z=%g: expansion=%g\n' % (z, root))

        sol = self.shoot_expansion_surf(y0, expansion_val = root, terminate_axis_repel = True)
        return root, sol

    def find_expansion_surf(self, expansion_val = 0.0, rtol = 1e-2,
                            bracket = None, surfaces = None,
                            surf_range = None, nb_surfaces = 128):
        def horizon_weight(z0):
            if self._verbose:
                sys.stderr.write('trying surface through z=%g\n' % z0)

            y0 = (0.0, z0, 1.0, 0.0)

            sol = self.shoot_expansion_surf(y0, expansion_val = expansion_val,
                                            terminate_axis_repel = True)
            if sol.result in (ShootResult.EXPIRED, ShootResult.BOUNDARY):
                x = sol.sol(sol.sol.t_max)[0]
                if self._verbose:
                    sys.stderr.write('hit boundary; final x=%g\n' % x)
                return -(x ** 2)
            if sol.result == ShootResult.AXIS:
                if self._verbose:
                    sys.stderr.write('hit z-axis\n')
                return 0.0
            if sol.result == ShootResult.AXIS_REPEL:
                y   = sol.sol(sol.sol.t_max)
                ret = -(y[0] ** 2) * np.sign((y[1] - z0) * y[3])
                if self._verbose:
                    sys.stderr.write('axis repel %s weight=%g\n' % (str(y), ret))
                return ret

            sys.stderr.write('Reached integration interval end\n')
            return np.nan

        # find the bracket if one is not supplied
        if bracket is None:
            if surfaces is None:
                if surf_range is None:
                    surf_range = 0.65 * self._Z[-1, 0]
                surfaces = np.linspace(surf_range, 0.0, nb_surfaces)

            bracket = [None, surfaces[0]]
            for z0 in surfaces:
                if self._verbose:
                    sys.stderr.write('Trying surface through z=%g\n' % z0)

                surf = self.shoot_expansion_surf(y0 = (0.0, z0, 1.0, 0.0),
                                                 expansion_val = expansion_val,
                                                 terminate_axis_repel = True)
                if (surf.result != ShootResult.AXIS_REPEL or
                    surf.sol(surf.sol.t_max)[3] < 0):
                    if self._verbose:
                        sys.stderr.write('No upwards turns at (z=%g)\n' % z0)
                    bracket[1] = z0
                    continue

                #ret = ahc.calc_expansion_surf(z, -0.1, tol = 1e-1)
                #if ret[0] > 0.0:
                #    sys.stderr.write('Upwards turn of the Θ=0 curve at (t=%g,z=%g), but expansion is positive\n' % (t, z))
                #    continue

                #if args.verbose:
                #    sys.stderr.write('Found negative expansion %g at (t=%g,z=%g)\n' % (ret[0], t, z))

                bracket[0] = z0
                break

            if bracket[0] is None:
                raise NoBracketFound()

        if self._verbose:
            sys.stderr.write('Looking for root in (%g, %g)\n' % (bracket[1], bracket[0]))

        hor_z = root_scalar(horizon_weight, method = 'brentq',
                            bracket = bracket, rtol = rtol).root

        if self._verbose:
            sys.stderr.write('Found root: %g\n' % hor_z)

        return self.shoot_expansion_surf(
                y0 = (0.0, hor_z, 1.0, 0.0),
                expansion_val = expansion_val,
                terminate_axis_repel = True)
