import numpy as np
import scipy.interpolate

def _invert_matrix(mat):
    inv = np.empty_like(mat)
    for i in xrange(mat.shape[-1]):
        try:
            inv[:, :, i] = np.linalg.inv(mat[:, :, i])
        except:
            #print i, mat[:, :, i]
            inv[:, :, i] = np.array([[1.0, 0.0, 0.0],
                                     [0.0, 1.0, 0.0],
                                     [0.0, 0.0, 1.0]])
    return inv

def _tensor_transform_cart_spherical_d2(theta, r, tens):
    z = np.zeros_like(theta)
    o = z + 1.0
    jac = np.array([[np.sin(theta), r * np.cos(theta),   z ],
                    [z,             z,                   r * np.sin(theta)],
                    [np.cos(theta), -r * np.sin(theta),  z ]])
    ret = np.zeros_like(tens)
    for i in xrange(3):
        for j in xrange(3):
            for k in xrange(3):
                for l in xrange(3):
                    ret[i, j] += jac[k, i] * jac[l, j] * tens[k, l]

    return ret;

def _calc_metric(X, Z, metric, curv, r, theta):
    R     = np.sqrt(X ** 2 + Z ** 2)
    Theta = np.where(np.abs(X) > 1e-12, np.arccos(Z / R), 0)

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

    metric_src_spherical = _tensor_transform_cart_spherical_d2(Theta, R, metric)
    curv_src_spherical   = _tensor_transform_cart_spherical_d2(Theta, R, curv)

    metric_val = np.empty((3, 3) + theta.shape)
    curv_val   = np.empty_like(metric_val)
    dg_cart    = np.empty((2,) + metric_val.shape)
    for i in xrange(3):
        for j in xrange(3):
            metric_interp = scipy.interpolate.RectBivariateSpline(Z[:, 0], X[0], metric_src_spherical[i, j])
            metric_val[i, j] = metric_interp.ev(z, x)
            dg_cart[0, i, j] = metric_interp.ev(z, x, 0, 1)
            dg_cart[1, i, j] = metric_interp.ev(z, x, 1, 0)

            curv_interp    = scipy.interpolate.RectBivariateSpline(Z[:, 0], X[0], curv_src_spherical[i, j])
            curv_val[i, j] = curv_interp.ev(z, x)

    dg = np.empty_like(dg_cart)
    dg[0] = dg_cart[0] * np.sin(theta) + dg_cart[1] * np.cos(theta)
    dg[1] = z * dg_cart[0] - x * dg_cart[1]

    metric_u_val = _invert_matrix(metric_val)

    Gamma_r_val     = np.zeros_like(metric_val)
    Gamma_theta_val = np.zeros_like(metric_val)
    for i in xrange(3):
        for j in xrange(3):
            val_r     = np.zeros_like(theta)
            val_theta = np.zeros_like(theta)
            if i != 2:
                val_r     += dg[i, j, 0]
                val_theta += dg[i, j, 1]
            if j != 2:
                val_r     += dg[j, i, 0]
                val_theta += dg[j, i, 1]
            val_r     -= dg[0, i, j]
            val_theta -= dg[1, i, j]
            Gamma_r_val[i, j]     = 0.5 * (metric_u_val[0, 0] * val_r     + metric_u_val[0, 1] * val_theta)
            Gamma_theta_val[i, j] = 0.5 * (metric_u_val[1, 1] * val_theta + metric_u_val[0, 1] * val_r)

    return metric_val, metric_u_val, curv_val, Gamma_r_val, Gamma_theta_val


def _horiz_axi_F(theta, u, args):
    r, dr = u

    X, Z, metric, curv = args
    metric_val, metric_u_val, curv_val, Gamma_r_val, Gamma_theta_val = _calc_metric(X, Z, metric, curv, r, theta)

    u2 = metric_u_val[1, 1] * dr * dr - 2 * metric_u_val[0, 1] * dr + metric_u_val[0, 0]

    df_ij = np.empty_like(metric_u_val)
    for i in xrange(3):
        for j in xrange(3):
            df_ij[i, j] = (metric_u_val[0, i] - metric_u_val[1, i] * dr) * (metric_u_val[0, j] - metric_u_val[1, j] * dr)

    term1 = u2 * metric_u_val - df_ij
    term2 = Gamma_r_val - dr * Gamma_theta_val + np.sqrt(u2) * curv_val

    return -np.einsum('ij...,ij...', term1, term2) / (metric_u_val[0, 0] * metric_u_val[1, 1] - metric_u_val[0, 1] * metric_u_val[0, 1])

def _horiz_axi_dF_r(theta, u, args):
    eps = 1e-4
    r, R = u

    rp = r + eps
    rm = r - eps
    Fp = _horiz_axi_F(theta, (rp, R), args)
    Fm = _horiz_axi_F(theta, (rm, R), args)
    return (Fp - Fm) / (2 * eps)

def _horiz_axi_dF_R(theta, u, args):
    eps = 1e-4
    r, R = u

    Rp = R + eps
    Rm = R - eps
    Fp = _horiz_axi_F(theta, (r, Rp), args)
    Fm = _horiz_axi_F(theta, (r, Rm), args)
    return (Fp - Fm) / (2 * eps)

"""
The right-hand side function (and its derivatives wrt u and u') defining the apparent horizon equation in axial
symmetry. Intended to be passed to the nonlin_solve_1d solver.

The extra args is a tuple of:
    - a meshgrid of the cartesian r coordinates
    - a meshgrid of the cartesian z coordinates
    - the spatial metric in cartesian coordinates at the provided coordinates -- a numpy aray of (3, 3) + z.shape
    - the extrinsic curvature in cartesian coordinates at the provided coordinates -- a numpy aray of (3, 3) + z.shape
"""
Fs = (_horiz_axi_F, _horiz_axi_dF_r, _horiz_axi_dF_R)