"""Mathematical utility functions for solvers."""

import numpy as np


def cross_product(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    r"""
    Compute cross product :math:`a \times b`.

    This implementation is faster than np.cross for large arrays.

    Args:
        a: First vector (shape (3, nx, ny, nz))
        b: Second vector (shape (3, nx, ny, nz))

    Returns:
        Cross product :math:`a \times b` (shape (3, nx, ny, nz))
    """
    return np.stack(
        [
            a[1] * b[2] - a[2] * b[1],  # x-component
            a[2] * b[0] - a[0] * b[2],  # y-component
            a[0] * b[1] - a[1] * b[0],  # z-component
        ],
        axis=0,
    )


def normalize(m_n: np.ndarray):
    r"""
    Normalize the magnetization array (in place).

    .. math::

        \mathbf{m}_n = \frac{\mathbf{m}_n}{|\mathbf{m}_n|}

    Args:
        m_n: Magnetization array at time step n (shape (3, nx, ny, nz)).
    """
    norm = np.sqrt(m_n[0] ** 2 + m_n[1] ** 2 + m_n[2] ** 2)
    m_n /= norm
