# Copyright (c) 2022 The Choclo Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
#
# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
#
"""
Forward modelling function for point sources
"""
from numba import jit

from ..constants import GRAVITATIONAL_CONST
from ..utils import distance_cartesian
from ._kernels import (
    kernel_e,
    kernel_ee,
    kernel_en,
    kernel_eu,
    kernel_n,
    kernel_nn,
    kernel_nu,
    kernel_pot,
    kernel_u,
    kernel_uu,
)


@jit(nopython=True)
def gravity_pot(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Gravitational potential field due to a point source

    Returns the gravitational potential field produced by a single point source
    on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    potential : float
        Gravitational potential field generated by the point source on the
        observation point in :math:`\text{J}/\text{kg}`.

    Notes
    -----
    Returns the gravitational potential field :math:`V(\mathbf{p})` on the
    observation point :math:`\mathbf{p} = (x_p, y_p, z_p)` generated by
    a single point source located in :math:`\mathbf{q} = (x_q, y_q, z_q)` and
    mass :math:`m`.

    .. math::

        V(\mathbf{p}) = G m \frac{1}{\lVert \mathbf{p} - \mathbf{q} \rVert_2^2}

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_pot(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_e(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Easting component of the gravitational acceleration due to a point source

    Returns the easting component of the gravitational acceleration produced by
    a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_e : float
        Easting component of the gravitational acceleration generated by the
        point source on the observation point in :math:`\text{m}/\text{s}^2`.

    Notes
    -----
    Returns the easting component :math:`g_x(\mathbf{p})` of the gravitational
    acceleration :math:`\mathbf{g}` on the observation point :math:`\mathbf{p}
    = (x_p, y_p, z_p)` generated by a single point source located in
    :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_x(\mathbf{p}) = - G m
            \frac{
                x_p - x_q
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^2
            }

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_e(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_n(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Northing component of the gravitational acceleration due to a point source

    Returns the northing component of the gravitational acceleration produced
    by a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_n : float
        Northing component of the gravitational acceleration generated by the
        point source on the observation point in :math:`\text{m}/\text{s}^2`.

    Notes
    -----
    Returns the northing component :math:`g_y(\mathbf{p})` of the gravitational
    acceleration :math:`\mathbf{g}` on the observation point :math:`\mathbf{p}
    = (x_p, y_p, z_p)` generated by a single point source located in
    :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_y(\mathbf{p}) = - G m
            \frac{
                y_p - y_q
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^2
            }

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_n(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_u(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Upward component of the gravitational acceleration due to a point source

    Returns the upward component of the gravitational acceleration produced by
    a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_u : float
        Upward component of the gravitational acceleration generated by the
        point source on the observation point in :math:`\text{m}/\text{s}^2`.

    Notes
    -----
    Returns the upward component :math:`g_z(\mathbf{p})` of the gravitational
    acceleration :math:`\mathbf{g}` on the observation point :math:`\mathbf{p}
    = (x_p, y_p, z_p)` generated by a single point source located in
    :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_z(\mathbf{p}) = - G m
            \frac{
                z_p - z_q
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^2
            }

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_u(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_ee(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Easting-easting component of the gravitational tensor due to a point source

    Returns the easting-easting component of the gravitational tensor produced
    by a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_ee : float
        Easting-easting component of the gravitational tensor generated by the
        point source on the observation point in :math:`\text{s}^{-2}`.

    Notes
    -----
    Returns the easting-easting component :math:`g_{xx}(\mathbf{p})` of the
    gravitational tensor :math:`\mathbf{T}` on the observation point
    :math:`\mathbf{p} = (x_p, y_p, z_p)` generated by a single point source
    located in :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_{xx}(\mathbf{p}) = G m
            \left[
            \frac{
                3 (x_p - x_q)^2
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^5
            }
            - \frac{
                1
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^3
            }
            \right]

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_ee(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_nn(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Northing-northing component of the gravitational tensor due to point source

    Returns the northing-northing component of the gravitational tensor
    produced by a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_nn : float
        Northing-northing component of the gravitational tensor generated by
        the point source on the observation point in :math:`\text{s}^{-2}`.

    Notes
    -----
    Returns the northing-northing component :math:`g_{yy}(\mathbf{p})` of the
    gravitational tensor :math:`\mathbf{T}` on the observation point
    :math:`\mathbf{p} = (x_p, y_p, z_p)` generated by a single point source
    located in :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_{yy}(\mathbf{p}) = G m
            \left[
            \frac{
                3 (y_p - y_q)^2
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^5
            }
            - \frac{
                1
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^3
            }
            \right]

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_nn(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_uu(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Upward-upward component of the gravitational tensor due to a point source

    Returns the upward-upward component of the gravitational tensor
    produced by a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_uu : float
        Upward-upward component of the gravitational tensor generated by
        the point source on the observation point in :math:`\text{s}^{-2}`.

    Notes
    -----
    Returns the upward-upward component :math:`g_{zz}(\mathbf{p})` of the
    gravitational tensor :math:`\mathbf{T}` on the observation point
    :math:`\mathbf{p} = (x_p, y_p, z_p)` generated by a single point source
    located in :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_{zz}(\mathbf{p}) = G m
            \left[
            \frac{
                3 (z_p - z_q)^2
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^5
            }
            - \frac{
                1
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^3
            }
            \right]

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_uu(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_en(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Easting-northing component of the gravitational tensor due to point source

    Returns the easting-northing component of the gravitational tensor
    produced by a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_en : float
        Easting-northing component of the gravitational tensor generated by
        the point source on the observation point in :math:`\text{s}^{-2}`.

    Notes
    -----
    Returns the easting-northing component :math:`g_{xy}(\mathbf{p})` of the
    gravitational tensor :math:`\mathbf{T}` on the observation point
    :math:`\mathbf{p} = (x_p, y_p, z_p)` generated by a single point source
    located in :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_{xy}(\mathbf{p}) = G m
            \left[
            \frac{
                3 (x_p - x_q) (y_p - y_q)
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^5
            }
            \right]

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_en(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_eu(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Easting-upward component of the gravitational tensor due to point source

    Returns the easting-upward component of the gravitational tensor
    produced by a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_eu : float
        Easting-upward component of the gravitational tensor generated by
        the point source on the observation point in :math:`\text{s}^{-2}`.

    Notes
    -----
    Returns the easting-upward component :math:`g_{xz}(\mathbf{p})` of the
    gravitational tensor :math:`\mathbf{T}` on the observation point
    :math:`\mathbf{p} = (x_p, y_p, z_p)` generated by a single point source
    located in :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_{xz}(\mathbf{p}) = G m
            \left[
            \frac{
                3 (x_p - x_q) (z_p - z_q)
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^5
            }
            \right]

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_eu(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel


@jit(nopython=True)
def gravity_nu(easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, mass):
    r"""
    Northing-upward component of the gravitational tensor due to point source

    Returns the northing-upward component of the gravitational tensor
    produced by a single point source on a single computation point

    Parameters
    ----------
    easting_p, northing_p, upward_p : float
        Easting, northing and upward coordinates of the observation point in
        meters.
    easting_q, northing_q, upward_q : float
        Easting, northing and upward coordinates of the point source in meters.
    mass : float
        Mass of the point source in kilograms.

    Returns
    -------
    g_nu : float
        Northing-upward component of the gravitational tensor generated by the
        point source on the observation point in :math:`\text{s}^{-2}`.

    Notes
    -----
    Returns the northing-upward component :math:`g_{yz}(\mathbf{p})` of the
    gravitational tensor :math:`\mathbf{T}` on the observation point
    :math:`\mathbf{p} = (x_p, y_p, z_p)` generated by a single point source
    located in :math:`\mathbf{q} = (x_q, y_q, z_q)` and mass :math:`m`.

    .. math::

        g_{yz}(\mathbf{p}) = G m
            \left[
            \frac{
                3 (y_p - y_q) (z_p - z_q)
            }{
                \lVert \mathbf{p} - \mathbf{q} \rVert_2^5
            }
            \right]

    where :math:`\lVert \cdot \rVert_2` refer to the :math:`L_2` norm (the
    Euclidean distance between :math:`\mathbf{p}` and :math:`\mathbf{q}`)
    and :math:`G` is the Universal Gravitational Constant.
    """
    distance = distance_cartesian(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q
    )
    kernel = kernel_nu(
        easting_p, northing_p, upward_p, easting_q, northing_q, upward_q, distance
    )
    return GRAVITATIONAL_CONST * mass * kernel
