#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numba
import numpy as np
from choclo.prism import gravity_u

@numba.jit(nopython=True, parallel=True)
def build_jacobian(coordinates, prisms):
    """
    Build a sensitivity matrix for gravity_u of a prism
    """
    # Unpack coordinates of the observation points
    easting, northing, upward = coordinates[:]
    # Initialize an empty 2d array for the sensitivity matrix
    n_coords = easting.size
    n_prisms = prisms.shape[0]
    jacobian = np.empty((n_coords, n_prisms), dtype=np.float64)
    # Compute the gravity_u field that each prism generate on every observation
    # point, considering that they have a unit density
    for i in numba.prange(len(easting)):
        for j in range(prisms.shape[0]):
            jacobian[i, j] = gravity_u(
                easting[i],
                northing[i],
                upward[i],
                prisms[j, 0],
                prisms[j, 1],
                prisms[j, 2],
                prisms[j, 3],
                prisms[j, 4],
                prisms[j, 5],
                1.0,
            )
    return jacobian


# In[2]:


easting = np.linspace(-5.0, 5.0, 21)
northing = np.linspace(-4.0, 4.0, 21)
easting, northing = np.meshgrid(easting, northing)
upward = 10 * np.ones_like(easting)

coordinates = (easting.ravel(), northing.ravel(), upward.ravel())


# In[3]:


prisms = np.array(
    [
        [-10.0, 0.0, -7.0, 0.0, -15.0, -10.0],
        [-10.0, 0.0, 0.0, 7.0, -25.0, -15.0],
        [0.0, 10.0, -7.0, 0.0, -20.0, -13.0],
        [0.0, 10.0, 0.0, 7.0, -12.0, -8.0],
    ]
)


# In[4]:


jacobian = build_jacobian(coordinates, prisms)
jacobian


# In[5]:


# Define densities for the prisms
densities = np.array([200.0, 300.0, -100.0, 400.0])

# Compute result
g_u = jacobian @ densities


# In[6]:


@numba.jit(nopython=True, parallel=True)
def gravity_upward_parallel(coordinates, prisms, densities):
    """
    Compute the upward component of the acceleration of a set of prisms
    """
    # Unpack coordinates of the observation points
    easting, northing, upward = coordinates[:]
    # Initialize a result array full of zeros
    result = np.zeros_like(easting, dtype=np.float64)
    # Compute the upward component that every prism generate on each
    # observation point
    for i in numba.prange(len(easting)):
        for j in range(prisms.shape[0]):
            result[i] += gravity_u(
                easting[i],
                northing[i],
                upward[i],
                prisms[j, 0],
                prisms[j, 1],
                prisms[j, 2],
                prisms[j, 3],
                prisms[j, 4],
                prisms[j, 5],
                densities[j],
            )
    return result


# In[7]:


expected = gravity_upward_parallel(coordinates, prisms, densities)
np.allclose(g_u, expected)