"""Quasi-steady electromagnetics calcs"""

from typing import Literal

import numpy as np
from interpn import MulticubicRectilinear
from numpy.typing import NDArray

from cfsem.bindings import (
    body_force_density_circular_filament_cartesian,
    body_force_density_linear_filament,
    filament_helix_path,
    flux_circular_filament,
    flux_density_circular_filament,
    flux_density_circular_filament_cartesian,
    flux_density_dipole,
    flux_density_dipole_hierarchical,
    flux_density_linear_filament,
    flux_density_linear_filament_hierarchical,
    flux_density_triangle_mesh_mapping,
    triangle_mesh_current_density,
    triangle_mesh_force_mapping,
    triangle_mesh_force_mapping_from_circular_filaments,
    triangle_mesh_force_mapping_from_dipoles,
    triangle_mesh_force_mapping_from_linear_filaments,
    triangle_mesh_flux_linkage_mapping_from_dipoles,
    triangle_mesh_inductance_mapping_from_circular_filaments,
    triangle_mesh_inductance_mapping_from_linear_filaments,
    triangle_mesh_inductance_matrix,
    triangle_mesh_quadrature_points,
    triangle_mesh_self_force_mapping,
    flux_density_triangle_mesh,
    flux_density_triangle_mesh_hierarchical,
    flux_density_point_segment,
    gs_operator_order2,
    gs_operator_order4,
    HierarchicalDiagnostics,
    inductance_linear_filaments,
    inductance_piecewise_linear_filaments,
    mutual_inductance_circular_to_linear,
    rotate_filaments_about_path,
    vector_potential_circular_filament,
    vector_potential_linear_filament,
    vector_potential_linear_filament_hierarchical,
    vector_potential_triangle_mesh_mapping,
    vector_potential_triangle_mesh,
    vector_potential_triangle_mesh_hierarchical,
    vector_potential_point_segment,
    vector_potential_dipole,
    vector_potential_dipole_hierarchical,
    SolveResult,
)
from cfsem.flux_solver import (
    calc_flux_density_from_flux,
    flux_solver,
    gradient_order4,
    solve_flux_axisymmetric,
)
from cfsem.types import Array3xN

from .cfsem import ellipe, ellipk
from .cfsem import DimensionalityError

MU_0 = 4.0 * np.pi * 1e-7 * (1.0 + 5.5e-10)
"""
Vacuum permeability with slight correction from latest measurements per NIST CODATA 2019 (SP-961),
https://www.physics.nist.gov/cuu/pdf/wall_2018.pdf .
"""

__all__ = [
    "flux_circular_filament",
    "flux_density_linear_filament",
    "flux_density_linear_filament_hierarchical",
    "flux_density_triangle_mesh_mapping",
    "triangle_mesh_current_density",
    "triangle_mesh_force_mapping",
    "triangle_mesh_force_mapping_from_circular_filaments",
    "triangle_mesh_force_mapping_from_dipoles",
    "triangle_mesh_force_mapping_from_linear_filaments",
    "triangle_mesh_flux_linkage_mapping_from_dipoles",
    "triangle_mesh_inductance_mapping_from_circular_filaments",
    "triangle_mesh_inductance_mapping_from_linear_filaments",
    "triangle_mesh_inductance_matrix",
    "triangle_mesh_quadrature_points",
    "triangle_mesh_self_force_mapping",
    "flux_density_triangle_mesh",
    "flux_density_triangle_mesh_hierarchical",
    "flux_density_circular_filament",
    "gs_operator_order2",
    "gs_operator_order4",
    "gradient_order4",
    "calc_flux_density_from_flux",
    "flux_solver",
    "solve_flux_axisymmetric",
    "filament_helix_path",
    "inductance_linear_filaments",
    "inductance_matrix_axisymmetric_coaxial_rectangular_coils",
    "inductance_piecewise_linear_filaments",
    "self_inductance_piecewise_linear_filaments",
    "self_inductance_axisymmetric_coil",
    "mutual_inductance_piecewise_linear_filaments",
    "flux_density_ideal_solenoid",
    "self_inductance_lyle6",
    "mutual_inductance_of_circular_filaments",
    "mutual_inductance_of_cylindrical_coils",
    "filament_coil",
    "self_inductance_circular_ring_wien",
    "self_inductance_annular_ring",
    "self_inductance_distributed_axisymmetric_conductor",
    "ellipe",
    "ellipk",
    "rotate_filaments_about_path",
    "vector_potential_linear_filament",
    "vector_potential_linear_filament_hierarchical",
    "vector_potential_triangle_mesh_mapping",
    "vector_potential_triangle_mesh",
    "vector_potential_triangle_mesh_hierarchical",
    "vector_potential_circular_filament",
    "flux_density_circular_filament_cartesian",
    "mutual_inductance_circular_to_linear",
    "flux_density_dipole",
    "flux_density_dipole_hierarchical",
    "vector_potential_dipole",
    "vector_potential_dipole_hierarchical",
    "flux_density_point_segment",
    "vector_potential_point_segment",
    "HierarchicalDiagnostics",
    "SolveResult",
    "MU_0",
    "body_force_density_circular_filament_cartesian",
    "body_force_density_linear_filament",
    "DimensionalityError",
]


def self_inductance_piecewise_linear_filaments(
    xyzp: Array3xN,
    wire_radius: float | NDArray = 0.0,
) -> float:
    """
    Estimate the self-inductance of one piecewise-linear current filament.

    Uses the vector-potential line-integral form
    $L = \\oint \\vec{A} \\cdot d\\vec{l}$ with the existing finite-radius
    linear-filament vector-potential kernel.

    Assumes:

    * Thin, well-behaved filaments
    * Uniform current distribution within segments
        * Low frequency operation; no skin effect
    * Vacuum permeability everywhere
    * Each filament has a constant current in all segments
      (otherwise we need an interaction matrix)

    Args:
        xyzp: [m] 3xN point series describing the filament
        wire_radius: [m] filament radius, scalar or array of length `N-1`

    Returns:
        [H] Scalar self-inductance
    """
    x, y, z = xyzp
    xyzfil = (x[:-1], y[:-1], z[:-1])
    dlxyzfil = (x[1:] - x[:-1], y[1:] - y[:-1], z[1:] - z[:-1])

    self_inductance = inductance_piecewise_linear_filaments(
        xyzfil,
        dlxyzfil,
        xyzfil,
        dlxyzfil,
        wire_radius=wire_radius,
    )

    return self_inductance  # [H]


def mutual_inductance_piecewise_linear_filaments(
    xyz0: Array3xN,
    xyz1: Array3xN,
    wire_radius: float | NDArray = 0.0,
) -> float:
    """
    Estimate the mutual inductance between two piecewise-linear current filaments.

    Uses the vector-potential line-integral form
    $M = \\oint \\vec{A}_{source} \\cdot d\\vec{l}_{target}$ with the finite-radius
    linear-filament vector-potential kernel.

    Assumes:

    * Thin, well-behaved filaments
    * Vacuum permeability everywhere
    * Each filament has a constant current in all segments
      (otherwise we need an interaction matrix)
    * All segments between the two filaments are distinct; no identical pairs

    Args:
        xyz0: [m] 3xN point series describing the first filament
        xyz1: [m] 3xM point series describing the second filament
        wire_radius: [m] source filament radius for `xyz0`, scalar or array of length `N-1`

    Returns:
        [H] Scalar mutual inductance between the two filaments
    """
    x0, y0, z0 = xyz0
    xyzfil0 = (x0[:-1], y0[:-1], z0[:-1])
    dlxyzfil0 = (x0[1:] - x0[:-1], y0[1:] - y0[:-1], z0[1:] - z0[:-1])

    x1, y1, z1 = xyz1
    xyzfil1 = (x1[:-1], y1[:-1], z1[:-1])
    dlxyzfil1 = (x1[1:] - x1[:-1], y1[1:] - y1[:-1], z1[1:] - z1[:-1])

    inductance = inductance_piecewise_linear_filaments(
        xyzfil0,
        dlxyzfil0,
        xyzfil1,
        dlxyzfil1,
        wire_radius=wire_radius,
    )

    return inductance  # [H]


def flux_density_ideal_solenoid(current: float, num_turns: float, length: float) -> float:
    """
    Axial B-field on centerline of an ideal (infinitely long) solenoid.

    This calc converges reasonably well for coil L/D > 20.

    Args:
        current: [A] solenoid current
        num_turns: [#] number of conductor turns
        length: [m] length of winding pack

    Returns:
        [T] B-field on axis (in the direction aligned with the axis)
    """
    b_on_axis = MU_0 * num_turns * current / length
    return b_on_axis  # [T]


def self_inductance_lyle6(r: float, dr: float, dz: float, n: float) -> float:
    """
    Self-inductance of a cylindrically-symmetric coil of rectangular
    cross-section, estimated to 6th order.

    This estimate is viable up to an L/D of about 1.5, above which it
    rapidly accumulates error and eventually produces negative values.

    References:
        [1] T. R. Lyle,
        “IX. On the self-inductance of circular coils of rectangular section,”
        Philosophical Transactions of the Royal Society of London.
        Series A, Containing Papers of a Mathematical or Physical Character,
        vol. 213, no. 497-508, pp. 421-435, Jan. 1914, doi: [10.1098/rsta.1914.0009](https://doi.org/10.1098/rsta.1914.0009)

    Args:
        r: [m] radius, coil center
        dr: [m] radial width of coil
        dz: [m] cylindrical height of coil
        n: number of turns

    Returns:
        [H] self-inductance
    """

    assert dr < 2.0 * r, "Axisymmetric coil edges can't extend to negative R"
    assert dz / r <= 3.5, "Coil geometry outside domain of validity of Lyle's formula"

    # Guarantee 64-bit floats needed for 6th-order shape term
    a = np.float64(r)
    b = np.float64(dz)
    c = np.float64(dr)

    # Build up reusable terms for calculation of shape parameter
    d = np.sqrt(b**2 + c**2)  # [m] diagonal length
    u = ((b / c) ** 2) * 2 * np.log(d / b)
    v = ((c / b) ** 2) * 2 * np.log(d / c)
    w = (b / c) * np.arctan(c / b)
    ww = (c / b) * np.arctan(b / c)

    bd2 = (b / d) ** 2
    cd2 = (c / d) ** 2
    da2 = (d / a) ** 2
    ml = np.log(8 * a / d)

    f = (
        ml
        + (1 + u + v - 8 * (w + ww)) / 12.0  # 0th order in d/a
        + (da2 * (cd2 * (221 + 60 * ml - 6 * v) + 3 * bd2 * (69 + 60 * ml + 10 * u - 64 * w)))
        / 5760.0  # 2nd order
        + (
            da2**2
            * (
                2 * cd2**2 * (5721 + 3080 * ml - 345 * v)
                + 5 * bd2 * cd2 * (407 + 5880 * ml + 6720 * u - 14336 * w)
                - 10 * bd2**2 * (3659 + 2520 * ml + 805 * u - 6144 * w)
            )
        )
        / 2.58048e7  # 4th order
        + (
            da2**3
            * (
                3 * cd2**3 * (4308631 + 86520 * ml - 10052 * v)
                - 14 * bd2**2 * cd2 * (617423 + 289800 * ml + 579600 * u - 1474560 * w)
                + 21 * bd2**3 * (308779 + 63000 * ml + 43596 * u - 409600 * w)
                + 42 * bd2 * cd2**2 * (-8329 + 46200 * ml + 134400 * u - 172032 * w)
            )
        )
        / 1.73408256e10  # 6th order
    )  # [nondim] shape parameter

    self_inductance = MU_0 * (n**2) * a * f

    assert self_inductance >= 0.0, "Coil geometry outside domain of validity of Lyle's formula"

    return self_inductance  # [H]


def mutual_inductance_of_circular_filaments(rzn1: NDArray, rzn2: NDArray, par: bool = True) -> float:
    """
    Analytic mutual inductance between a pair of ideal cylindrically-symmetric coaxial filaments.

    This is equivalent to taking the flux produced by each circular filament from
    either collection of filaments to the other. Mutual inductance is reflexive,
    so the order of the inputs is not important.

    Args:
        rzn1 (array): 3x1 array (r [m], z [m], n []) coordinates and number of turns
        rzn2 (array): 3x1 array (r [m], z [m], n []) coordinates and number of turns
        par: Whether to use CPU parallelism

    Returns:
        float: [H] mutual inductance
    """

    m = mutual_inductance_of_cylindrical_coils(rzn1.reshape((3, 1)), rzn2.reshape((3, 1)), par)

    return m  # [H]


def inductance_matrix_axisymmetric_coaxial_rectangular_coils(
    r: list[float],
    z: list[float],
    dr: list[float],
    dz: list[float],
    td: list[float],
    nr: list[int],
    nz: list[int],
) -> NDArray:
    """Inductance matrix for a set of coils with rectangular cross-sections and prescribed turn
    density. Self-inductances are calculated using Lyle's 6th-order formula, mutual inductance
    calculations use filamentized coil model. The filamentization inputs nr and nt are used primarily
    to define the accuracy of the mutual inductance calculations, but do not impact assumed total current
    or current density in each coil.

    Args:
        r: [m] radii, coil center
        z: [m] axial positions, coil center
        dr: [m] radial sizes of each coil pack
        dz: [m] axial sizes of each coil pack
        td: [turns/m^2] turn density in each coil pack (same as current density at 1A per turn)
        nr: radial discretizations
        nz: axial discretizations

    Returns:
        m: Mutual inductance matrix in [H] with size (ncoil, ncoil)
    """

    # Check that all input lists have the same length
    assert (
        len(np.unique([len(r), len(z), len(dr), len(dz), len(td), len(nr), len(nz)])) == 1
    ), "All input lists must have the same length."

    # Check for sizes to be positive
    assert all(size > 0 for size in dr), "All coil radial sizes must be positive."
    assert all(size > 0 for size in dz), "All coil axial sizes must be positive."

    # Check for discretizations to be positive integers
    assert all(
        isinstance(n, int) and n > 0 for n in nr
    ), "All radial discretizations must be positive integers."
    assert all(
        isinstance(n, int) and n > 0 for n in nz
    ), "All axial discretizations must be positive integers."

    # Make sure that rectangular coil don't overlap
    nc = len(r)
    # Calculate turn number per coil
    nt = [td[c] * dr[c] * dz[c] for c in range(nc)]

    for c1 in range(nc):
        # Corner points of coil 1
        r1_min = r[c1] - dr[c1] / 2
        r1_max = r[c1] + dr[c1] / 2
        z1_min = z[c1] - dz[c1] / 2
        z1_max = z[c1] + dz[c1] / 2

        for c2 in range(c1 + 1, nc):
            # Corner points of coil 2
            r2_min = r[c2] - dr[c2] / 2
            r2_max = r[c2] + dr[c2] / 2
            z2_min = z[c2] - dz[c2] / 2
            z2_max = z[c2] + dz[c2] / 2

            # Check for radial overlap
            radial_overlap = (
                (r1_min > r2_min and r1_min < r2_max)
                or (r1_max > r2_min and r1_max < r2_max)
                or (r2_min > r1_min and r2_min < r1_max)
                or (r2_max > r1_min and r2_max < r1_max)
            )

            # No need to check for axial overlap if no radial overlap
            if not radial_overlap:
                continue

            # Check for axial overlap
            axial_overlap = (
                (z1_min > z2_min and z1_min < z2_max)
                or (z1_max > z2_min and z1_max < z2_max)
                or (z2_min > z1_min and z2_min < z1_max)
                or (z2_max > z1_min and z2_max < z1_max)
            )

            assert not axial_overlap, f"Rectangular coils {c1} and {c2} overlap."

    # Create filaments for each coil
    filaments_list = []
    for c in range(nc):
        filaments = filament_coil(r=r[c], z=z[c], w=dr[c], h=dz[c], nt=nt[c], nr=nr[c], nz=nz[c])
        filaments_list.append(filaments)

    # Build symmetric inductance matrix
    m = np.zeros((nc, nc), dtype=np.float64)
    for c1 in range(nc):
        for c2 in range(c1, nc):
            if c1 == c2:
                # Use Lyle's formula for self-inductance
                m[c1, c2] = self_inductance_lyle6(
                    r=r[c1],
                    dr=dr[c1],
                    dz=dz[c1],
                    n=nt[c1],  # Number of turns in the coil
                )
            else:
                # Mutual inductance between different filamentized coils
                m[c1, c2] = mutual_inductance_of_cylindrical_coils(
                    f1=filaments_list[c1].T,
                    f2=filaments_list[c2].T,
                )
                m[c2, c1] = m[c1, c2]  # Symmetric matrix

    return m  # [H]


def self_inductance_axisymmetric_coil(
    f: NDArray,
    section_kind: Literal["rectangular", "circular", "annular"],
    section_size: tuple[float, float] | float,
    par: bool = True,
) -> float:
    """
    Self-inductance of a coaxial collection of ideal circular filaments.

    Args:
        f: 3 x N array of filament definitions like (r [m], z [m], n []) coordinates and number of turns
        section_kind: Kind of conductor cross-section to use for self-inductance correction
        section_size: Size of the conductor cross-section
            * For "rectangular", a tuple (width [m], height [m])
            * For "circular", a single float (radius [m])
            * For "annular", a tuple (inner radius [m], outer radius [m])

    Returns:
        float: Self-inductance [H]
    """

    # Validate valid section_kind and section_size
    assert section_kind in (
        "rectangular",
        "circular",
        "annular",
    ), "section_kind must be one of 'rectangular', 'circular', or 'annular'"

    if section_kind == "rectangular":
        assert isinstance(
            section_size, tuple
        ), "For rectangular section, section_size must be a tuple (conductor width and height)"
        assert len(section_size) == 2, "For rectangular section, section_size must be a tuple of length 2"
        assert (
            section_size[0] > 0.0 and section_size[1] > 0.0
        ), "Conductor section sizes (width and height) must be positive"

    if section_kind == "circular":
        assert (
            isinstance(section_size, float) and section_size > 0.0
        ), "For circular section, section size must be a non-negative float (conductor radius)"

    if section_kind == "annular":
        assert (
            isinstance(section_size, tuple) and len(section_size) == 2
        ), "For annular section, section_size must be a tuple (inner and outer conductor radius)"
        assert section_size[0] >= 0.0 and section_size[1] > section_size[0], (
            "For annular section, "
            "inner radius must be non-negative and outer radius must be larger than inner radius"
        )

    m = 0.0
    rs, zs, ns = f

    # Iterate over filaments
    for i in range(f.shape[1]):
        # Get mutual inductance of i-th filament to all other filaments
        m_contribs = ns * flux_circular_filament(ns[i], rs[i], zs[i], rs, zs, par)

        # Get self-inductance of i-th filament
        if section_kind == "rectangular":
            assert isinstance(section_size, tuple)
            m_self = self_inductance_lyle6(
                r=rs[i],  # coil center radius
                dr=section_size[0],  # width of the rectangular conductor section
                dz=section_size[1],  # height of the rectangular conductor section
                n=ns[i],
            )
        elif section_kind == "circular":
            # Wien formula for circular cross-section
            assert isinstance(section_size, float)
            m_self = (
                self_inductance_circular_ring_wien(
                    major_radius=rs[i],  # coil center radius
                    minor_radius=np.atleast_1d(section_size),
                )
                * ns[i] ** 2
            )
            m_self = float(m_self.item())  # Convert to float

        elif section_kind == "annular":
            # Wien formula for annular cross-section
            assert isinstance(section_size, tuple)
            m_self = (
                self_inductance_annular_ring(
                    r=rs[i],  # major radius
                    a=section_size[0],  # inner minor radius (tube inside radius)
                    b=section_size[1],  # outer minor radius (tube outside radius)
                )
                * ns[i] ** 2
            )

        # Replace self-pairing with self-inductance
        m_contribs[i] = m_self
        m += np.sum(m_contribs)

    return m  # [H]


def mutual_inductance_of_cylindrical_coils(f1: NDArray, f2: NDArray, par: bool = True) -> float:
    """
    Analytical mutual inductance between two coaxial collections of ideal filaments.

    Each collection typically represents a discretized "real" cylindrically-symmetric
    coil of rectangular cross-section, but could have any cross-section as long as it
    maintains cylindrical symmetry.

    Args:
        f1: 3 x N array of filament definitions like (r [m], z [m], n [])
        f2: 3 x N array of filament definitions like (r [m], z [m], n [])
        par: Whether to use CPU parallelism

    Returns:
        [H] mutual inductance of the two discretized coils
    """
    r1, z1, n1 = f1
    r2, z2, n2 = f2

    # Using n2 as the current per filament is equivalent to examining a 1A reference current,
    # which gives us the flux per amp (inductance)
    m = np.sum(n1 * flux_circular_filament(n2, r2, z2, r1, z1, par))  # [H] total mutual inductance

    return m  # [H]


def filament_coil(r: float, z: float, w: float, h: float, nt: float, nr: int, nz: int) -> NDArray:
    """
    Create an array of filaments from coil cross-section, evenly spaced
    _inside_ the winding pack. No filaments are coincident with the coil surface.

    Args:
        r: [m] radius, coil center
        z: [m] axial position, coil center
        w: [m] width of coil pack
        h: [m] height of coil pack
        nt: turns
        nr: radial discretizations
        nz: axial discretizations

    Returns:
        (nr*nz) x 3, (r,z,n) of each filament
    """

    # Build a 2D mesh of points evenly spaced in the interior of the bounding rectangle
    rs = np.linspace(r - w * (nr - 1) / nr / 2, r + w * (nr - 1) / nr / 2, nr)  # [m]
    zs = np.linspace(z - h * (nz - 1) / nz / 2, z + h * (nz - 1) / nz / 2, nz)  # [m]
    rmesh, zmesh = np.meshgrid(rs, zs, indexing="ij")  # [m]

    # Number of turns attributed to each point is not necessarily an integer
    n = np.full_like(rmesh.ravel(), float(nt) / (nr * nz))

    # Pack filament locations and number of turns into an array together
    filaments = np.dstack([rmesh.ravel(), zmesh.ravel(), n]).reshape(nr * nz, 3)

    return filaments  # [m], [m], [dimensionless]


def self_inductance_circular_ring_wien(major_radius: NDArray, minor_radius: NDArray) -> NDArray:
    """
    Wien's formula for the self-inductance of a circular ring
    with thin circular cross section.

    Uses equation 7 from reference [1].

    References:
        [1] E Rosa and L Cohen, "On the Self-Inductance of Circles,"
        Bulletin of the Bureau of Standards, 1908.
        [Online]. Available:
        <https://nvlpubs.nist.gov/nistpubs/bulletin/04/nbsbulletinv4n1p149_A2b.pdf>

    Args:
        major_radius: [m] Major radius of the ring.
        minor_radius: [m] Radius of the ring's cross section.

    Returns:
        [H] self-inductance
    """
    ar = major_radius / minor_radius  # [], a / rho, dimensionless
    ra2 = (minor_radius / major_radius) ** 2  # [], (rho / a)^2, dimensionless
    # Equation 7 in Rosa & Cohen lacks the factor of 1e-7
    # because it is not in SI base units.
    self_inductance = MU_0 * major_radius * ((1 + 0.125 * ra2) * np.log(8 * ar) - 0.0083 * ra2 - 1.75)  # [H]
    return self_inductance  # [H]


def self_inductance_distributed_axisymmetric_conductor(
    current: float,
    grid: tuple[NDArray, NDArray],
    mesh: tuple[NDArray, NDArray],
    b_part: tuple[NDArray, NDArray],
    psi_part: NDArray,
    mask: NDArray,
    edge_path: tuple[NDArray, NDArray],
) -> tuple[float, float, float]:
    """
    Calculation of a distributed conductor's self-inductance from two components:

    * External inductance: the portion related to the poloidal flux exactly at the conductor edge
    * Internal inductance: the portion related to the poloidal magnetic field inside the conductor,
      where the filamentized method used for coils does not apply due to the parallel arrangement
      and variable current density.

    Note: the B-field and flux inputs are _not_ the total from all sources - they are only the
    contribution from the distributed conductor under examination.

    This calculation was developed for use with tokamak plasmas, but applies similarly to
    other kinds of distributed axisymmetric conductor.

    Assumptions:

    * Cylindrically-symmetric, distributed, single-winding, contiguous conductor
    * No high-magnetic-permeability materials in the vicinity
    * Isopotential on the edge contour
        * This is slightly less restrictive than isopotential on the section,
          but notably does _not_ allow the calc to be used with, for example,
          large, shell conductors where different regions are meaningfully
          independent of each other.
    * Conductor interior does not touch the edge of the computational domain
        * At least one grid cell of padding is needed to support finite differences

    References:
        [1] S. Ejima, R. W. Callis, J. L. Luxon, R. D. Stambaugh, T. S. Taylor, and J. C. Wesley,
            “Volt-second analysis and consumption in Doublet III plasmas,”
            Nucl. Fusion, vol. 22, no. 10, pp. 1313-1319, Oct. 1982,
            doi: [10.1088/0029-5515/22/10/006](https://doi.org/10.1088/0029-5515/22/10/006)

        [2] J. A. Romero and J.-E. Contributors, “Plasma internal inductance dynamics in a tokamak,”
            arXiv.org. Accessed: Dec. 21, 2023. [Online]. Available: https://arxiv.org/abs/1009.1984v1
            doi: [10.1088/0029-5515/50/11/115002](https://doi.org/10.1088/0029-5515/50/11/115002)

        [3] J. T. Wai and E. Kolemen, “GSPD: An algorithm for time-dependent tokamak equilibria design.”
            arXiv, Jun. 22, 2023. Accessed: Sep. 15, 2023. [Online]. Available: https://arxiv.org/abs/2306.13163
            doi: [10.48550/arXiv.2306.13163](https://doi.org/10.48550/arXiv.2306.13163)

    Args:
        current: [A] total toroidal current in this conductor
        grid: [m] (1 X nr) grids of (R coords, Z coords)
        mesh: [m] (nr X nz) meshgrids of (R coords, Z coords)
        b_part: [T] (nr X nz) Flux density (R-component, Z-component) due to this conductor
        psi_part: [V-s] or [T-m^2] (nr X nz) this conductor's poloidal flux field
        mask: (nr X nz) positive mask of the conductor's interior region
        edge_path: [m] (2 x N) closed (r, z) path along conductor edge

    Returns:
        (Lt, Li, Le) Total, internal, and external self-inductance components
    """
    # Unpack
    rgrid, zgrid = grid  # [m]
    rmesh, zmesh = mesh  # [m]
    br, bz = b_part  # [T]

    # Set up
    nr = rgrid.size
    nz = zgrid.size
    psi_interpolator = MulticubicRectilinear.new([rgrid, zgrid], psi_part.ravel())

    # Same-length diffs assuming last grid cell is the same size
    # as the previous one. In general the conductor can't touch the
    # edge of the computational domain without breaking the solver,
    # so it's ok to allow some potential error at the last index.
    drmesh = np.zeros_like(rmesh)  # [m]
    drmesh[: nr - 1, :] = np.diff(rmesh, axis=0)
    drmesh[-1, :] = drmesh[-2, :]

    dzmesh = np.zeros_like(zmesh)  # [m]
    dzmesh[:, : nz - 1] = np.diff(zmesh, axis=1)
    dzmesh[:, -1] = dzmesh[:, -2]

    # Toroidal volume of each cell in the mesh.
    #
    # Ideally we'd use a two-sided difference for
    # calculating dr and dz in order to properly handle cell volume
    # for nonuniform grids, but (1) that's a lot of extra array handling
    # for minimal real benefit and (2) the grid is almost always uniform
    volmesh = 2.0 * np.pi * rmesh * drmesh * dzmesh  # [m^3]

    # Stored poloidal magnetic energy inside the conductor volume.
    #
    # Because E ~ B^2 and B = sqrt(Br^2 + Bz^2 + Btor^2) -> B^2 = Br^2 + Bz^2 + Btor^2,
    # the components of magnetic energy (and induction) due to a current on different axes
    # are separable, and we can ignore the toroidal field entirely when treating the
    # poloidal inductance.
    #
    # That doesn't mean there isn't store energy related to the conductor's toroidal field,
    # only that it can be separated from the poloidal inductance.
    wmag_pol = (1.0 / (2.0 * MU_0)) * np.sum((br**2 + bz**2) * volmesh * mask)  # [J]

    # Internal inductance.
    #
    # Because E = (1/2) * L * I^2, we can superpose multiple inductance terms
    # for a given current-carrying element, and say E = (1/2) * (Li + Le) * I^2 .
    #
    # Now that we have the internal stored magnetic energy, we can calculate
    # the part of the self-inductance related to that energy directly.
    internal_inductance = 2.0 * wmag_pol / current**2  # [H]

    # External inductance.
    #
    # This is a hack based on Poynting's Theorem relating surface conditions to
    # magnetic energy.
    #
    #   Do weighted average of flux along the edge contour
    #   to adjust for the length of individual segments
    edge_dr = np.diff(edge_path[0])  # [m]
    edge_dz = np.diff(edge_path[1])  # [m]
    edge_dl = (edge_dr**2 + edge_dz**2) ** 0.5  # [m] length of each segment
    edge_length = np.sum(edge_dl)  # [m] total length of conductor limit contour
    edge_psi = psi_interpolator.eval([x[:-1] for x in edge_path])  # [V-s]
    edge_psi_mean = np.sum(edge_psi * edge_dl) / edge_length  # [V-s]
    #   Take the inductance
    external_inductance = edge_psi_mean / current  # [H]

    # Total inductance.
    total_inductance = internal_inductance + external_inductance  # [H]

    return (
        float(total_inductance),
        float(internal_inductance),
        float(external_inductance),
    )


def self_inductance_annular_ring(r: float, a: float, b: float) -> float:
    """
    Low-frequency self-inductance of a thick-walled tube bent in a circle.

    Uses Wien's method, per the 1912 NIST handbook [1], Eqn. 64 on pg. 112,
    with a correction to a misprint in term 5 and a unit conversion factor
    in the final expression.

    This is an approximation that drops terms of order higher than `(a/r)^2`
    and `(b/r)^2`.

    References:
        [1] E. B. Rosa and F. W. Grover,
            “Formulas and tables for the calculation of mutual and self-inductance (Revised),”
            BULL. NATL. BUR. STAND., vol. 8, no. 1, p. 1, Jan. 1912,
            doi: [10.6028/bulletin.185](https://doi.org/10.6028/bulletin.185)

    Args:
        r: [m] major radius of loop
        a: [m] inner minor radius (tube inside radius)
        b: [m] outer minor radius (tube outside radius)

    Returns:
        [H] self-inductance [H]
    """

    if not ((r > 0.0) and (a > 0.0) and (b > 0.0)):
        raise ValueError("All radii must be positive and nonzero.")
    if not (a < b):
        raise ValueError("Inner minor radius must be less than outer minor radius.")
    if not (b < r):
        raise ValueError("Outer minor radius must be less than major radius.")

    term1 = (1.0 + (a**2 + b**2) / (8.0 * r**2)) * np.log(8.0 * r / b)
    term2 = -1.75 + (2.0 * b**2 + a**2) / (32.0 * r**2)
    term3 = -0.5 * a**2 / (b**2 - a**2)
    term4 = (a**4 / ((b**2 - a**2) ** 2)) * (1.0 + a**2 / (8.0 * r**2)) * np.log(b / a)
    term5 = -(a**4 + a**2 * b**2 + b**4) / (48.0 * r**2 * (b**2 - a**2))

    self_inductance = MU_0 * r * (term1 + term2 + term3 + term4 + term5)  # [H]

    return self_inductance  # [H]
