# 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)
#
"""
Kernel functions for rectangular prisms
"""
import numpy as np
from numba import jit


@jit(nopython=True)
def kernel_pot(easting, northing, upward, radius):
    r"""
    Kernel for the potential field due to a rectangular prism

    Evaluates the integration kernel for the potential field generated by
    a prism [Nagy2000] on a single vertex of the prism. The coordinates that
    must be passed are shifted coordinates: the coordinates of the vertex from
    a Cartesian coordinate system whose origin is located in the observation
    point.

    This function makes use of a safe natural logarithmic function and a safe
    arctangent function [Fukushima2020]_ that guarantee a good accuracy on
    every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the numerical kernel function for the potential field due to
        a rectangular prism evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_V(x, y, z) &=
            x y \, \operatorname{safe-ln} (z + r)
            + y z \, \operatorname{safe-ln} (x + r)
            + z x \, \operatorname{safe-ln} (y + r) \\
            & - \frac{x^2}{2} \operatorname{safe-arctan} \left( yz, xr \right)
            - \frac{y^2}{2} \operatorname{safe-arctan} \left( zx, yr \right)
            - \frac{z^2}{2} \operatorname{safe-arctan} \left( xy, zr \right)

    where

    .. math::

        \operatorname{safe-ln}(x) =
        \begin{cases}
            0 & |x| < 10^{-10} \\
            \ln (x)
        \end{cases}

    and

    .. math::

        \operatorname{safe-arctan} \left( y, x \right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    kernel = (
        easting * northing * _safe_log(upward, radius)
        + northing * upward * _safe_log(easting, radius)
        + easting * upward * _safe_log(northing, radius)
        - 0.5 * easting**2 * _safe_atan2(upward * northing, easting * radius)
        - 0.5 * northing**2 * _safe_atan2(upward * easting, northing * radius)
        - 0.5 * upward**2 * _safe_atan2(easting * northing, upward * radius)
    )
    return kernel


@jit(nopython=True)
def kernel_e(easting, northing, upward, radius):
    r"""
    Kernel for easting component of the gradient due to a rectangular prism

    Evaluates the integration kernel for the easting component of the gradient
    of the potential field generated by a prism [Nagy2000]_ on a single vertex
    of the prism. The coordinates that must be passed are shifted coordinates:
    the coordinates of the vertex from a Cartesian coordinate system whose
    origin is located in the observation point.

    This function makes use of a safe natural logarithmic function and a safe
    arctangent function [Fukushima2020]_ that guarantee a good accuracy on
    every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the numerical kernel function for the easting component of the
        gradient of the potential field due to a rectangular prism evaluated on
        a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_x(x, y, z) =
            -\left[
            y \, \operatorname{safe-ln} (z + r)
            + z \, \operatorname{safe-ln} (y + r)
            - x \, \operatorname{safe-arctan} \left( yz, xr \right)
            \right]

    where

    .. math::

        \operatorname{safe-ln}(x) =
        \begin{cases}
            0 & |x| < 10^{-10} \\
            \ln (x)
        \end{cases}

    and

    .. math::

        \operatorname{safe-arctan} \left( y, x \right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    .. important::

        In the first equation a minus sign has been added to the one obtained
        by [Nagy2000]_ in order to compute the numerical kernel for the
        **eastward** component instead for the westward one.

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    kernel = -(
        northing * _safe_log(upward, radius)
        + upward * _safe_log(northing, radius)
        - easting * _safe_atan2(northing * upward, easting * radius)
    )
    return kernel


@jit(nopython=True)
def kernel_n(easting, northing, upward, radius):
    r"""
    Kernel for northing component of the gradient due to a rectangular prism

    Evaluates the integration kernel for the northing component of the gradient
    of the potential field generated by a prism [Nagy2000]_ on a single vertex
    of the prism. The coordinates that must be passed are shifted coordinates:
    the coordinates of the vertex from a Cartesian coordinate system whose
    origin is located in the observation point.

    This function makes use of a safe natural logarithmic function and a safe
    arctangent function [Fukushima2020]_ that guarantee a good accuracy on
    every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the numerical kernel function for the northing component of
        the gradient of the potential field due to a rectangular prism
        evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_y(x, y, z) =
            -\left[
            z \, \operatorname{safe-ln} (x + r)
            + x \, \operatorname{safe-ln} (z + r)
            - y \, \operatorname{safe-arctan} \left( zx, yr \right)
            \right]

    where

    .. math::

        \operatorname{safe-ln}(x) =
        \begin{cases}
            0 & |x| < 10^{-10} \\
            \ln (x)
        \end{cases}

    and

    .. math::

        \operatorname{safe-arctan} \left( y, x \right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    .. important::

        In the first equation a minus sign has been added to the one obtained
        by [Nagy2000]_ in order to compute the numerical kernel for the
        **northward** component instead for the southward one.

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    kernel = -(
        upward * _safe_log(easting, radius)
        + easting * _safe_log(upward, radius)
        - northing * _safe_atan2(easting * upward, northing * radius)
    )
    return kernel


@jit(nopython=True)
def kernel_u(easting, northing, upward, radius):
    r"""
    Kernel for upward component of the gradient due to a rectangular prism

    Evaluates the integration kernel for the upward component of the gradient
    of the potential field generated by a prism [Nagy2000]_ on a single vertex
    of the prism. The coordinates that must be passed are shifted coordinates:
    the coordinates of the vertex from a Cartesian coordinate system whose
    origin is located in the observation point.

    This function makes use of a safe natural logarithmic function and a safe
    arctangent function [Fukushima2020]_ that guarantee a good accuracy on
    every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the numerical kernel function for the upward component of the
        gradient of the potential field due to a rectangular prism evaluated on
        a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_z(x, y, z) =
            - \left[
            x \, \operatorname{safe-ln} (y + r)
            + y \, \operatorname{safe-ln} (x + r)
            - z \, \operatorname{safe-arctan} \left( xy, zr \right)
            \right]

    where

    .. math::

        \operatorname{safe-ln}(x) =
        \begin{cases}
            0 & |x| < 10^{-10} \\
            \ln (x)
        \end{cases}

    and

    .. math::

        \operatorname{safe-arctan} \left( y, x \right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    .. important::

        In the first equation a minus sign has been added to the one obtained
        by [Nagy2000]_ in order to compute the numerical kernel for the
        **upward** component instead for the downward one.


    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    # The minus sign is to return the kernel for the upward component instead
    # of the downward one.
    kernel = -(
        easting * _safe_log(northing, radius)
        + northing * _safe_log(easting, radius)
        - upward * _safe_atan2(easting * northing, upward * radius)
    )
    return kernel


@jit(nopython=True)
def kernel_ee(easting, northing, upward, radius):
    r"""
    Kernel for easting-easting component of the tensor due to a prism

    Evaluates the integration kernel for the easting-easting component of the
    tensor generated by a prism [Nagy2000]_ on a single vertex of the prism.
    The coordinates that must be passed are shifted coordinates: the
    coordinates of the vertex from a Cartesian coordinate system whose origin
    is located in the observation point.

    This function makes use of a safe arctangent function [Fukushima2020]_ that
    guarantee a good accuracy on every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the kernel function for the easting-easting component of the
        tensor due to a rectangular prism evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_{xx}(x, y, z) = - \operatorname{safe-arctan} \left( yz, xr \right)

    where

    .. math::

        \operatorname{safe-arctan} \left( y, x \right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    return -_safe_atan2(northing * upward, easting * radius)


@jit(nopython=True)
def kernel_nn(easting, northing, upward, radius):
    r"""
    Kernel for northing-northing component of the tensor due to a prism

    Evaluates the integration kernel for the northing-northing component of the
    tensor generated by a prism [Nagy2000]_ on a single vertex of the prism.
    The coordinates that must be passed are shifted coordinates: the
    coordinates of the vertex from a Cartesian coordinate system whose origin
    is located in the observation point.

    This function makes use of a safe arctangent function [Fukushima2020]_ that
    guarantee a good accuracy on every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the kernel function for the northing-northing component of the
        tensor due to a rectangular prism evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_{yy}(x, y, z) = - \operatorname{safe-arctan} \left( zx, yr \right)

    where

    .. math::

        \operatorname{safe-arctan} \left( y, x \right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    return -_safe_atan2(easting * upward, northing * radius)


@jit(nopython=True)
def kernel_uu(easting, northing, upward, radius):
    r"""
    Kernel for upward-upward component of the tensor due to a prism

    Evaluates the integration kernel for the upward-upward component of the
    tensor generated by a prism [Nagy2000]_ on a single vertex of the prism.
    The coordinates that must be passed are shifted coordinates: the
    coordinates of the vertex from a Cartesian coordinate system whose origin
    is located in the observation point.

    This function makes use of a safe arctangent function [Fukushima2020]_ that
    guarantee a good accuracy on every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the kernel function for the upward-upward component of the
        tensor due to a rectangular prism evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_{zz}(x, y, z) = - \operatorname{safe-arctan} \left( xy, zr \right)

    where

    .. math::

        \operatorname{safe-arctan} \left( y, x \right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    return -_safe_atan2(easting * northing, upward * radius)


@jit(nopython=True)
def kernel_en(easting, northing, upward, radius):
    r"""
    Kernel for easting-northing component of the tensor due to a prism

    Evaluates the integration kernel for the easting-northing component of the
    tensor generated by a prism [Nagy2000]_ on a single vertex of the prism.
    The coordinates that must be passed are shifted coordinates: the
    coordinates of the vertex from a Cartesian coordinate system whose origin
    is located in the observation point.

    This function makes use of a safe natural logarithmic function
    [Fukushima2020]_ that guarantee a good accuracy on every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the kernel function for the easting-northing component of the
        tensor due to a rectangular prism evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_{xy}(x, y, z) = \operatorname{safe-ln} \left( z + r \right)

    where

    .. math::

        \operatorname{safe-ln}(x) =
        \begin{cases}
            0 & |x| < 10^{-10} \\
            \ln (x)
        \end{cases}

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    return _safe_log(upward, radius)


@jit(nopython=True)
def kernel_eu(easting, northing, upward, radius):
    r"""
    Kernel for easting-upward component of the tensor due to a prism

    Evaluates the integration kernel for the easting-upward component of the
    tensor generated by a prism [Nagy2000]_ on a single vertex of the prism.
    The coordinates that must be passed are shifted coordinates: the
    coordinates of the vertex from a Cartesian coordinate system whose origin
    is located in the observation point.

    This function makes use of a safe natural logarithmic function
    [Fukushima2020]_ that guarantee a good accuracy on every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the kernel function for the easting-upward component of the
        tensor due to a rectangular prism evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_{xz}(x, y, z) = \operatorname{safe-ln} \left( y + r \right)

    where

    .. math::

        \operatorname{safe-ln}(x) =
        \begin{cases}
            0 & |x| < 10^{-10} \\
            \ln (x)
        \end{cases}

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    return _safe_log(northing, radius)


@jit(nopython=True)
def kernel_nu(easting, northing, upward, radius):
    r"""
    Kernel for northing-upward component of the tensor due to a prism

    Evaluates the integration kernel for the northing-upward component of the
    tensor generated by a prism [Nagy2000]_ on a single vertex of the prism.
    The coordinates that must be passed are shifted coordinates: the
    coordinates of the vertex from a Cartesian coordinate system whose origin
    is located in the observation point.

    This function makes use of a safe natural logarithmic function
    [Fukushima2020]_ that guarantee a good accuracy on every observation point.

    Parameters
    ----------
    easting : float
        Shifted easting coordinate of the vertex of the prism. Must be in
        meters.
    northing : float
        Shifted northing coordinate of the vertex of the prism. Must be in
        meters.
    upward : float
        Shifted upward coordinate of the vertex of the prism. Must be in
        meters.
    radius : float
        Square root of the sum of the squares of the ``easting``, ``northing``
        and ``upward`` shifted coordinates.

    Returns
    -------
    kernel : float
        Value of the kernel function for the northing-upward component of the
        tensor due to a rectangular prism evaluated on a single vertex.

    Notes
    -----
    Computes the following numerical kernel on the passed *shifted
    coordinates*:

    .. math::

        k_{yz}(x, y, z) = \operatorname{safe-ln} \left( x + r \right)

    where

    .. math::

        \operatorname{safe-ln}(x) =
        \begin{cases}
            0 & |x| < 10^{-10} \\
            \ln (x)
        \end{cases}

    References
    ----------
    - [Nagy2000]_
    - [Nagy2002]_
    - [Fukushima2020]_
    """
    return _safe_log(easting, radius)


@jit(nopython=True)
def _safe_atan2(y, x):
    r"""
    Principal value of the arctangent expressed as a two variable function

    This modification has to be made to the arctangent function so the
    gravitational field of the prism satisfies the Poisson's equation.
    Therefore, it guarantees that the fields satisfies the symmetry properties
    of the prism. This modified function has been defined according to
    [Fukushima2020]_.

    Notes
    -----

    .. math::

        \operatorname{safe-arctan} \left(y, x\right) =
        \begin{cases}
            \text{arctan}\left( \frac{y}{x} \right) & x \ne 0 \\
            \frac{\pi}{2} & x = 0 \quad \text{and} \quad y > 0 \\
            -\frac{\pi}{2} & x = 0 \quad \text{and} \quad y < 0 \\
            0 & x = 0 \quad \text{and} \quad y = 0 \\
        \end{cases}

    References
    ----------
    - [Fukushima2020]_
    """
    if x == 0:
        if y > 0:
            result = np.pi / 2
        elif y < 0:
            result = -np.pi / 2
        else:
            result = 0
        return result
    return np.arctan(y / x)


@jit(nopython=True)
def _safe_log(x, r):
    r"""
    Safe log function to use in the prism kernels

    Evaluates the :math:`\ln{x + r}` where :math:`x` is one of the shifted
    coordinate of the prism vertex and :math:`r` is the Euclidean distance
    (always non-negative) from the vertex to the observation point.

    .. math::

        \text{safe_ln}(x, r) =
        \begin{cases}
            0 & x = 0, r = 0 \\
            \ln(x + r) & x \ge 0 \\
            \ln((r^2 - x^2) / (r - x)) & x < 0, r \ne -x \\
            -\ln(-2 x) & x < 0, r = -x
        \end{cases}

    This function returns 0 when the observation point is located on the vertex
    of the prism (:math:`r=0`); and two modified versions in case that
    :math:`x` is negative: if :math:`x = -r` then the :math:`\ln{x + r}` can be
    replaced by :math:`-\ln{|x| + r} = -\ln(-2x)`, and for any other negative
    value of :math:`x` it returns :math:`\ln((r^2 - x^2) / (r - x))` which
    helps by reducing the floating point errors ([Fukushima2020_]).
    This modified version was inspired by [Nagy2000] and [Fukushima2020_]:

    References
    ----------
    - [Fukushima2020]_
    """
    if r == 0:
        return 0
    if x < 0:
        if r == -x:
            result = -np.log(-2 * x)
        else:
            result = np.log((r**2 - x**2) / (r - x))
        return result
    return np.log(x + r)
