ruvllm-esp32 0.3.1

//! Hyperbolic Embeddings for RuvLLM ESP32
//!
//! Implements hyperbolic geometry distance metrics optimized for microcontrollers.
//! Hyperbolic spaces are ideal for hierarchical data (taxonomies, knowledge graphs)
//! as they naturally represent tree-like structures with exponentially growing space.
//!
//! # Models
//!
//! ## Poincaré Ball Model
//! - Points in unit ball: ||x|| < 1
//! - Conformal (preserves angles)
//! - Intuitive visualization
//! - Distance: d(x,y) = arcosh(1 + 2||x-y||² / ((1-||x||²)(1-||y||²)))
//!
//! ## Lorentz (Hyperboloid) Model
//! - Points on hyperboloid: -x₀² + x₁² + ... + xₙ² = -1, x₀ > 0
//! - More numerically stable
//! - Better for optimization
//! - Distance: d(x,y) = arcosh(-⟨x,y⟩_L) where ⟨x,y⟩_L = -x₀y₀ + Σxᵢyᵢ
//!
//! # INT8 Implementation
//!
//! We use fixed-point arithmetic scaled to [-127, 127] range:
//! - Poincaré: Scale factor maps to radius < 1 (use 100 = 0.787 radius)
//! - Lorentz: First component is timelike (computed from spatial)

use heapless::Vec as HVec;
use libm::{acoshf, sqrtf, coshf, sinhf};

/// Scale factor for INT8 to float conversion
/// 127 in INT8 = 0.787 in Poincaré ball (safe margin from boundary)
const POINCARE_SCALE: f32 = 127.0 / 0.787;

/// Curvature of hyperbolic space (negative)
/// Default: -1.0 (can be adjusted for different geometries)
const DEFAULT_CURVATURE: f32 = -1.0;

/// Hyperbolic embedding configuration
#[derive(Debug, Clone, Copy)]
pub struct HyperbolicConfig {
    /// Curvature of the hyperbolic space (negative value)
    pub curvature: f32,
    /// Dimension of the embedding
    pub dim: usize,
    /// Epsilon for numerical stability
    pub eps: f32,
}

impl Default for HyperbolicConfig {
    fn default() -> Self {
        Self {
            curvature: DEFAULT_CURVATURE,
            dim: 32,
            eps: 1e-5,
        }
    }
}

// ============================================================================
// Poincaré Ball Model
// ============================================================================

/// Poincaré distance between two INT8 vectors
///
/// Formula: d(x,y) = (1/√|c|) * arcosh(1 + 2c||x-y||² / ((1-c||x||²)(1-c||y||²)))
///
/// For c = -1: d(x,y) = arcosh(1 + 2||x-y||² / ((1-||x||²)(1-||y||²)))
///
/// Returns distance scaled to i32 (multiply by 1000 for precision)
pub fn poincare_distance_i8(a: &[i8], b: &[i8]) -> i32 {
    poincare_distance_i8_with_curvature(a, b, DEFAULT_CURVATURE)
}

/// Poincaré distance with custom curvature
pub fn poincare_distance_i8_with_curvature(a: &[i8], b: &[i8], curvature: f32) -> i32 {
    let c = -curvature.abs(); // Ensure negative
    let sqrt_c = sqrtf(-c);

    // Convert to float and scale to Poincaré ball
    let scale = 1.0 / POINCARE_SCALE;

    // Calculate ||x||², ||y||², ||x-y||²
    let mut norm_a_sq: f32 = 0.0;
    let mut norm_b_sq: f32 = 0.0;
    let mut diff_sq: f32 = 0.0;

    for (x, y) in a.iter().zip(b.iter()) {
        let xf = (*x as f32) * scale;
        let yf = (*y as f32) * scale;

        norm_a_sq += xf * xf;
        norm_b_sq += yf * yf;
        diff_sq += (xf - yf) * (xf - yf);
    }

    // Clamp norms to stay inside ball (numerical stability)
    let max_norm = 1.0 - 1e-5;
    norm_a_sq = norm_a_sq.min(max_norm * max_norm);
    norm_b_sq = norm_b_sq.min(max_norm * max_norm);

    // Calculate distance
    // d = (1/√|c|) * arcosh(1 + 2|c| * ||x-y||² / ((1 - |c|||x||²)(1 - |c|||y||²)))
    let numerator = 2.0 * (-c) * diff_sq;
    let denom_a = 1.0 - (-c) * norm_a_sq;
    let denom_b = 1.0 - (-c) * norm_b_sq;
    let denominator = denom_a * denom_b;

    // Avoid division by zero
    if denominator < 1e-10 {
        return i32::MAX / 2; // Large distance for boundary points
    }

    let arg = 1.0 + numerator / denominator;

    // arcosh(x) = ln(x + sqrt(x² - 1)) for x >= 1
    let arg_clamped = arg.max(1.0); // Ensure valid for arcosh
    let dist = acoshf(arg_clamped) / sqrt_c;

    // Scale to i32 (multiply by 1000 for 3 decimal places)
    (dist * 1000.0) as i32
}

/// Convert Euclidean INT8 vector to Poincaré ball
/// Projects onto ball with exponential map at origin
pub fn to_poincare_i8(euclidean: &[i8]) -> HVec<i8, 64> {
    let mut result: HVec<i8, 64> = HVec::new();

    // Calculate norm
    let mut norm_sq: f32 = 0.0;
    for x in euclidean {
        let xf = *x as f32;
        norm_sq += xf * xf;
    }
    let norm = sqrtf(norm_sq);

    if norm < 1e-6 {
        // Origin maps to origin
        for _ in 0..euclidean.len() {
            let _ = result.push(0);
        }
        return result;
    }

    // Exponential map at origin: exp_0(v) = tanh(||v||/2) * v/||v||
    let scale = (norm / (2.0 * POINCARE_SCALE)).tanh() * POINCARE_SCALE / norm;

    for x in euclidean {
        let mapped = ((*x as f32) * scale).clamp(-127.0, 127.0) as i8;
        let _ = result.push(mapped);
    }

    result
}

// ============================================================================
// Lorentz (Hyperboloid) Model
// ============================================================================

/// Lorentz inner product: ⟨x,y⟩_L = -x₀y₀ + x₁y₁ + ... + xₙyₙ
///
/// For points on hyperboloid: ⟨x,x⟩_L = -1/c (= 1 for c = -1)
fn lorentz_inner_product(a: &[f32], b: &[f32]) -> f32 {
    if a.is_empty() || b.is_empty() {
        return 0.0;
    }

    // First component is timelike (negative sign)
    let mut result = -a[0] * b[0];

    // Remaining components are spacelike (positive sign)
    for (x, y) in a[1..].iter().zip(b[1..].iter()) {
        result += x * y;
    }

    result
}

/// Lorentz distance between two INT8 vectors
///
/// Formula: d(x,y) = (1/√|c|) * arcosh(-c * ⟨x,y⟩_L)
///
/// For c = -1: d(x,y) = arcosh(⟨x,y⟩_L) where ⟨x,y⟩_L uses Minkowski metric
///
/// Note: Input vectors should be in Lorentz format where first component
/// is the timelike coordinate. If input is spatial only, use `lorentz_distance_spatial_i8`.
pub fn lorentz_distance_i8(a: &[i8], b: &[i8]) -> i32 {
    lorentz_distance_i8_with_curvature(a, b, DEFAULT_CURVATURE)
}

/// Lorentz distance with custom curvature
pub fn lorentz_distance_i8_with_curvature(a: &[i8], b: &[i8], curvature: f32) -> i32 {
    let c = -curvature.abs();
    let sqrt_c = sqrtf(-c);

    // Convert to float
    let scale = 1.0 / 127.0;
    let a_f: HVec<f32, 65> = a.iter().map(|&x| x as f32 * scale).collect();
    let b_f: HVec<f32, 65> = b.iter().map(|&x| x as f32 * scale).collect();

    // Calculate Lorentz inner product
    let inner = lorentz_inner_product(&a_f, &b_f);

    // Distance = arcosh(-c * ⟨x,y⟩_L) / √|c|
    let arg = (-c * inner).max(1.0); // Clamp for numerical stability
    let dist = acoshf(arg) / sqrt_c;

    (dist * 1000.0) as i32
}

/// Lorentz distance from spatial coordinates only
///
/// Automatically computes the timelike component from spatial coordinates
/// using the hyperboloid constraint: x₀ = √(1/|c| + ||x_spatial||²)
pub fn lorentz_distance_spatial_i8(a: &[i8], b: &[i8]) -> i32 {
    lorentz_distance_spatial_i8_with_curvature(a, b, DEFAULT_CURVATURE)
}

/// Lorentz distance from spatial coordinates with custom curvature
pub fn lorentz_distance_spatial_i8_with_curvature(a: &[i8], b: &[i8], curvature: f32) -> i32 {
    let c = -curvature.abs();
    let sqrt_c = sqrtf(-c);
    let k = 1.0 / (-c); // = 1 for c = -1

    let scale = 1.0 / POINCARE_SCALE; // Use same scale as Poincaré for consistency

    // Calculate spatial norms
    let mut norm_a_sq: f32 = 0.0;
    let mut norm_b_sq: f32 = 0.0;
    let mut spatial_dot: f32 = 0.0;

    for (x, y) in a.iter().zip(b.iter()) {
        let xf = (*x as f32) * scale;
        let yf = (*y as f32) * scale;

        norm_a_sq += xf * xf;
        norm_b_sq += yf * yf;
        spatial_dot += xf * yf;
    }

    // Compute timelike components: x₀ = √(k + ||x||²)
    let t_a = sqrtf(k + norm_a_sq);
    let t_b = sqrtf(k + norm_b_sq);

    // Lorentz inner product: -t_a*t_b + spatial_dot
    let inner = -t_a * t_b + spatial_dot;

    // Distance = arcosh(-c * inner) / √|c| = arcosh(inner) for c = -1
    let arg = (-c * inner).max(1.0);
    let dist = acoshf(arg) / sqrt_c;

    (dist * 1000.0) as i32
}

/// Convert Euclidean INT8 vector to Lorentz hyperboloid
///
/// Maps point x to (√(1 + ||x||²), x) on hyperboloid
pub fn to_lorentz_i8(spatial: &[i8]) -> HVec<i8, 65> {
    let mut result: HVec<i8, 65> = HVec::new();

    let scale = 1.0 / POINCARE_SCALE;

    // Calculate spatial norm squared
    let mut norm_sq: f32 = 0.0;
    for x in spatial {
        let xf = (*x as f32) * scale;
        norm_sq += xf * xf;
    }

    // Timelike component: t = √(1 + ||x||²)
    let t = sqrtf(1.0 + norm_sq);
    let t_scaled = (t * 127.0).clamp(-127.0, 127.0) as i8;
    let _ = result.push(t_scaled);

    // Spatial components (already scaled)
    for x in spatial {
        let _ = result.push(*x);
    }

    result
}

// ============================================================================
// Conversions between models
// ============================================================================

/// Convert Poincaré ball point to Lorentz hyperboloid
///
/// Formula: (x₀, x₁...) = ((1 + ||p||²)/(1 - ||p||²), 2p/(1 - ||p||²))
pub fn poincare_to_lorentz(poincare: &[f32]) -> HVec<f32, 65> {
    let mut result: HVec<f32, 65> = HVec::new();

    let mut norm_sq: f32 = 0.0;
    for x in poincare {
        norm_sq += x * x;
    }

    // Clamp to stay inside ball
    norm_sq = norm_sq.min(1.0 - 1e-5);

    let denom = 1.0 - norm_sq;

    // Timelike component
    let t = (1.0 + norm_sq) / denom;
    let _ = result.push(t);

    // Spatial components
    let spatial_scale = 2.0 / denom;
    for x in poincare {
        let _ = result.push(x * spatial_scale);
    }

    result
}

/// Convert Lorentz hyperboloid point to Poincaré ball
///
/// Formula: p = x_spatial / (1 + x₀)
pub fn lorentz_to_poincare(lorentz: &[f32]) -> HVec<f32, 64> {
    let mut result: HVec<f32, 64> = HVec::new();

    if lorentz.is_empty() {
        return result;
    }

    let t = lorentz[0];
    let scale = 1.0 / (1.0 + t);

    for x in &lorentz[1..] {
        let _ = result.push(x * scale);
    }

    result
}

// ============================================================================
// Hyperbolic operations
// ============================================================================

/// Möbius addition in Poincaré ball: x ⊕ y
///
/// Used for translation in hyperbolic space
pub fn mobius_add(x: &[f32], y: &[f32], curvature: f32) -> HVec<f32, 64> {
    let c = -curvature.abs();
    let mut result: HVec<f32, 64> = HVec::new();

    let mut x_sq: f32 = 0.0;
    let mut y_sq: f32 = 0.0;
    let mut xy: f32 = 0.0;

    for (a, b) in x.iter().zip(y.iter()) {
        x_sq += a * a;
        y_sq += b * b;
        xy += a * b;
    }

    let c_abs = -c;
    let num_factor = 1.0 + 2.0 * c_abs * xy + c_abs * y_sq;
    let denom = 1.0 + 2.0 * c_abs * xy + c_abs * c_abs * x_sq * y_sq;

    if denom.abs() < 1e-10 {
        return result;
    }

    let y_factor = (1.0 - c_abs * x_sq) / denom;
    let x_factor = num_factor / denom;

    for (a, b) in x.iter().zip(y.iter()) {
        let val = a * x_factor + b * y_factor;
        let _ = result.push(val);
    }

    result
}

/// Hyperbolic midpoint between two points (Poincaré ball)
pub fn hyperbolic_midpoint(a: &[i8], b: &[i8]) -> HVec<i8, 64> {
    let scale = 1.0 / POINCARE_SCALE;

    let a_f: HVec<f32, 64> = a.iter().map(|&x| x as f32 * scale).collect();
    let b_f: HVec<f32, 64> = b.iter().map(|&x| x as f32 * scale).collect();

    // For midpoint, we use: mid = (a ⊕ b) / 2 in hyperbolic sense
    // Simplified: geodesic midpoint
    let sum = mobius_add(&a_f, &b_f, DEFAULT_CURVATURE);

    // Scale back to INT8
    sum.iter()
        .map(|&x| ((x * 0.5) * POINCARE_SCALE).clamp(-127.0, 127.0) as i8)
        .collect()
}

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

    #[test]
    fn test_poincare_distance_zero() {
        // Same point should have zero distance
        let a = [0i8, 0, 0, 0];
        let b = [0i8, 0, 0, 0];
        let dist = poincare_distance_i8(&a, &b);
        assert!(dist < 10, "Distance at origin should be ~0, got {}", dist);
    }

    #[test]
    fn test_poincare_distance_symmetric() {
        let a = [10i8, 20, 30, 40];
        let b = [50i8, 60, 70, 80];

        let d1 = poincare_distance_i8(&a, &b);
        let d2 = poincare_distance_i8(&b, &a);

        assert_eq!(d1, d2, "Distance should be symmetric");
    }

    #[test]
    fn test_poincare_distance_triangle_inequality() {
        let a = [10i8, 0, 0, 0];
        let b = [0i8, 10, 0, 0];
        let c = [0i8, 0, 10, 0];

        let ab = poincare_distance_i8(&a, &b);
        let bc = poincare_distance_i8(&b, &c);
        let ac = poincare_distance_i8(&a, &c);

        assert!(ac <= ab + bc + 1, "Triangle inequality violated");
    }

    #[test]
    fn test_lorentz_distance_spatial() {
        // Use larger separation to ensure measurable distance
        let a = [10i8, 20, 30];
        let b = [60i8, 70, 80];

        let dist = lorentz_distance_spatial_i8(&a, &b);
        assert!(dist >= 0, "Distance should be non-negative, got {}", dist);

        // Same point should have zero distance
        let zero_dist = lorentz_distance_spatial_i8(&a, &a);
        assert!(zero_dist < 10, "Same point distance should be ~0, got {}", zero_dist);
    }

    #[test]
    fn test_lorentz_distance_symmetric() {
        let a = [10i8, 20, 30];
        let b = [50i8, 60, 70];

        let d1 = lorentz_distance_spatial_i8(&a, &b);
        let d2 = lorentz_distance_spatial_i8(&b, &a);

        assert_eq!(d1, d2, "Lorentz distance should be symmetric");
    }

    #[test]
    fn test_to_poincare_origin() {
        let euclidean = [0i8, 0, 0, 0];
        let poincare = to_poincare_i8(&euclidean);

        for x in poincare.iter() {
            assert_eq!(*x, 0, "Origin should map to origin");
        }
    }

    #[test]
    fn test_to_lorentz() {
        let spatial = [50i8, 50, 50];
        let lorentz = to_lorentz_i8(&spatial);

        // First component (timelike) should be >= 127 (scaled 1.0)
        assert!(lorentz[0] > 0, "Timelike component should be positive");
        assert_eq!(lorentz.len(), spatial.len() + 1, "Should add timelike component");
    }

    #[test]
    fn test_poincare_to_lorentz_roundtrip() {
        let original = [0.3f32, 0.2, 0.1];
        let lorentz = poincare_to_lorentz(&original);
        let back = lorentz_to_poincare(&lorentz);

        for (a, b) in original.iter().zip(back.iter()) {
            assert!((a - b).abs() < 0.01, "Roundtrip should preserve values");
        }
    }

    #[test]
    fn test_hyperbolic_midpoint() {
        let a = [20i8, 0, 0, 0];
        let b = [-20i8, 0, 0, 0];

        let mid = hyperbolic_midpoint(&a, &b);

        // Midpoint should be near origin
        let norm: i32 = mid.iter().map(|&x| (x as i32).abs()).sum();
        assert!(norm < 50, "Midpoint of symmetric points should be near origin");
    }

    #[test]
    fn test_boundary_behavior() {
        // Points near boundary should have large distances
        let center = [0i8, 0, 0, 0];
        let near_boundary = [120i8, 0, 0, 0]; // Close to ||x|| = 1

        let dist = poincare_distance_i8(&center, &near_boundary);
        assert!(dist > 500, "Distance to boundary should be large");
    }
}