"""Magnetic field calculation routines"""

from numpy import float64, ascontiguousarray, zeros, hstack, newaxis, array, pi
from numpy.typing import NDArray

# Create bindings for calculation engine written in Rust
from ._oersted import (
    _bfield_direct,
    _bfield_octree,
    _bfield_dualtree,
    _bfield_hexahedron,
    _hfield_dipole,
    _hfield_tetrahedrons,
    _hfield_tetrahedrons_direct,
    _hfield_dipole_tetrahedrons,
)

# For typing; currently unused
Nx3Array = NDArray[float64]


def bfield_direct(
    centroids: NDArray[float64], vol: NDArray[float64], jdensity: NDArray[float64], targets: NDArray[float64], nthreads: int = 0
) -> NDArray[float64]:
    """Compute the magnetic field generated by a collection of current-carrying
        finite elements at `targets` using the direct summation Biot-Savart method.

    Args:
        centroids: [m] Nx3 array, location of source element centroids in 3D (cartesian) space
        vol: [m^3] N-length array, volume of each source element
        jdensity: [A/m^2] Nx3 array, current density vectors for each source element
        targets: [m] Nx3 array, describes location of target points at which to calculate magnetic fields
        nthreads: number of threads to use for the calculation (default: all available)

    Returns:
        [T] Nx3 array, magnetic flux density at each specified target point
    """

    ntargets = targets.shape[0]
    bx = zeros(ntargets)
    by = zeros(ntargets)
    bz = zeros(ntargets)

    _bfield_direct(
        ascontiguousarray(centroids[:, 0]),
        ascontiguousarray(centroids[:, 1]),
        ascontiguousarray(centroids[:, 2]),
        ascontiguousarray(vol[:]),
        ascontiguousarray(jdensity[:, 0]),
        ascontiguousarray(jdensity[:, 1]),
        ascontiguousarray(jdensity[:, 2]),
        ascontiguousarray(targets[:, 0]),
        ascontiguousarray(targets[:, 1]),
        ascontiguousarray(targets[:, 2]),
        ascontiguousarray(bx[:]),
        ascontiguousarray(by[:]),
        ascontiguousarray(bz[:]),
        nthreads,
    )

    return hstack((bx[:, newaxis], by[:, newaxis], bz[:, newaxis]))


def bfield_octree(
    centroids: NDArray[float64],
    vol: NDArray[float64],
    jdensity: NDArray[float64],
    targets: NDArray[float64],
    theta: float = 0.5,
    leaf_threshold: int = 16,
    nthreads: int = 0,
) -> NDArray[float64]:
    """Compute the magnetic field generated by a collection of current-carrying
        finite elements at `targets` using the Barnes-Hut (octree) Biot-Savart method.

    Args:
        centroids: [m] Nx3 array, location of source element centroids in 3D (cartesian) space
        vol: [m^3] N-length array, volume of each source element
        jdensity: [A/m^2] Nx3 array, current density vectors for each source element
        targets: [m] Nx3 array, describes location of target points at which to calculate magnetic fields
        theta: angle-opening parameter for Barnes-Hut; smaller values correspond to lower error
            at the expense of longer runtimes. Suggested values in the range of [0.1, 0.5]
        leaf_threshold: how many individual sources to aggregate per leaf
        nthreads: number of threads to use for the calculation (default: all available)

    Returns:
        [T] Nx3 array, magnetic flux density at each specified target point

    """

    ntargets = targets.shape[0]
    bx = zeros(ntargets)
    by = zeros(ntargets)
    bz = zeros(ntargets)

    _bfield_octree(
        ascontiguousarray(centroids[:, 0]),
        ascontiguousarray(centroids[:, 1]),
        ascontiguousarray(centroids[:, 2]),
        ascontiguousarray(vol[:]),
        ascontiguousarray(jdensity[:, 0]),
        ascontiguousarray(jdensity[:, 1]),
        ascontiguousarray(jdensity[:, 2]),
        ascontiguousarray(targets[:, 0]),
        ascontiguousarray(targets[:, 1]),
        ascontiguousarray(targets[:, 2]),
        ascontiguousarray(bx[:]),
        ascontiguousarray(by[:]),
        ascontiguousarray(bz[:]),
        theta,
        leaf_threshold,
        nthreads,
    )

    return (4 * pi * 10**-7) * hstack((bx[:, newaxis], by[:, newaxis], bz[:, newaxis]))


def bfield_dualtree(
    centroids: NDArray[float64],
    vol: NDArray[float64],
    jdensity: NDArray[float64],
    targets: NDArray[float64],
    theta_source: float = 0.5,
    theta_target: float = 0.1,
    leaf_threshold: int = 16,
    nthreads: int = 0,
) -> NDArray[float64]:
    """Compute the magnetic field generated by a collection of current-carrying
        finite elements at `targets` using the dual Barnes-Hut (octree) Biot-Savart method.

    Args:
        centroids: [m] Nx3 array, location of source element centroids in 3D (cartesian) space
        vol: [m^3] N-length array, volume of each source element
        jdensity: [A/m^2] Nx3 array, current density vectors for each source element
        targets: [m] Nx3 array, describes location of target points at which to calculate magnetic fields
        theta_source, theta_target: angle-opening parameter for Barnes-Hut; smaller values correspond to lower error
            at the expense of longer runtimes. Suggested values in the range of [0.1, 0.5]
        leaf_threshold: how many individual sources to aggregate per leaf
        nthreads: number of threads to use for the calculation (default: all available)

    Returns:
        [T] Nx3 array, magnetic flux density at each specified target point

    """

    ntargets = targets.shape[0]
    bx = zeros(ntargets)
    by = zeros(ntargets)
    bz = zeros(ntargets)

    _bfield_dualtree(
        ascontiguousarray(centroids[:, 0]),
        ascontiguousarray(centroids[:, 1]),
        ascontiguousarray(centroids[:, 2]),
        ascontiguousarray(vol[:]),
        ascontiguousarray(jdensity[:, 0]),
        ascontiguousarray(jdensity[:, 1]),
        ascontiguousarray(jdensity[:, 2]),
        ascontiguousarray(targets[:, 0]),
        ascontiguousarray(targets[:, 1]),
        ascontiguousarray(targets[:, 2]),
        ascontiguousarray(bx[:]),
        ascontiguousarray(by[:]),
        ascontiguousarray(bz[:]),
        theta_source,
        theta_target,
        leaf_threshold,
        nthreads,
    )

    return hstack((bx[:, newaxis], by[:, newaxis], bz[:, newaxis]))


def bfield_hexahedron(nodes: NDArray[float64], jdensity: NDArray[float64], target: NDArray[float64]) -> NDArray[float64]:

    b = _bfield_hexahedron(
        ascontiguousarray(nodes[:, 0]),
        ascontiguousarray(nodes[:, 1]),
        ascontiguousarray(nodes[:, 2]),
        ascontiguousarray(jdensity),
        ascontiguousarray(target),
    )

    return array([b[0], b[1], b[2]])


def bfield_tetrahedrons(
    nodes: NDArray[float64],
    centroids: NDArray[float64],
    vol: NDArray[float64],
    jdensity: NDArray[float64],
    targets: NDArray[float64],
    theta: float = 0.5,
    leaf_threshold: int = 1,
    nthreads: int = 0,
) -> NDArray[float64]:
    """Compute the magnetic flux density at a set of target points
        using a tetrahedral finite element mesh as a near-field source,
        and a point approximation for the far-field.

    Args:
        nodes: (12*N,) nodal coordinates of each element
        vol: (N,) volume of each element
        jdensity: (N,3) current density vector assumed constant over each element
        targets: (M,3) target point locations in 3d space
        theta: angle-opening criteria for barnes-hut

    Returns:
        (N,3) magnetic flux density at each target point

    """

    ntargets = targets.shape[0]
    hx = zeros(ntargets)
    hy = zeros(ntargets)
    hz = zeros(ntargets)

    _hfield_tetrahedrons(
        ascontiguousarray(nodes[:]),
        ascontiguousarray(centroids.ravel()),
        ascontiguousarray(vol[:]),
        ascontiguousarray(jdensity.ravel()),
        ascontiguousarray(targets[:, 0]),
        ascontiguousarray(targets[:, 1]),
        ascontiguousarray(targets[:, 2]),
        ascontiguousarray(hx[:]),
        ascontiguousarray(hy[:]),
        ascontiguousarray(hz[:]),
        theta,
        leaf_threshold,
        nthreads,
    )

    return (4 * pi * 10**-7) * hstack((hx[:, newaxis], hy[:, newaxis], hz[:, newaxis]))


def bfield_tetrahedrons_direct(
    nodes: NDArray[float64],
    centroids: NDArray[float64],
    vol: NDArray[float64],
    jdensity: NDArray[float64],
    targets: NDArray[float64],
    nthreads: int = 0,
) -> NDArray[float64]:
    """Compute the magnetic flux density at a set of target points
        using a tetrahedral finite element mesh as a source. This
        function performs the analytic (exact) integral for every
        element in the mesh.

    Args:
        nodes: (12*N,) nodal coordinates of each element
        centroids: (N,3) coordinates of each element centroid
        vol: (N,) volume of each element
        jdensity: (N,3) current density vector assumed constant over each element
        targets: (M,3) target point locations in 3d space

    Returns:
        (N,3) magnetic flux density at each target point

    """

    ntargets = targets.shape[0]
    hx = zeros(ntargets)
    hy = zeros(ntargets)
    hz = zeros(ntargets)

    _hfield_tetrahedrons_direct(
        ascontiguousarray(nodes[:]),
        ascontiguousarray(centroids.ravel()),
        ascontiguousarray(vol[:]),
        ascontiguousarray(jdensity.ravel()),
        ascontiguousarray(targets[:, 0]),
        ascontiguousarray(targets[:, 1]),
        ascontiguousarray(targets[:, 2]),
        ascontiguousarray(hx[:]),
        ascontiguousarray(hy[:]),
        ascontiguousarray(hz[:]),
        nthreads,
    )

    return (4 * pi * 10**-7) * hstack((hx[:, newaxis], hy[:, newaxis], hz[:, newaxis]))


def hfield_dipole(
    centroids: NDArray[float64],
    vol: NDArray[float64],
    moments: NDArray[float64],
    targets: NDArray[float64],
    theta: float = 0.5,
    leaf_threshold: int = 16,
    nthreads: int = 0,
) -> NDArray[float64]:
    """Compute the magnetic H-field generated by a collection of magnetized
        finite elements at `targets` using the Barnes-Hut (octree) Biot-Savart method.

    Args:
        centroids: [m] Nx3 array, location of source element centroids in 3D (cartesian) space
        vol: [m^3] N-length array, volume of each source element
        jdensity: [A/m^2] Nx3 array, current density vectors for each source element
        targets: [m] Nx3 array, describes location of target points at which to calculate magnetic fields
        theta: angle-opening parameter for Barnes-Hut; smaller values correspond to lower error
            at the expense of longer runtimes. Suggested values in the range of [0.1, 0.5]
        leaf_threshold: how many individual sources to aggregate per leaf
        nthreads: number of threads to use for the calculation (default: all available)

    Returns:
        [A/m] Nx3 array, magnetic field intensity [H-field] at each specified target point

    """

    ntargets = targets.shape[0]
    hx = zeros(ntargets)
    hy = zeros(ntargets)
    hz = zeros(ntargets)

    _hfield_dipole(
        ascontiguousarray(centroids[:, 0]),
        ascontiguousarray(centroids[:, 1]),
        ascontiguousarray(centroids[:, 2]),
        ascontiguousarray(vol[:]),
        ascontiguousarray(moments[:, 0]),
        ascontiguousarray(moments[:, 1]),
        ascontiguousarray(moments[:, 2]),
        ascontiguousarray(targets[:, 0]),
        ascontiguousarray(targets[:, 1]),
        ascontiguousarray(targets[:, 2]),
        ascontiguousarray(hx[:]),
        ascontiguousarray(hy[:]),
        ascontiguousarray(hz[:]),
        theta,
        leaf_threshold,
        nthreads,
    )

    return hstack((hx[:, newaxis], hy[:, newaxis], hz[:, newaxis]))


def hfield_dipole_tetrahedrons(
    nodes: NDArray[float64],
    centroids: NDArray[float64],
    vol: NDArray[float64],
    moments: NDArray[float64],
    targets: NDArray[float64],
    theta: float = 0.5,
    leaf_threshold=1,
    nthreads: int = 0,
) -> NDArray[float64]:
    """Compute the magnetic field intensity at a set of target points
        using a tetrahedral finite element mesh as a near-field dipole source,
        and a point approximation for the far-field.

    Args:
        nodes: (12*N,) nodal coordinates of each element
        vol: (N,) volume of each element
        moments: (N,3) current density vector assumed constant over each element
        targets: (M,3) target point locations in 3d space
        theta: angle-opening criteria for barnes-hut

    Returns:
        (N,3) magnetic flux density at each target point

    """

    ntargets = targets.shape[0]
    hx = zeros(ntargets)
    hy = zeros(ntargets)
    hz = zeros(ntargets)

    _hfield_dipole_tetrahedrons(
        ascontiguousarray(nodes[:]),
        ascontiguousarray(centroids.ravel()),
        ascontiguousarray(vol[:]),
        ascontiguousarray(moments.ravel()),
        ascontiguousarray(targets[:, 0]),
        ascontiguousarray(targets[:, 1]),
        ascontiguousarray(targets[:, 2]),
        ascontiguousarray(hx[:]),
        ascontiguousarray(hy[:]),
        ascontiguousarray(hz[:]),
        theta,
        leaf_threshold,
        nthreads,
    )

    return (4 * pi * 10**-7) * hstack((hx[:, newaxis], hy[:, newaxis], hz[:, newaxis]))
