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

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

import itertools as it
import numpy as np
import scipy.integrate

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):
    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):
    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):
    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 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

    _diff_op        = None
    _diff2_op       = None
    _interp_stencil = 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

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

        self._X            = X
        self._Z            = Z
        self._metric_cart  = metric_cart
        self._curv_cart    = curv_cart

        self.Fs = [self._F, self._dF_r_fd, self._dF_R_fd]

    @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:
            self._metric_u_cart = utils.matrix_invert(self.metric_cart)
        return self._metric_u_cart
    @property
    def dmetric_cart(self):
        if self._dmetric_cart is None:
            self._dmetric_cart = curvature._calc_dmetric(self.X, self.Z, self.metric_cart, self._diff_op)
        return self._dmetric_cart
    @property
    def d2metric_cart(self):
        if self._d2metric_cart is None:
            self._d2metric_cart = curvature._calc_d2metric(self.X, self.Z, self.metric_cart,
                                                           self.dmetric_cart, self._diff_op, self._diff2_op)
        return self._d2metric_cart
    @property
    def dmetric_u_cart(self):
        if self._dmetric_u_cart is None:
            self._dmetric_u_cart = -np.einsum('ij...,km...,ljk...->lim...',
                                              self.metric_u_cart, self.metric_u_cart,
                                              self.dmetric_cart)
        return self._dmetric_u_cart
    @property
    def Gamma_cart(self):
        if self._Gamma_cart is None:
            self._Gamma_cart = curvature._calc_christoffel(self.metric_cart, self.metric_u_cart,
                                                           self.dmetric_cart)
        return self._Gamma_cart
    @property
    def dGamma_cart(self):
        if self._dGamma_cart is None:
            self._dGamma_cart = curvature._calc_dchristoffel(self.metric_cart,  self.metric_u_cart,
                                                             self.dmetric_cart, self.dmetric_u_cart,
                                                             self.d2metric_cart)
        return self._dGamma_cart
    @property
    def dcurv_cart(self):
        if self._dcurv_cart is None:
            self._dcurv_cart = curvature._calc_dmetric(self.X, self.Z, self.curv_cart, self._diff_op)
        return self._dcurv_cart

    def _interp_var(self, theta, r, x, z, var):
        origin = (self.Z[0, 0], self.X[0, 0])
        step   = (self.Z[1, 0] - self.Z[0, 0], self.X[0, 1] - self.X[0, 0])
        return interp.interp2d(origin, step, var, (z, x), self._interp_order)

    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._interp_var(theta, r, x, z, self.metric_cart[idx])
        return _tensor_transform_cart_spherical_d2(theta, r, x, z, metric_cart_dst)

    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._interp_var(theta, r, x, z, self.metric_u_cart[idx])
        return _tensor_transform_cart_spherical_u2(theta, r, x, z, metric_u_cart_dst)

    def _spherical_dmetric_u(self, theta, r, x, z, metric, metric_u, 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 -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._interp_var(theta, r, x, z, self.curv_cart[idx])
        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._interp_var(theta, r, x, z, self.curv_cart[idx])
        for idx in it.product(range(3), repeat = 3):
            dcurv_cart_dst[idx] = self._interp_var(theta, r, x, z, self.dcurv_cart[idx])
        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._interp_var(theta, r, x, z, self.Gamma_cart[idx])
        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._interp_var(theta, r, x, z, self.Gamma_cart[idx])
        for idx in it.product(range(3), repeat = 4):
            dGamma_cart_dst[idx] = self._interp_var(theta, r, x, z, self.dGamma_cart[idx])

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

    def _F(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)

        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)

        term1 = u2 * metric_u - df_ij
        term2 = Gamma[0] - dr * Gamma[1] + u * curv

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

        # handle the symmetry axis
        g = theta / np.pi
        for idx in np.where(np.abs(g.round() - g) < 1e-12)[0]:
            stencil    = 8

            coords     = np.empty(stencil * 2)
            vals       = np.empty(stencil * 2)
            fact       = np.empty(stencil * 2)
            if idx + stencil < ret.shape[0]:
                for i in range(stencil):
                    coords[2 * i]     = theta[idx + i + 1]
                    coords[2 * i + 1] = 2 * theta[idx] - coords[2 * i]

                    vals[2 * i]       = ret[idx + i + 1]
                    vals[2 * i + 1]   = ret[idx + i + 1]
            elif idx - stencil >= 0:
                for i in range(stencil):
                    coords[2 * i]     = theta[idx - i - 1]
                    coords[2 * i + 1] = 2 * theta[idx] - coords[2 * i]

                    vals[2 * i]       = ret[idx - i - 1]
                    vals[2 * i + 1]   = ret[idx - i - 1]
            else:
                continue

            for i in range(len(fact)):
                f = 1.0
                for j in range(len(fact)):
                    if i == j:
                        continue
                    f *= (theta[idx] - coords[j]) / (coords[i] - coords[j])
                fact[i] = f

            ret[idx] = np.dot(fact, vals)

        return ret

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

        rp = r + eps
        rm = r - eps
        Fp = self._F(theta, (rp, R))
        Fm = self._F(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_u = self._spherical_dmetric_u(theta, r, x, z, metric, metric_u, Gamma)
        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._F(theta, (r, Rp))
        Fm = self._F(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 i, j in it.product(range(3), repeat = 2):
            metric_u[i, j] = self._interp_var(theta, r, x, z, self.metric_u_cart[i, j])

        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