Utils Module

def cart2pol(x, y):
    """Converts from cartesian to polar coordinates

    Args:
        x (ndarray): x coordinates
        y (ndarray): y coordinates

    Returns:
        tuple: rho, phi
    """
    rho = _np.sqrt(x ** 2 + y ** 2)
    phi = _np.arctan2(y, x)
    return (rho, phi)

def cart2sph(x, y, z):
    """Converts from cartesian to spherical coordinates

    Args:
        x (ndarray): x coordinates
        y (ndarray): y coordinates
        z (ndarray): z coordinates

    Returns:
        tuple: r, theta, phi
    """
    r = _np.sqrt(x ** 2 + y ** 2 + z ** 2)
    phi = _np.arctan2(y, x)

    # Hide the warning for situtations where there is a divide by zero.
    # This returns a NaN in the array, which is ignored for plotting.
    with _np.errstate(divide="ignore", invalid="ignore"):
        theta = _np.arccos(z / r)
    return (r, theta, phi)

def get_unit_value_meter(unit):
    """Returns a queried metre unit as a number
    Example:
        factor = get_unit_value_meter('cm')
        print(f"factor for 'cm' is {factor}")

    Args:
        unit (string): SI length unit

    Returns:
        float: SI prefix factor
    """
    si_prefixes = {
        "Ym": 1e24,
        "Zm": 1e21,
        "Em": 1e18,
        "Pm": 1e15,
        "Tm": 1e12,
        "Gm": 1e9,
        "Mm": 1e6,
        "km": 1e3,
        "hm": 1e2,
        "dam": 1e1,
        "m": 1,
        "dm": 1e-1,
        "cm": 1e-2,
        "mm": 1e-3,
        "µm": 1e-6,
        "um": 1e-6,
        "nm": 1e-9,
        "Ang": 1e-10,
        "pm": 1e-12,
        "fm": 1e-15,
        "am": 1e-18,
        "zm": 1e-21,
        "ym": 1e-24,
    }

    return si_prefixes.get(unit, None)

def get_unit_value_tesla(unit):
    """Returns a queried magnetic flux density unit as a number
    Example:
        factor = get_unit_value_meter('mT')
        print(f"factor for 'mT' is {factor}")

    Args:
        unit (string): SI length unit

    Returns:
        float: SI prefix factor
    """

    si_prefixes = {
        "YT": 1e24,
        "ZT": 1e21,
        "ET": 1e18,
        "PT": 1e15,
        "TT": 1e12,
        "GT": 1e9,
        "MT": 1e6,
        "kT": 1e3,
        "hT": 1e2,
        "daT": 1e1,
        "T": 1,
        "dT": 1e-1,
        "cT": 1e-2,
        "mT": 1e-3,
        "µT": 1e-6,
        "uT": 1e-6,
        "nT": 1e-9,
        "pT": 1e-12,
        "fT": 1e-15,
        "aT": 1e-18,
        "zT": 1e-21,
        "yT": 1e-24,
    }

    return si_prefixes.get(unit, None)

def pol2cart(rho, phi):
    """Converts from polar to cartesian coordinates

    Args:
        rho (ndarray): radial coordinates
        phi (ndarray): azimuthal coordinates

    Returns:
        tuple: x,y
    """
    x = rho * _np.cos(phi)
    y = rho * _np.sin(phi)
    return (x, y)

def sph2cart(r, theta, phi):
    """Converts from spherical to cartesian coordinates

    Args:
        r (ndarray): radial coordinates
        theta (ndarray): azimuthal angles
        phi (ndarray): polar angle

    Returns:
        tuple: x,y,z
    """
    x = r * _np.sin(theta) * _np.cos(phi)
    y = r * _np.sin(theta) * _np.sin(phi)
    z = r * _np.cos(theta)
    return x, y, z

def sphere_sph2cart(Br, Btheta, theta, phi):
    """Converts magnetic field of a sphere from spherical to cartesian coordinates

    Args:
        Br (ndarray): radial vector component
        Btheta (ndarray): polar vector component
        theta (ndarray): azimuthal angles
        phi (ndarray): polar angle

    Returns:
        tuple: Bx,By,Bz
    """
    Bx = Br * _np.sin(theta) * _np.cos(phi) + Btheta * _np.cos(theta) * _np.cos(phi)

    By = Br * _np.sin(theta) * _np.sin(phi) + Btheta * _np.cos(theta) * _np.sin(phi)

    Bz = Br * _np.cos(theta) - Btheta * _np.sin(theta)
    return Bx, By, Bz

def vector_pol2cart(Brho, Bphi, phi):
    """Converts Vectors from polar to cartesian coordinates

    Args:
        Brho (ndarray): radial vector component
        Bphi (ndarray): azimuthal vector component
        phi (ndarray): azimuthal coordinates

    Returns:
        tuple: Bx, By
    """
    Bx = Brho * _np.cos(phi) - Bphi * _np.sin(phi)
    By = Brho * _np.sin(phi) + Bphi * _np.cos(phi)
    return Bx, By

def vector_sph2cart(Br, Btheta, Bphi, theta, phi):
    """Converts Vectors from spherical to cartesian coordinates

    Args:
        Br (ndarray): radial vector component
        Btheta (ndarray): polar vector component
        Bphi (ndarray): azimuthal vector component
        theta (ndarray): azimuthal angles
        phi (ndarray): polar angle

    Returns:
        tuple: Bx,By,Bz
    """
    Bx = (
        Br * _np.sin(theta) * _np.cos(phi)
        + Btheta * _np.cos(theta) * _np.cos(phi)
        - Bphi * _np.sin(phi)
    )

    By = (
        Br * _np.sin(theta) * _np.sin(phi)
        + Btheta * _np.cos(theta) * _np.sin(phi)
        + Bphi * _np.cos(phi)
    )

    Bz = Br * _np.cos(theta) - Btheta * _np.sin(theta)
    return Bx, By, Bz

def FgradB_2D(B, x, y, chi_m, c):
    """Calculates the magnetic field gradient force for a 2D field.

    Args:
        B (Field2): Magnetic field vector
        x (ndarray): x coordinates
        y (ndarray): y coordinates

    Returns:
        Field2: Magnetic field gradient force vector
    """

    BgB = Field2(_np.zeros_like(B.n), _np.zeros_like(B.n))
    FB = Field2(_np.zeros_like(B.n), _np.zeros_like(B.n))

    dB = gradB_2D(B, x, y)
    BgB.n = dB.n * B.n
    BgB.x = dB.x * B.n
    BgB.y = dB.y * B.n
    FB.n = (1 / MU0) * chi_m * c * BgB.n
    FB.x = (1 / MU0) * chi_m * c * BgB.x
    FB.y = (1 / MU0) * chi_m * c * BgB.y
    return FB

def get_field_2D(Point_Array2):
    """Calculates magnetic field at an array of points due to every instantated
    `Magnet2D` magnet.

    Args:
        Point_Array2 (Point_Array2): array of x,y points and associated unit, defaults to 'mm'

    Returns:
        Field2: array of Bx,By,|B| values and associated unit (defaults to 'T')
    """
    from ..magnets import Magnet2D

    # Empty data structure
    B = _allocate_field_array2(Point_Array2.x, Point_Array2.y)

    for magnet in Magnet2D.instances:
        Bx, By = magnet.get_field(Point_Array2.x, Point_Array2.y)
        B.x += Bx
        B.y += By

    B.calc_norm()
    return B

def gradB_2D(B, x, y):
    """Calculates the magnetic field gradient for a 2D field.

    Args:
        B (Field2): Magnetic field vector
        x (ndarray): x coordinates
        y (ndarray): y coordinates

    Returns:
        Field2: Magnetic field gradient vector
    """

    dB = Field2(_np.zeros_like(B), _np.zeros_like(B))
    Nx = x.shape[0]
    Ny = x.shape[1]
    dx = (x.max() - x.min()) / Nx
    dy = (y.max() - y.min()) / Ny
    dB.x, dB.y = _np.gradient(B, dx, dy)
    dB.calc_norm()
    return dB

def grid2D(xmax, ymax, **kwargs):
    """Generates grid of x and y points

    Args:
        xmax (float): maximum x value
        ymax (float): maximum y value

    Kwargs:
        num_points (int): Number of points in each direction. Defaults to 100
        xmin (float): minimum x value. Defaults to -xmax
        ymin (float): minimum y value. Defaults to -ymax
        unit (str): unit length. Defaults to 'mm'

    Returns:
        Point_Array2: array of x and y values of shape (num_points, num_points) and associated unit
    """
    num_points = kwargs.pop("num_points", 100)
    xmin = kwargs.pop("xmin", -1 * xmax)
    ymin = kwargs.pop("ymin", -1 * ymax)
    unit = kwargs.pop("unit", "mm")
    NPJ = num_points * 1j
    x, y = _np.mgrid[xmin:xmax:NPJ, ymin:ymax:NPJ]
    return Point_Array2(x, y, unit=unit)

def rotate_points_2D(x, y, alpha):
    """Counter-clockwise rotation of points x,y

    Rotates 2D coordinates using a rotation matrix

    Args:
        x (ndarray): array of x coordinates
        y (ndarray): array of x coordinates
        alpha (float): rotation angle w.r.t. x-axis

    Returns:
        tuple: (x', y') rotated array of points
    """
    x = _np.atleast_1d(x)
    y = _np.atleast_1d(y)
    if len(x) != len(y):
        raise Exception("Must have same number of points in x and y")

    rot_matrix = _np.array(
        [[_np.cos(alpha), -_np.sin(alpha)], [_np.sin(alpha), _np.cos(alpha)]]
    )
    stacked_points = _np.column_stack((_np.ravel(x), _np.ravel(y)))
    rotated_points = _np.dot(rot_matrix, stacked_points.T)
    x_rotated = rotated_points[0, :]
    y_rotated = rotated_points[1, :]

    return _np.reshape(x_rotated, x.shape), _np.reshape(y_rotated, y.shape)

def get_field_3D(points):
    """Calculates magnetic field at an array of points due to every instantated
    `Magnet3D` magnet.

    Args:
        Point_Array3 (Point_Array3): array of x,y,z points and associated unit, defaults to 'mm'

    Returns:
        Field3: array of Bx,By,Bz,|B| values and associated unit (defaults to 'T')
    """
    from ..magnets import Magnet3D

    B = _allocate_field_array3(points.x, points.y, points.z)

    for magnet in Magnet3D.instances:
        Bx, By, Bz = magnet.get_field(points.x, points.y, points.z)
        B.x += Bx.reshape(B.x.shape)
        B.y += By.reshape(B.y.shape)
        B.z += Bz.reshape(B.z.shape)

    B.calc_norm()
    return B

def grid3D(xmax, ymax, zmax, **kwargs):
    """Generates grid of x, y, z points

    Args:
        xmax (float): maximum x value
        ymax (float): maximum y value
        zmax (float): maximum y value

    Kwargs:
        num_points (int): Number of points in each direction. Defaults to 100
        xmin (float): minimum x value. Defaults to -xmax
        ymin (float): minimum y value. Defaults to -ymax
        zmin (float): minimum y value. Defaults to -zmax
        unit (string): unit length. Defaults to 'mm'

    Returns:
        Point_Array2: array of x and y values of shape (num_points, num_points) and associated unit
    """
    num_points = kwargs.pop("num_points", None)

    xmin = kwargs.pop("xmin", -1 * xmax)
    ymin = kwargs.pop("ymin", -1 * ymax)
    zmin = kwargs.pop("zmin", -1 * zmax)
    unit = kwargs.pop("unit", "mm")
    # NPJ = num_points * 1j

    if num_points is None:
        num_points_x = kwargs.pop("num_points_x", 100)
        num_points_y = kwargs.pop("num_points_y", 100)
        num_points_z = kwargs.pop("num_points_z", 100)
    else:
        num_points_x = num_points
        num_points_y = num_points
        num_points_z = num_points

    x, y, z = _np.mgrid[
        xmin : xmax : num_points_x * 1j,
        ymin : ymax : num_points_y * 1j,
        zmin : zmax : num_points_z * 1j,
    ]

    return Point_Array3(x, y, z, unit=unit)

def line3D(start, end, num_points=100, **kwargs):
    """Generates a line of points

    Args:
        start (tuple): Starting point (x1,y1,z1)
        end (tuple): End point (x2,y2,z2)
        num_points (int): number of points to generate. Defaults to 100
        unit (str, optional): length scale units. Defaults to "mm".

    Returns:
        Point_Array3: array of x, y, and z values of shape (num_points) and associated unit
    """

    unit = kwargs.pop("unit", "mm")

    return Point_Array3(
        x=_np.linspace(start[0], end[0], num_points),
        y=_np.linspace(start[1], end[1], num_points),
        z=_np.linspace(start[2], end[2], num_points),
        unit=unit,
    )

def plane3D(origin, point_a, point_b, num_points=100):
    """Generates a plane of points defined by

    Args:
        origin (tuple): Origin of plane
        point_a (tuple): Point defining end of vector in x-direction
        point_b (tuple): Point defining end of vector in y-direction
        num_points (int, optional): Number of points in each direction, for a total of num_points^2. Defaults to 100.

    Returns:
        Point_Array3: Struct containing the x,y,z coordinates
    """

    vector_a = point_a - origin
    vector_b = point_b - origin
    num_points_j = num_points * 1j
    x_elements, y_elements = _np.mgrid[0:1:num_points_j, 0:1:num_points_j]
    x = origin[0] + x_elements * vector_a[0] + y_elements * vector_b[0]
    y = origin[1] + x_elements * vector_a[1] + y_elements * vector_b[1]
    z = origin[2] + x_elements * vector_a[2] + y_elements * vector_b[2]

    return Point_Array3(x=x, y=y, z=z)

def point3D(point, **kwargs):
    """Returns a single point

    Args:
        point (tuple):
        unit (str, optional): length scale units. Defaults to "mm".

    Returns:
        Point_Array3: struct of x, y, and z values of shape (1) and associated unit
    """

    unit = kwargs.pop("unit", "mm")

    return Point_Array3(
        x=point[0],
        y=point[1],
        z=point[2],
        unit=unit,
    )

def slice3D(plane="xy", max1=1.0, max2=1.0, slice_value=0.0, unit="mm", **kwargs):
    """Generates a planar slice of values

    Args:
        plane (str, optional): plane. Defaults to "xy".
        max1 (float, optional): maximum along axis 1. Defaults to 1.0.
        max2 (float, optional): maximum along axis 2. Defaults to 1.0.
        slice_value (float, optional): constant value for third axis. Defaults to 0.0.
        unit (str, optional): length scale units. Defaults to "mm".

    Kwargs:
        num_points (int): Number of points in each direction. Defaults to 100
        min1 (float): minimum along axis 1. Defaults to -min1
        min2 (float): minimum along axis 2. Defaults to -min2

    Raises:
        Exception: plane type, must be one of 'xy', 'xz, 'yz', or 'custom'

    Returns:
        Point_Array3: array of x, y, and z values of shape (num_points, num_points) and associated unit
    """
    num_points = kwargs.pop("num_points", 100)
    min1 = kwargs.pop("min1", -1 * max1)
    min2 = kwargs.pop("min2", -1 * max2)
    NPj = num_points * 1j

    if plane.lower() == "xy":
        x, y = _np.mgrid[min1:max1:NPj, min2:max2:NPj]
        z = _np.asarray([slice_value])
        z = _np.tile(z, x.shape)

    elif plane.lower() == "xz":
        x, z = _np.mgrid[min1:max1:NPj, min2:max2:NPj]
        y = _np.asarray([slice_value])
        y = _np.tile(y, x.shape)

    elif plane.lower() == "yz":
        y, z = _np.mgrid[min1:max1:NPj, min2:max2:NPj]
        x = _np.asarray([slice_value])
        x = _np.tile(x, y.shape)

    elif plane.lower() == "custom":
        x = kwargs.pop("custom_x", _np.array([0.0]))
        y = kwargs.pop("custom_y", _np.array([0.0]))
        z = kwargs.pop("custom_z", _np.array([0.0]))

    else:
        raise Exception("plane must be one of 'xy', 'xz, 'yz', or 'custom'")

    return Point_Array3(x, y, z, unit=unit)

def align_triangle_to_y(triangle, rot_axis, norm_vec):
    """Rotates and translates a triangle in lie in the xz plane

    Args:
        triangle (ndarray): vertices of a triangle
        rot_axis (ndarray): axis about which to rotate triangle
        norm_vec (ndarray): normal to triangle

    Returns:
        tuple: aligned_triangle (ndarray, rotated triangle), first_rotation (quaternion, align to y-axis)
    """
    y_axis = _np.array([0, 1, 0])

    if _np.linalg.norm(rot_axis) < ALIGN_CUTOFF:

        # Check if parallel or anti-parallel
        if check_sign(y_axis, norm_vec):
            # Parallel
            first_rotation = Quaternion()
            aligned_triangle = triangle

        else:
            # Anti-parallel
            first_rotation = Quaternion.q_angle_from_axis(PI, y_axis)
            aligned_triangle = rotate_points(triangle, first_rotation)

    else:
        angle = -_np.arccos(_np.dot(y_axis, norm_vec))
        first_rotation = Quaternion.q_angle_from_axis(angle, rot_axis)
        aligned_triangle = rotate_points(triangle, first_rotation)
    return aligned_triangle, first_rotation

def align_triangle_xz(triangle, longest_side):
    """Returns quaternions needed to rotate a triangle lying in the xz plane to be
    aligned with its longest side along x and altitude along z

    Args:
        triangle (ndarray): vertices of triangle
        longest_side (int): index of the longest side of triangle

    Returns:
        tuple: second_rotation (quaternion, align to x-axis), third_rotation (quaternion, align to z-axis):
    """
    x_axis = _np.array([1, 0, 0])
    y_axis = _np.array([0, 1, 0])
    z_axis = _np.array([0, 0, 1])

    # side_list = [0, 1, 2]
    # side_list.pop(longest_side)

    vec_x = return_axis_vector(triangle, longest_side)
    rot_axis = _np.cross(x_axis, vec_x)

    # Check aligment of base of triangle with x-axis
    if _np.linalg.norm(rot_axis) < ALIGN_CUTOFF:

        # Check if parallel or anti-parallel
        if check_sign(x_axis, vec_x):
            # Parallel
            second_rotation = Quaternion()
            tri_x = triangle
        else:
            # Anti-parallel
            second_rotation = Quaternion.q_angle_from_axis(PI, y_axis)
            tri_x = rotate_points(triangle, second_rotation)

    else:
        angle = -_np.arccos(_np.dot(x_axis, vec_x))
        second_rotation = Quaternion.q_angle_from_axis(angle, rot_axis)
        tri_x = rotate_points(triangle, second_rotation)

    vec_z = return_z_vector(tri_x, longest_side)
    rot_axis = _np.cross(z_axis, vec_z)

    # Check aligment of triangle altitude with z-axis
    if _np.all(_np.fabs([rot_axis]) < ALIGN_CUTOFF):

        # Check if parallel anti-parallel
        if check_sign(z_axis, vec_z):
            # Parallel
            third_rotation = Quaternion()
        else:
            # Anti-parallel
            third_rotation = Quaternion.q_angle_from_axis(PI, y_axis)

    else:
        angle = -_np.arccos(_np.dot(z_axis, vec_z))
        third_rotation = Quaternion.q_angle_from_axis(angle, rot_axis)

    return second_rotation, third_rotation

def altitude(a, b, c):
    """Gets altitude to side `a` of a triangle.

    Args:
        a (float): longest side
        b (float): triangle side
        c (float): triangle side

    Returns:
        float: altitude to side `a`
    """
    s = (a + b + c) / 2
    return 2 * _np.sqrt(s * (s - a) * (s - b) * (s - c)) / a