"""Straight line, 3D distance.

The distance between two XYZ coordinates is straightforward, just the magnitude of
p2-p1, but in order to keep from retyping it in many places, we'll keep a public
method here as well.

This method is unitless, as long as both units match, it will return whatever units you
used as inputs.

Supports vectorized operations: pass arrays of shape (N, 3) for pairwise distances.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def dist_xyz(
    xyz1: NDArray[np.floating] | Sequence,
    xyz2: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating] | float:
    """Get the 3D, straight line distance.

    Returns the magnitude of (xyz2-xyz1) in the same units as were passed in.

    Supports vectorized operations for pairwise distances between corresponding points.

    Args:
        xyz1: XYZ coordinate(s). Shape (3,) for single point or (N, 3) for N points.
        xyz2: XYZ coordinate(s) in the same units as xyz1. Must match shape of xyz1.

    Returns:
        Distance(s) between the points. Returns a float for single point pairs,
        or ndarray of shape (N,) for vectorized input.
    """
    arr1 = np.asarray(xyz1)
    arr2 = np.asarray(xyz2)
    result = np.linalg.norm(arr2 - arr1, axis=-1)
    if result.ndim == 0:
        return float(result)
    return result
