import numpy             as np
from   numpy import pi   as PI

import scipy.integrate   as integrate
import scipy.interpolate as interp

def _sphere_area_int(theta, r, x_vals, z_vals, metric):
    st = np.sin(theta)
    ct = np.cos(theta)

    x = r * st
    z = r * ct

    metric_val = np.empty((3, 3))
    for i in range(3):
        for j in range(3):
            metric_interp = interp.RectBivariateSpline(x_vals, z_vals, metric[i, j])
            metric_val[i, j] = metric_interp(x, z)

    transform = np.array([[st,  r * ct, 0.0],
                          [0.0, 0.0, r * st],
                          [ct, -r * st, 0.0]])

    ms = np.einsum('ki,lj,kl', transform, transform, metric_val)
    return np.sqrt(ms[1, 1] * ms[2, 2])

def mass_adm_sphere(x, z, metric, r, dr = 0.1):
    """
    Calculate an estimate of the ADM mass based on falloff of sphere area
    (equation (A.26) in [1]).

    [1] Alcubierre (2008): Introduction to 3+1 Numerical Relativity

    :param array_like x: 1D array of x coordinates at which the values of the
                         metric are provided
    :param array_like z: 1D array of z coordinates at which the values of the
                         metric are provided
    :param array_like metric: 4D array of metric components at the grid defined by x
                              and z. metric[i, j, k, l] is the ij-th component
                              at spatial point (X=x[k], Z=z[l]).
    :param float r: The radius at which the mass is to be evaluated.
    :param float dr: optional step used for calculating dA/dr with finite
                     differences.
    :return: Estimate of the ADM mass of the spacetime.
    :rtype: float
    """
    A     = 2 * 2 * PI * integrate.quad(_sphere_area_int, 0, PI / 2, (r,      x, z, metric))[0]
    Am1   = 2 * 2 * PI * integrate.quad(_sphere_area_int, 0, PI / 2, (r - dr, x, z, metric))[0]
    Ap1   = 2 * 2 * PI * integrate.quad(_sphere_area_int, 0, PI / 2, (r + dr, x, z, metric))[0]
    dA_dr = (Ap1 - Am1) / (2.0 * dr)

    X, Z      = np.meshgrid(x, z)
    R         = np.sqrt(X**2 + Z**2)
    grr       = ((X / R) ** 2) * metric[0, 0] + ((Z / R) ** 2) * metric[2, 2] + 2.0 * (X * Z / (R * R)) * metric[0, 2]
    grr       = np.where(np.isnan(grr), 0.0, grr)

    dst_theta = np.linspace(0, PI / 2, 128)
    grr_vals  = interp.RectBivariateSpline(x, z, grr)(r * np.sin(dst_theta), r * np.cos(dst_theta), grid = False)

    return (np.sqrt(A / (16.0 * PI)) *
            (1.0 - (dA_dr ** 2) / (16 * PI * A * np.average(grr_vals))))