'''
This code presents an approach for implementing the gravitational field
produced by a sphere by using the analytical formulas.

References

'''

import numpy as np
from numba import jit

#: The gravitational constant in m^3 kg^{-1} s^{-1}
GRAVITATIONAL_CONST = 0.00000000006673

def gravitational(coordinates,spheres,density,field):
    """
    Computing the gravitational field generated by a sphere.
    """

    kernels = {"potential":kernel_potential,
               "g_z": kernel_g_z,
               "g_x": kernel_g_x,
               "g_y": kernel_g_y}

    # Checking input field
    if field not in kernels:
        raise ValueError("Gravitational field {} not recognized".format(field))

    # Checking input parameters
    if coordinates.ndim != 2:
        raise ValueError(
            "coordinates ndim ({}) ".format(coordinates.ndim)
            + "not equal to 2"
        )
    if coordinates.shape[0] != 3:
        raise ValueError(
            "Number of lines in coordinates ({}) ".format(coordinates.shape[0])
            + "not equal to 3"
        )
    if spheres.ndim != 2:
        raise ValueError(
            "prisms ndim ({}) ".format(spheres.ndim)
            + "not equal to 2"
        )
    if spheres.shape[1] != 4:
        raise ValueError(
            "Number of columns in spheres ({}) ".format(spheres.shape[1])
            + "not equal to 6"
        )
    if density.ndim != 1:
        raise ValueError(
            "density ndim ({}) ".format(density.ndim)
            + "not equal to 1"
        )
    if density.size != spheres.shape[0]:
        raise ValueError(
            "Number of elements in density ({}) ".format(density.size)
            + "mismatch the number of spheres ({})".format(spheres.shape[0])
        )

    # Create the array to store the result
    result = np.zeros(coordinates[0].size, dtype="float64")

    # Compute gravitational field
    jit_gravitational(coordinates, spheres, density, kernels[field], result)
    result *= GRAVITATIONAL_CONST
    # Convert from m/s^2 to mGal
    if field in ["potential", "g_z","g_x","g_y"]:
        result *= 1e5
    return result

@jit(nopython=True)
def jit_gravitational(coordinates,spheres,density,kernel,out):
    """
    Compute gravitational field at the observation points
    """

    # Iterate over coordinates
    for l in range(coordinates[0].size):
        # Iterate over spheres:
        for m in range(spheres.shape[0]):
            Y = spheres[m,0] - coordinates[0,l]
            X = spheres[m,1] - coordinates[1,l]
            Z = spheres[m,2] - coordinates[2,l]
            r3 = spheres[m,3]*spheres[m,3]*spheres[m,3]
            Volume = (4/3)*np.pi*r3
            out[l] += (density[m]*Volume*kernel(Y,X,Z))

@jit(nopython=True)
def kernel_potential(Y, X, Z):
    """
    Kernel function for potential gravitational field generated by a sphere
    """

    radius = np.sqrt(Y ** 2 + X ** 2 + Z ** 2)
    kernel = 1/radius

    return kernel

@jit(nopython=True)
def kernel_g_x(Y, X, Z):
    """
    Kernel for x component of gravitational acceleration of a sphere
    """
    radius = np.sqrt(Y ** 2 + X ** 2 + Z ** 2)
    r3 = radius*radius*radius
    kernel = X /r3
    return kernel

@jit(nopython=True)
def kernel_g_y(Y, X, Z):
    """
    Kernel for y component of gravitational acceleration of a sphere
    """
    radius = np.sqrt(Y ** 2 + X ** 2 + Z ** 2)
    r3 = radius*radius*radius
    kernel = Y /r3
    return kernel

@jit(nopython=True)
def kernel_g_z(Y, X, Z):
    """
    Kernel for z component of gravitational acceleration of a sphere
    """
    radius = np.sqrt(Y ** 2 + X ** 2 + Z ** 2)
    r3 = radius*radius*radius
    kernel = Z/r3
    return kernel
