helena 0.1.0

//! Data-preservation metrics: does the generated output preserve meaningful
//! structure from the source data? (PRD §15.4, §16.2.)
//!
//! These are total, deterministic, dependency-free measures over *paired
//! per-item observations* — a source-side quantity and an output-side quantity
//! for the same items, in item order — so the core stays decoupled from how
//! those quantities are extracted (a future eval CLI supplies them: a source
//! attribute from the conditioning, an output summary from the generated
//! latent or decoded audio). This is the geometry-preservation family that
//! needs no learned predictor; the recovery-task metrics PRD §15.4 frames as
//! "auxiliary predictors" (class recovery, anomaly detectability) require a
//! trained model and are a separate slice.

use crate::{Error, Result};

/// Linear (Pearson) and rank-order (Spearman) correlation between a source
/// attribute and an output attribute across items (PRD §16.2 "correlation
/// preservation", §15.4 "rank-order preservation"). Both lie in `[-1, 1]`.
///
/// Pearson captures a linear relationship; Spearman captures any monotone one,
/// so a faithful but nonlinear control shows a high Spearman with a lower
/// Pearson. A constant series has no relationship to measure and is rejected.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct CorrelationPreservation {
    pearson: f32,
    spearman: f32,
    count: usize,
}

impl CorrelationPreservation {
    /// Correlate paired per-item scalars. Requires equal-length, finite series
    /// of at least two items, each with non-zero spread.
    pub fn measure(source: &[f32], output: &[f32]) -> Result<Self> {
        if source.len() != output.len() {
            return Err(Error::validation(format!(
                "paired series must have equal length; source {} vs output {}",
                source.len(),
                output.len()
            )));
        }
        if source.len() < 2 {
            return Err(Error::validation(
                "correlation needs at least two paired observations",
            ));
        }
        ensure_finite(source, "source")?;
        ensure_finite(output, "output")?;

        let linear = pearson(source, output)?;
        let rank_order = pearson(&ranks(source), &ranks(output))?;
        Ok(Self {
            pearson: linear as f32,
            spearman: rank_order as f32,
            count: source.len(),
        })
    }

    /// Linear correlation of source and output attributes.
    pub fn pearson(&self) -> f32 {
        self.pearson
    }

    /// Rank-order correlation, insensitive to monotone reparameterization.
    pub fn spearman(&self) -> f32 {
        self.spearman
    }

    /// Number of paired observations measured.
    pub fn count(&self) -> usize {
        self.count
    }
}

/// Mean fraction of each item's `k` nearest source-space neighbors that remain
/// among its `k` nearest output-space neighbors (PRD §16.2 "embedding
/// neighborhood preservation", §15.4 "cluster preservation in embedding
/// space"). In `[0, 1]`: `1.0` iff every item's source neighborhood is
/// preserved exactly.
///
/// Source and output embeddings may have different dimensions (they live in
/// different spaces); only the item count must match. Distance ties are broken
/// by item index so the score is reproducible.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct NeighborhoodPreservation {
    overlap: f32,
    k: usize,
    count: usize,
}

impl NeighborhoodPreservation {
    /// Measure neighborhood overlap. Requires equal-length, non-empty,
    /// dimension-consistent, finite embedding sets and `1 <= k < items`.
    pub fn measure(source: &[Vec<f32>], output: &[Vec<f32>], k: usize) -> Result<Self> {
        let n = source.len();
        if n != output.len() {
            return Err(Error::validation(format!(
                "paired embeddings must have equal length; source {n} vs output {}",
                output.len()
            )));
        }
        if k == 0 {
            return Err(Error::validation("neighborhood size k must be positive"));
        }
        if k >= n {
            return Err(Error::validation(format!(
                "neighborhood size k = {k} needs at least {} items, got {n}",
                k + 1
            )));
        }
        ensure_embeddings(source, "source")?;
        ensure_embeddings(output, "output")?;

        let mut total = 0.0f64;
        for i in 0..n {
            let src = nearest(source, i, k);
            let out = nearest(output, i, k);
            let shared = src.iter().filter(|j| out.contains(j)).count();
            total += shared as f64 / k as f64;
        }
        Ok(Self {
            overlap: (total / n as f64) as f32,
            k,
            count: n,
        })
    }

    /// Mean preserved-neighbor fraction across items, in `[0, 1]`.
    pub fn overlap(&self) -> f32 {
        self.overlap
    }

    /// Neighborhood size the score was measured at.
    pub fn k(&self) -> usize {
        self.k
    }

    /// Number of items measured.
    pub fn count(&self) -> usize {
        self.count
    }
}

/// Pearson correlation in `f64`, rejecting a constant (zero-variance) series.
/// The result is clamped to `[-1, 1]` to absorb floating-point drift past the
/// Cauchy-Schwarz bound.
fn pearson(x: &[f32], y: &[f32]) -> Result<f64> {
    let n = x.len() as f64;
    let mx = x.iter().map(|&v| v as f64).sum::<f64>() / n;
    let my = y.iter().map(|&v| v as f64).sum::<f64>() / n;
    let (mut sxy, mut sxx, mut syy) = (0.0f64, 0.0f64, 0.0f64);
    for (&xi, &yi) in x.iter().zip(y) {
        let (dx, dy) = (xi as f64 - mx, yi as f64 - my);
        sxy += dx * dy;
        sxx += dx * dx;
        syy += dy * dy;
    }
    if sxx <= 0.0 || syy <= 0.0 {
        return Err(Error::validation(
            "a series is constant; correlation is undefined",
        ));
    }
    Ok((sxy / (sxx * syy).sqrt()).clamp(-1.0, 1.0))
}

/// Fractional (average) ranks of `xs`, ties sharing their mean rank. Inputs are
/// pre-validated finite, so the comparison never observes a `NaN`.
fn ranks(xs: &[f32]) -> Vec<f32> {
    let mut order: Vec<usize> = (0..xs.len()).collect();
    order.sort_by(|&a, &b| xs[a].partial_cmp(&xs[b]).expect("finite values compare"));
    let mut ranks = vec![0.0f32; xs.len()];
    let mut i = 0;
    while i < order.len() {
        let mut j = i;
        while j + 1 < order.len() && xs[order[j + 1]] == xs[order[i]] {
            j += 1;
        }
        // 1-based ranks i+1..=j+1 average to (i + j) / 2 + 1.
        let avg = (i + j) as f32 / 2.0 + 1.0;
        for &idx in &order[i..=j] {
            ranks[idx] = avg;
        }
        i = j + 1;
    }
    ranks
}

/// The `k` items nearest to `set[i]` (excluding `i`) by squared Euclidean
/// distance, ties broken by item index for a reproducible neighborhood.
fn nearest(set: &[Vec<f32>], i: usize, k: usize) -> Vec<usize> {
    let mut others: Vec<(f64, usize)> = (0..set.len())
        .filter(|&j| j != i)
        .map(|j| (sq_distance(&set[i], &set[j]), j))
        .collect();
    others.sort_by(|a, b| {
        a.0.partial_cmp(&b.0)
            .expect("finite distances compare")
            .then(a.1.cmp(&b.1))
    });
    others.into_iter().take(k).map(|(_, j)| j).collect()
}

/// Squared Euclidean distance between equal-dimension embeddings, in `f64`.
fn sq_distance(a: &[f32], b: &[f32]) -> f64 {
    a.iter()
        .zip(b)
        .map(|(&x, &y)| {
            let d = x as f64 - y as f64;
            d * d
        })
        .sum()
}

fn ensure_finite(xs: &[f32], label: &str) -> Result<()> {
    if let Some(bad) = xs.iter().find(|v| !v.is_finite()) {
        return Err(Error::validation(format!(
            "{label} series must be finite; found {bad}"
        )));
    }
    Ok(())
}

/// Validate an embedding set: a shared positive dimension and finite values.
fn ensure_embeddings(set: &[Vec<f32>], label: &str) -> Result<()> {
    let dim = set[0].len();
    if dim == 0 {
        return Err(Error::validation(format!(
            "{label} embeddings must be non-empty"
        )));
    }
    for (i, v) in set.iter().enumerate() {
        if v.len() != dim {
            return Err(Error::validation(format!(
                "{label} embeddings must share a dimension; item 0 has {dim}, item {i} has {}",
                v.len()
            )));
        }
        ensure_finite(v, label)?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    #[test]
    fn correlation_perfect_and_anti() {
        let c = CorrelationPreservation::measure(&[1.0, 2.0, 3.0], &[2.0, 4.0, 6.0]).unwrap();
        assert!((c.pearson() - 1.0).abs() < 1e-6);
        assert!((c.spearman() - 1.0).abs() < 1e-6);
        assert_eq!(c.count(), 3);

        let anti = CorrelationPreservation::measure(&[1.0, 2.0, 3.0], &[6.0, 4.0, 2.0]).unwrap();
        assert!((anti.pearson() + 1.0).abs() < 1e-6);
        assert!((anti.spearman() + 1.0).abs() < 1e-6);
    }

    #[test]
    fn spearman_captures_monotone_nonlinear() {
        // y = x^2 over positive x is monotone but not linear: rank-order is
        // perfectly preserved while the linear fit is not.
        let c = CorrelationPreservation::measure(&[1.0, 2.0, 3.0, 4.0], &[1.0, 4.0, 9.0, 16.0])
            .unwrap();
        assert!((c.spearman() - 1.0).abs() < 1e-6);
        assert!(c.pearson() < 1.0 - 1e-4);
    }

    #[test]
    fn correlation_rejects_degenerate_inputs() {
        assert!(CorrelationPreservation::measure(&[1.0, 2.0], &[1.0]).is_err());
        assert!(CorrelationPreservation::measure(&[1.0], &[1.0]).is_err());
        assert!(CorrelationPreservation::measure(&[2.0, 2.0, 2.0], &[1.0, 2.0, 3.0]).is_err());
        assert!(CorrelationPreservation::measure(&[1.0, 2.0], &[f32::NAN, 1.0]).is_err());
    }

    #[test]
    fn neighborhood_perfect_when_identical() {
        let pts = vec![vec![0.0], vec![1.0], vec![2.0], vec![5.0]];
        let n = NeighborhoodPreservation::measure(&pts, &pts, 1).unwrap();
        assert_eq!(n.overlap(), 1.0);
        assert_eq!(n.k(), 1);
        assert_eq!(n.count(), 4);
    }

    #[test]
    fn neighborhood_partial_overlap() {
        // Collinear 0,1,2,3; the output swaps the last two coordinates, so the
        // two interior items keep their nearest neighbor and the outer two do
        // not: overlap = (1 + 1 + 0 + 0) / 4.
        let source = vec![vec![0.0], vec![1.0], vec![2.0], vec![3.0]];
        let output = vec![vec![0.0], vec![1.0], vec![3.0], vec![2.0]];
        let n = NeighborhoodPreservation::measure(&source, &output, 1).unwrap();
        assert!((n.overlap() - 0.5).abs() < 1e-6);
    }

    #[test]
    fn neighborhood_rejects_bad_shapes_and_k() {
        let pts = vec![vec![0.0], vec![1.0], vec![2.0]];
        assert!(NeighborhoodPreservation::measure(&pts, &pts, 0).is_err());
        assert!(NeighborhoodPreservation::measure(&pts, &pts, 3).is_err());
        assert!(NeighborhoodPreservation::measure(&pts, &pts[..2], 1).is_err());

        let ragged = vec![vec![0.0, 1.0], vec![1.0], vec![2.0, 3.0]];
        assert!(NeighborhoodPreservation::measure(&ragged, &pts, 1).is_err());

        let nonfinite = vec![vec![0.0], vec![f32::INFINITY], vec![2.0]];
        assert!(NeighborhoodPreservation::measure(&nonfinite, &pts, 1).is_err());
    }

    proptest! {
        #[test]
        fn correlation_identity_is_one(xs in proptest::collection::vec(-100.0f32..100.0, 2..32)) {
            prop_assume!(xs.iter().any(|&v| v != xs[0]));
            let c = CorrelationPreservation::measure(&xs, &xs).unwrap();
            prop_assert!((c.pearson() - 1.0).abs() < 1e-3);
            prop_assert!((c.spearman() - 1.0).abs() < 1e-3);
        }

        #[test]
        fn correlation_sign_follows_affine_slope(
            xs in proptest::collection::vec(-50.0f32..50.0, 2..32),
            a in prop_oneof![-10.0f32..-0.1, 0.1f32..10.0],
            b in -5.0f32..5.0,
        ) {
            prop_assume!(xs.iter().any(|&v| v != xs[0]));
            let ys: Vec<f32> = xs.iter().map(|&x| a * x + b).collect();
            let c = CorrelationPreservation::measure(&xs, &ys).unwrap();
            prop_assert!((c.pearson() - a.signum()).abs() < 1e-3, "pearson {}", c.pearson());
            prop_assert!((c.spearman() - a.signum()).abs() < 1e-3, "spearman {}", c.spearman());
        }

        #[test]
        fn correlation_within_bounds(
            xs in proptest::collection::vec(-50.0f32..50.0, 2..32),
            ys in proptest::collection::vec(-50.0f32..50.0, 2..32),
        ) {
            let n = xs.len().min(ys.len());
            prop_assume!(xs[..n].iter().any(|&v| v != xs[0]));
            prop_assume!(ys[..n].iter().any(|&v| v != ys[0]));
            let c = CorrelationPreservation::measure(&xs[..n], &ys[..n]).unwrap();
            prop_assert!((-1.0..=1.0).contains(&c.pearson()));
            prop_assert!((-1.0..=1.0).contains(&c.spearman()));
        }

        #[test]
        fn neighborhood_identity_is_one(
            set in proptest::collection::vec(
                proptest::collection::vec(-20.0f32..20.0, 1..6),
                3..12,
            ),
            k in 1usize..3,
        ) {
            // Equal-dimension embeddings: take the common prefix length.
            let dim = set[0].len();
            let pts: Vec<Vec<f32>> = set.iter().filter(|v| v.len() == dim).cloned().collect();
            prop_assume!(pts.len() > k);
            let n = NeighborhoodPreservation::measure(&pts, &pts, k).unwrap();
            prop_assert_eq!(n.overlap(), 1.0);
        }

        #[test]
        fn neighborhood_within_bounds(
            base in proptest::collection::vec(
                proptest::collection::vec(-20.0f32..20.0, 2..3),
                4..12,
            ),
            shuffle in proptest::collection::vec(-20.0f32..20.0, 8..24),
            k in 1usize..3,
        ) {
            let source = base.clone();
            let output: Vec<Vec<f32>> = source
                .iter()
                .enumerate()
                .map(|(i, _)| vec![shuffle[(2 * i) % shuffle.len()], shuffle[(2 * i + 1) % shuffle.len()]])
                .collect();
            prop_assume!(source.len() > k);
            let n = NeighborhoodPreservation::measure(&source, &output, k).unwrap();
            prop_assert!((0.0..=1.0).contains(&n.overlap()));
        }
    }
}