sphereql-core 0.3.0

use crate::conversions::spherical_to_cartesian;
use crate::error::SphereQlError;
use crate::types::{CartesianPoint, SphericalPoint};

/// Returns the angular separation (in radians) between two spherical points.
///
/// Uses the Vincenty formula for numerical stability at all separations.
///
/// ```
/// use sphereql_core::{SphericalPoint, angular_distance};
///
/// let a = SphericalPoint::new_unchecked(1.0, 0.0, 0.0);
/// let b = SphericalPoint::new_unchecked(1.0, 0.0, std::f64::consts::FRAC_PI_2);
/// let dist = angular_distance(&a, &b);
/// assert!((dist - std::f64::consts::FRAC_PI_2).abs() < 1e-10);
/// ```
#[must_use]
pub fn angular_distance(a: &SphericalPoint, b: &SphericalPoint) -> f64 {
    // Hottest call in the workspace. `unit_cartesian` is #[inline] and avoids
    // allocating SphericalPoint/CartesianPoint temporaries that the old
    // spherical_to_cartesian path produced per call.
    let [ax, ay, az] = a.unit_cartesian();
    let [bx, by, bz] = b.unit_cartesian();

    // Vincenty formula: numerically stable for all angular separations
    let cross_x = ay * bz - az * by;
    let cross_y = az * bx - ax * bz;
    let cross_z = ax * by - ay * bx;
    let cross_mag = (cross_x * cross_x + cross_y * cross_y + cross_z * cross_z).sqrt();

    let dot = ax * bx + ay * by + az * bz;

    cross_mag.atan2(dot)
}

/// Cosine proxy distance between two unit Cartesian direction vectors.
///
/// Returns `1 - dot(a, b)`, which is in [0, 2] and is monotone with angular
/// distance: 0 when identical, 1 at 90°, 2 when antipodal. This is much
/// cheaper than [`angular_distance`] since it avoids spherical-to-Cartesian
/// conversion, cross product, sqrt, and atan2.
///
/// Use this when only the **ordering** matters (k-NN heaps, candidate pruning).
/// Convert only the final k results to actual angular distance.
///
/// ```
/// use sphereql_core::cosine_proxy;
///
/// let a = [1.0, 0.0, 0.0];
/// let b = [1.0, 0.0, 0.0];
/// assert!(cosine_proxy(&a, &b) < 1e-10);
///
/// let c = [-1.0, 0.0, 0.0];
/// assert!((cosine_proxy(&a, &c) - 2.0).abs() < 1e-10);
/// ```
#[must_use]
#[inline]
pub fn cosine_proxy(a: &[f64; 3], b: &[f64; 3]) -> f64 {
    let dot = a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
    1.0 - dot.clamp(-1.0, 1.0)
}

/// Returns the great-circle (arc) distance between two points on a sphere of given `radius`.
///
/// ```
/// use sphereql_core::{SphericalPoint, great_circle_distance};
/// use std::f64::consts::FRAC_PI_2;
///
/// let a = SphericalPoint::new_unchecked(1.0, 0.0, 0.0);
/// let b = SphericalPoint::new_unchecked(1.0, 0.0, FRAC_PI_2);
/// let dist = great_circle_distance(&a, &b, 6371.0);
/// assert!((dist - 6371.0 * FRAC_PI_2).abs() < 1e-6);
/// ```
#[must_use]
pub fn great_circle_distance(a: &SphericalPoint, b: &SphericalPoint, radius: f64) -> f64 {
    radius * angular_distance(a, b)
}

/// Returns the straight-line (chord) distance between two spherical points.
///
/// Unlike angular or great-circle distances which measure along the sphere
/// surface, this computes the Euclidean distance through 3D space.
///
/// ```
/// use sphereql_core::{SphericalPoint, chord_distance};
///
/// let p = SphericalPoint::new_unchecked(1.0, 0.0, 0.5);
/// assert!(chord_distance(&p, &p) < 1e-10);
/// ```
#[must_use]
pub fn chord_distance(a: &SphericalPoint, b: &SphericalPoint) -> f64 {
    let ac = spherical_to_cartesian(a);
    let bc = spherical_to_cartesian(b);
    euclidean_distance(&ac, &bc)
}

/// Returns the Euclidean (L2) distance between two Cartesian points.
#[must_use]
pub fn euclidean_distance(a: &CartesianPoint, b: &CartesianPoint) -> f64 {
    let dx = a.x - b.x;
    let dy = a.y - b.y;
    let dz = a.z - b.z;
    (dx * dx + dy * dy + dz * dz).sqrt()
}

/// Cosine similarity between two high-dimensional vectors.
///
/// Returns `dot(a, b) / (‖a‖ × ‖b‖)`, in [-1, 1]. Returns 0.0 if either
/// vector has zero norm. This operates on the **original** embedding space,
/// not the projected sphere.
///
/// # Errors
///
/// Returns [`SphereQlError::DimensionMismatch`] when the input slices have
/// different lengths. Use this fallible form at any boundary that ingests
/// vectors from an external source (remote vector stores, user payloads).
///
/// ```
/// use sphereql_core::cosine_similarity;
///
/// let a = vec![1.0, 0.0, 0.0];
/// let b = vec![1.0, 0.0, 0.0];
/// assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 1e-10);
/// ```
pub fn cosine_similarity(a: &[f64], b: &[f64]) -> Result<f64, SphereQlError> {
    if a.len() != b.len() {
        return Err(SphereQlError::DimensionMismatch {
            expected: a.len(),
            actual: b.len(),
        });
    }
    let (mut dot, mut norm_a, mut norm_b) = (0.0, 0.0, 0.0);
    for (&x, &y) in a.iter().zip(b.iter()) {
        dot += x * y;
        norm_a += x * x;
        norm_b += y * y;
    }
    let denom = norm_a.sqrt() * norm_b.sqrt();
    if denom < f64::EPSILON {
        return Ok(0.0);
    }
    Ok((dot / denom).clamp(-1.0, 1.0))
}

/// Compute pairwise cosine similarities for a set of vectors.
///
/// Returns a flat `Vec<f64>` containing the upper triangle of the n×n
/// similarity matrix in row-major order: `[sim(0,1), sim(0,2), ..., sim(0,n-1), sim(1,2), ...]`.
/// Length is `n * (n - 1) / 2`.
///
/// All vectors must have the same dimensionality. Returns
/// `Err(DimensionMismatch)` if any vector differs in length from the first.
pub fn pairwise_cosine_similarities(vectors: &[Vec<f64>]) -> Result<Vec<f64>, SphereQlError> {
    if vectors.is_empty() {
        return Ok(Vec::new());
    }
    let dim = vectors[0].len();
    let n = vectors.len();

    for v in vectors.iter().skip(1) {
        if v.len() != dim {
            return Err(SphereQlError::DimensionMismatch {
                expected: dim,
                actual: v.len(),
            });
        }
    }

    let norms: Vec<f64> = vectors
        .iter()
        .map(|v| v.iter().map(|x| x * x).sum::<f64>().sqrt())
        .collect();

    let mut result = Vec::with_capacity(n * (n - 1) / 2);
    for i in 0..n {
        for j in (i + 1)..n {
            let dot: f64 = vectors[i]
                .iter()
                .zip(vectors[j].iter())
                .map(|(a, b)| a * b)
                .sum();
            let denom = norms[i] * norms[j];
            let sim = if denom < f64::EPSILON {
                0.0
            } else {
                dot / denom
            };
            result.push(sim.clamp(-1.0, 1.0));
        }
    }

    Ok(result)
}

/// Index into the upper-triangle flat array returned by [`pairwise_cosine_similarities`].
/// Given `i < j`, returns the index into the flat vector.
///
/// Panics in debug builds if `i >= j` or `j >= n`.
#[inline]
pub fn upper_triangle_index(i: usize, j: usize, n: usize) -> usize {
    debug_assert!(i < j && j < n);
    i * n - i * (i + 1) / 2 + j - i - 1
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use std::f64::consts::{FRAC_PI_2, PI};

    fn point(theta: f64, phi: f64) -> SphericalPoint {
        SphericalPoint::new_unchecked(1.0, theta, phi)
    }

    #[test]
    fn angular_distance_same_point() {
        let p = point(0.5, 1.0);
        assert_relative_eq!(angular_distance(&p, &p), 0.0, epsilon = 1e-12);
    }

    #[test]
    fn angular_distance_antipodal() {
        let a = point(0.0, FRAC_PI_2);
        let b = point(PI, PI - FRAC_PI_2);
        assert_relative_eq!(angular_distance(&a, &b), PI, epsilon = 1e-12);
    }

    #[test]
    fn angular_distance_90_degrees() {
        let a = point(0.0, 0.0);
        let b = point(0.0, FRAC_PI_2);
        assert_relative_eq!(angular_distance(&a, &b), FRAC_PI_2, epsilon = 1e-12);
    }

    #[test]
    fn great_circle_on_unit_sphere() {
        let a = point(0.0, 0.0);
        let b = point(0.0, FRAC_PI_2);
        assert_relative_eq!(
            great_circle_distance(&a, &b, 1.0),
            FRAC_PI_2,
            epsilon = 1e-12
        );
    }

    #[test]
    fn great_circle_with_radius() {
        let a = point(0.0, 0.0);
        let b = point(0.0, FRAC_PI_2);
        let r = 6371.0;
        assert_relative_eq!(
            great_circle_distance(&a, &b, r),
            r * FRAC_PI_2,
            epsilon = 1e-9
        );
    }

    #[test]
    fn chord_distance_same_point() {
        let p = point(1.0, 0.5);
        assert_relative_eq!(chord_distance(&p, &p), 0.0, epsilon = 1e-12);
    }

    #[test]
    fn chord_distance_antipodal_unit_sphere() {
        let a = point(0.0, FRAC_PI_2);
        let b = point(PI, PI - FRAC_PI_2);
        assert_relative_eq!(chord_distance(&a, &b), 2.0, epsilon = 1e-12);
    }

    #[test]
    fn chord_distance_90_degrees_unit_sphere() {
        let a = point(0.0, 0.0);
        let b = point(0.0, FRAC_PI_2);
        assert_relative_eq!(chord_distance(&a, &b), 2.0_f64.sqrt(), epsilon = 1e-12);
    }

    #[test]
    fn euclidean_distance_basic() {
        let a = CartesianPoint::new(0.0, 0.0, 0.0);
        let b = CartesianPoint::new(1.0, 0.0, 0.0);
        assert_relative_eq!(euclidean_distance(&a, &b), 1.0, epsilon = 1e-12);
    }

    #[test]
    fn euclidean_distance_3d() {
        let a = CartesianPoint::new(1.0, 2.0, 3.0);
        let b = CartesianPoint::new(4.0, 6.0, 3.0);
        assert_relative_eq!(euclidean_distance(&a, &b), 5.0, epsilon = 1e-12);
    }

    #[test]
    fn vincenty_stability_near_zero() {
        let a = point(0.0, FRAC_PI_2);
        let b = point(1e-15, FRAC_PI_2);
        let dist = angular_distance(&a, &b);
        assert!(dist >= 0.0);
        assert!(dist < 1e-10);
    }

    #[test]
    fn vincenty_stability_near_pi() {
        let a = point(0.0, 1e-15);
        let b = point(PI, PI - 1e-15);
        let dist = angular_distance(&a, &b);
        assert_relative_eq!(dist, PI, epsilon = 1e-10);
    }

    // --- cosine_proxy tests ---

    #[test]
    fn cosine_proxy_same_direction() {
        let a = [1.0, 0.0, 0.0];
        assert!(cosine_proxy(&a, &a) < 1e-12);
    }

    #[test]
    fn cosine_proxy_orthogonal() {
        let a = [1.0, 0.0, 0.0];
        let b = [0.0, 1.0, 0.0];
        assert_relative_eq!(cosine_proxy(&a, &b), 1.0, epsilon = 1e-12);
    }

    #[test]
    fn cosine_proxy_antipodal() {
        let a = [1.0, 0.0, 0.0];
        let b = [-1.0, 0.0, 0.0];
        assert_relative_eq!(cosine_proxy(&a, &b), 2.0, epsilon = 1e-12);
    }

    #[test]
    fn cosine_proxy_monotone_with_angular_distance() {
        let q = point(0.5, 1.0);
        let a = point(0.6, 1.0);
        let b = point(2.0, 1.0);
        let q_cart = q.unit_cartesian();
        let a_cart = a.unit_cartesian();
        let b_cart = b.unit_cartesian();

        let proxy_a = cosine_proxy(&q_cart, &a_cart);
        let proxy_b = cosine_proxy(&q_cart, &b_cart);
        let angular_a = angular_distance(&q, &a);
        let angular_b = angular_distance(&q, &b);

        assert!(
            (proxy_a < proxy_b) == (angular_a < angular_b),
            "cosine proxy must preserve distance ordering"
        );
    }

    // --- cosine_similarity tests ---

    #[test]
    fn cosine_similarity_identical() {
        let a = vec![1.0, 2.0, 3.0];
        assert_relative_eq!(cosine_similarity(&a, &a).unwrap(), 1.0, epsilon = 1e-12);
    }

    #[test]
    fn cosine_similarity_opposite() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![-1.0, 0.0, 0.0];
        assert_relative_eq!(cosine_similarity(&a, &b).unwrap(), -1.0, epsilon = 1e-12);
    }

    #[test]
    fn cosine_similarity_orthogonal() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![0.0, 1.0, 0.0];
        assert_relative_eq!(cosine_similarity(&a, &b).unwrap(), 0.0, epsilon = 1e-12);
    }

    #[test]
    fn cosine_similarity_zero_vector() {
        let a = vec![0.0, 0.0, 0.0];
        let b = vec![1.0, 2.0, 3.0];
        assert_relative_eq!(cosine_similarity(&a, &b).unwrap(), 0.0, epsilon = 1e-12);
    }

    #[test]
    fn cosine_similarity_dimension_mismatch() {
        let a = vec![1.0, 0.0, 0.0];
        let b = vec![1.0, 0.0];
        let err = cosine_similarity(&a, &b).unwrap_err();
        assert!(matches!(
            err,
            SphereQlError::DimensionMismatch {
                expected: 3,
                actual: 2
            }
        ));
    }

    #[test]
    fn pairwise_cosine_similarities_basic() {
        let vecs = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]];
        let sims = pairwise_cosine_similarities(&vecs).unwrap();
        assert_eq!(sims.len(), 3);
        assert!((sims[upper_triangle_index(0, 1, 3)] - 0.0).abs() < 1e-10);
        assert!(
            (sims[upper_triangle_index(0, 2, 3)] - std::f64::consts::FRAC_1_SQRT_2).abs() < 1e-10
        );
        assert!(
            (sims[upper_triangle_index(1, 2, 3)] - std::f64::consts::FRAC_1_SQRT_2).abs() < 1e-10
        );
    }

    #[test]
    fn pairwise_cosine_similarities_empty() {
        let sims = pairwise_cosine_similarities(&[]).unwrap();
        assert!(sims.is_empty());
    }

    #[test]
    fn pairwise_cosine_similarities_single() {
        let sims = pairwise_cosine_similarities(&[vec![1.0, 2.0]]).unwrap();
        assert!(sims.is_empty());
    }

    #[test]
    fn pairwise_cosine_similarities_dimension_mismatch() {
        let vecs = vec![vec![1.0, 0.0], vec![1.0, 0.0, 0.0]];
        assert!(pairwise_cosine_similarities(&vecs).is_err());
    }
}