holographic-memory 0.6.3

// Copyright 2024-2026 WritersLogic Contributors
// SPDX-License-Identifier: Apache-2.0

//! Hopfield-Fenchel-Young (HFYN) energy-based associative memory retrieval.
//!
//! Implements sparse Hopfield retrieval using Tsallis α-entropy transformations
//! (Santos et al., JMLR 2025). For sparse binary vectors, the update rule is:
//!
//!   scores_i = β · |query ∩ pattern_i|       (intersection count)
//!   weights  = entmax_α(scores)               (sparse attention)
//!   result   = patterns with non-zero weights
//!
//! Key properties:
//! - **Exact retrieval**: for well-separated patterns, converges to a single
//!   pattern in one step (margin m = 1/(α-1), requires Δ ≥ m/β).
//! - **Sparse output**: α > 1 produces exact zeros, naturally determining
//!   the number of relevant results (no fixed k needed).
//! - **Exponential capacity**: O(ζ^(D/2)) patterns recoverable for dimension D.

use std::borrow::Borrow;

use crate::core::entangled::EntangledHVec;
use crate::core::types::RetrievalResult;

/// Hopfield retrieval configuration.
#[derive(Clone, Debug)]
pub struct HopfieldConfig {
    /// Inverse temperature. Higher β = sharper retrieval (fewer results).
    /// Must be > 0. Typical range: 10.0–500.0 (scores are Jaccard in `[0, 1]`).
    pub beta: f64,
    /// Tsallis entropy parameter controlling sparsity of attention weights.
    /// α=1: softmax (dense, no exact zeros — equivalent to standard attention).
    /// α=1.5: 1.5-entmax (moderately sparse).
    /// α=2: sparsemax (maximally sparse, margin m=1).
    pub alpha: f64,
    /// Maximum Hopfield iterations for query refinement.
    /// 1 = single-step retrieval (sufficient for well-separated patterns).
    pub max_iter: usize,
}

impl Default for HopfieldConfig {
    fn default() -> Self {
        Self {
            beta: 100.0,
            alpha: 2.0,
            max_iter: 1,
        }
    }
}

/// Perform Hopfield-HFYN retrieval against a set of stored patterns.
///
/// Returns patterns with non-zero entmax weights, sorted by weight descending.
/// The number of results is determined by the sparsity of the entmax output,
/// not by a fixed k — but at most `max_results` are returned.
pub fn hopfield_query<S: AsRef<str>, V: Borrow<EntangledHVec>>(
    query: &EntangledHVec,
    patterns: &[(S, V)],
    config: &HopfieldConfig,
    max_results: usize,
) -> Vec<RetrievalResult> {
    if max_results == 0 || patterns.is_empty() || config.beta <= 0.0 {
        return Vec::new();
    }

    // Jaccard rather than raw overlap: multi-scale encoding produces vectors of varying
    // density, and raw counts favour the denser ones.
    let scores: Vec<f64> = patterns
        .iter()
        .map(|(_, pat)| config.beta * query.similarity(pat.borrow()))
        .collect();

    let weights = entmax(&scores, config.alpha);

    let mut results: Vec<RetrievalResult> = weights
        .iter()
        .enumerate()
        .filter(|(_, &w)| w > 1e-12)
        .map(|(i, &w)| RetrievalResult {
            id: patterns[i].0.as_ref().to_owned(),
            similarity: w,
        })
        .collect();

    results.sort_unstable_by(|a, b| {
        b.similarity
            .partial_cmp(&a.similarity)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    results.truncate(max_results);
    results
}

/// Iterative Hopfield retrieval with query refinement.
///
/// Each iteration reconstructs the query as a weighted combination of patterns
/// (via weighted bundle), then re-queries. Converges to an attractor basin.
pub fn hopfield_query_iterative<S: AsRef<str>, V: Borrow<EntangledHVec>>(
    query: &EntangledHVec,
    patterns: &[(S, V)],
    config: &HopfieldConfig,
    max_results: usize,
) -> Vec<RetrievalResult> {
    if max_results == 0 || patterns.is_empty() {
        return Vec::new();
    }
    if config.max_iter <= 1 {
        return hopfield_query(query, patterns, config, max_results);
    }

    let mut current_query = query.clone();
    let mut scores = Vec::with_capacity(patterns.len());
    for _ in 0..config.max_iter - 1 {
        scores.clear();
        scores.extend(
            patterns
                .iter()
                .map(|(_, pat)| config.beta * current_query.similarity(pat.borrow())),
        );

        let weights = entmax(&scores, config.alpha);

        // Reconstruct query as weighted bundle of attended patterns
        current_query = weighted_reconstruct(patterns, &weights, query.dim);
        if current_query.indices.is_empty() {
            break;
        }
    }

    drop(scores);
    hopfield_query(&current_query, patterns, config, max_results)
}

/// α-entmax over `scores`: softmax for α ≤ 1, 1.5-entmax for α = 1.5, sparsemax otherwise.
fn entmax(scores: &[f64], alpha: f64) -> Vec<f64> {
    if alpha <= 1.0 + f64::EPSILON {
        softmax(scores)
    } else if (alpha - 1.5).abs() < f64::EPSILON {
        entmax_15(scores)
    } else {
        sparsemax(scores)
    }
}

/// Sparsemax transformation (α=2 entmax).
///
/// Projects scores onto the probability simplex with maximum sparsity.
/// Margin m = 1/(α-1) = 1 for α=2.
///
/// Algorithm (Martins & Astudillo, 2016):
/// 1. Sort scores descending
/// 2. Find support size k = max{j : 1 + j·z_(j) > Σ_{i≤j} z_(i)}
/// 3. Threshold τ = (Σ_{i≤k} z_(i) - 1) / k
/// 4. p_i = max(z_i - τ, 0)
fn sparsemax(scores: &[f64]) -> Vec<f64> {
    let n = scores.len();
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![1.0];
    }

    let mut sorted_indices: Vec<usize> = (0..n).collect();
    sorted_indices.sort_unstable_by(|&a, &b| {
        scores[b]
            .partial_cmp(&scores[a])
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let mut cumsum = 0.0;
    let mut support_size = 0;

    for (j, &idx) in sorted_indices.iter().enumerate() {
        cumsum += scores[idx];
        // 1-indexed: j+1
        if 1.0 + (j + 1) as f64 * scores[idx] > cumsum {
            support_size = j + 1;
        } else {
            break;
        }
    }

    if support_size == 0 {
        // All scores identical or degenerate — uniform
        let uniform = 1.0 / n as f64;
        return vec![uniform; n];
    }

    let tau_sum: f64 = sorted_indices[..support_size]
        .iter()
        .map(|&i| scores[i])
        .sum();
    let tau = (tau_sum - 1.0) / support_size as f64;

    scores.iter().map(|&s| (s - tau).max(0.0)).collect()
}

/// 1.5-entmax transformation (α=1.5).
///
/// Intermediate sparsity between softmax and sparsemax.
/// Margin m = 1/(α-1) = 2 for α=1.5.
///
/// Uses bisection to find threshold τ such that Σ max(z_i - τ, 0)^(1/(α-1)) = 1.
/// For α=1.5: Σ max(z_i - τ, 0)^2 = 1.
fn entmax_15(scores: &[f64]) -> Vec<f64> {
    let n = scores.len();
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![1.0];
    }

    let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);

    // Bisection for τ: find τ such that Σ max(s_i - τ, 0)^2 = 1
    let mut lo = max_score - 1.0;
    let mut hi = max_score;

    while sum_clipped_sq(scores, lo) < 1.0 {
        lo -= 1.0;
    }

    for _ in 0..64 {
        let mid = (lo + hi) / 2.0;
        let s = sum_clipped_sq(scores, mid);
        if s > 1.0 {
            lo = mid;
        } else {
            hi = mid;
        }
    }

    let tau = (lo + hi) / 2.0;
    scores.iter().map(|&s| (s - tau).max(0.0).powi(2)).collect()
}

fn sum_clipped_sq(scores: &[f64], tau: f64) -> f64 {
    scores.iter().map(|&s| (s - tau).max(0.0).powi(2)).sum()
}

/// Standard softmax (α=1 limit). Dense, no exact zeros.
fn softmax(scores: &[f64]) -> Vec<f64> {
    if scores.is_empty() {
        return Vec::new();
    }
    let max = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let exps: Vec<f64> = scores.iter().map(|&s| (s - max).exp()).collect();
    let sum: f64 = exps.iter().sum();
    if sum == 0.0 {
        let uniform = 1.0 / scores.len() as f64;
        return vec![uniform; scores.len()];
    }
    exps.iter().map(|&e| e / sum).collect()
}

/// Reconstruct a query vector from weighted patterns.
///
/// For sparse binary vectors: count weighted frequency of each index across
/// patterns with non-zero weight, keep indices with highest weighted frequency
/// up to the target active count.
fn weighted_reconstruct<S, V: Borrow<EntangledHVec>>(
    patterns: &[(S, V)],
    weights: &[f64],
    dim: usize,
) -> EntangledHVec {
    use crate::core::entangled::DEFAULT_RHO_DENOM;

    let target = (dim / DEFAULT_RHO_DENOM).max(1);
    let mut freq: crate::core::hash::FastMap<u32, f64> =
        crate::core::hash::FastMap::with_capacity_and_hasher(target * 2, Default::default());

    for (i, (_, pat)) in patterns.iter().enumerate() {
        let w = weights[i];
        if w < 1e-12 {
            continue;
        }
        for &idx in &pat.borrow().indices {
            *freq.entry(idx).or_insert(0.0) += w;
        }
    }

    let mut scored: Vec<(u32, f64)> = freq.into_iter().collect();
    if scored.len() > target {
        scored.select_nth_unstable_by(target - 1, |a, b| {
            b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
        });
        scored.truncate(target);
    }

    let mut indices: Vec<u32> = scored.into_iter().map(|(idx, _)| idx).collect();
    indices.sort_unstable();
    EntangledHVec::from_indices(indices, dim)
}

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

    fn make_pattern(dim: usize, seed: u64) -> EntangledHVec {
        EntangledHVec::new_deterministic(dim, seed)
    }

    type Transform = fn(&[f64]) -> Vec<f64>;
    const TRANSFORMS: [(&str, Transform); 3] = [
        ("softmax", softmax),
        ("entmax_15", entmax_15),
        ("sparsemax", sparsemax),
    ];

    #[test]
    fn entmax_family_outputs_probability_distributions() {
        let inputs: [&[f64]; 4] = [
            &[5.0],
            &[3.0, 3.0, 3.0],
            &[10.0, 1.0, 1.0, 1.0],
            &[5.0, 3.0, 1.0, 0.5, 0.1],
        ];
        for (name, transform) in TRANSFORMS {
            for scores in inputs {
                let weights = transform(scores);
                assert_eq!(weights.len(), scores.len(), "{name} {scores:?}");
                assert!(weights.iter().all(|&w| w >= 0.0), "{name} {scores:?}");
                let sum: f64 = weights.iter().sum();
                assert!((sum - 1.0).abs() < 1e-6, "{name} {scores:?} sums to {sum}");
                if scores.windows(2).all(|pair| pair[0] == pair[1]) {
                    let uniform = 1.0 / scores.len() as f64;
                    assert!(
                        weights.iter().all(|&w| (w - uniform).abs() < 1e-6),
                        "{name} {scores:?} is not uniform: {weights:?}"
                    );
                }
            }
        }
    }

    #[test]
    fn only_softmax_keeps_dominated_scores_nonzero() {
        let scores = [10.0, 1.0, 1.0, 1.0];
        for ((name, transform), expected_support) in TRANSFORMS.into_iter().zip([4, 1, 1]) {
            let support = transform(&scores).iter().filter(|&&w| w > 1e-10).count();
            assert_eq!(support, expected_support, "{name}");
        }
    }

    #[test]
    fn test_hopfield_exact_retrieval() {
        let dim = 16384;
        let p1 = make_pattern(dim, 100);
        let p2 = make_pattern(dim, 200);
        let p3 = make_pattern(dim, 300);

        let patterns = vec![
            ("p1".to_string(), p1.clone()),
            ("p2".to_string(), p2.clone()),
            ("p3".to_string(), p3),
        ];

        let config = HopfieldConfig {
            beta: 100.0,
            alpha: 2.0,
            max_iter: 1,
        };

        let results = hopfield_query(&p1, &patterns, &config, 10);
        assert!(!results.is_empty(), "Should retrieve at least one pattern");
        assert_eq!(results[0].id, "p1", "Should retrieve the exact pattern");

        let results = hopfield_query(&p2, &patterns, &config, 10);
        assert!(!results.is_empty());
        assert_eq!(results[0].id, "p2");
    }

    #[test]
    fn test_hopfield_sparse_output() {
        let dim = 16384;
        let patterns: Vec<(String, EntangledHVec)> = (0..20)
            .map(|i| (format!("p{}", i), make_pattern(dim, i * 100 + 1)))
            .collect();

        let config = HopfieldConfig {
            beta: 100.0,
            alpha: 2.0,
            max_iter: 1,
        };

        let query = patterns[5].1.clone();
        let results = hopfield_query(&query, &patterns, &config, 20);

        assert!(
            results.len() < 10,
            "Sparsemax should produce sparse output, got {} results",
            results.len()
        );
    }

    #[test]
    fn test_hopfield_iterative_convergence() {
        let dim = 16384;
        let p1 = make_pattern(dim, 100);
        let p2 = make_pattern(dim, 200);

        let noise = make_pattern(dim, 999);
        let noisy_query = p1.bind(&noise);

        let patterns = vec![("p1".to_string(), p1.clone()), ("p2".to_string(), p2)];

        let config = HopfieldConfig {
            beta: 100.0,
            alpha: 2.0,
            max_iter: 3,
        };

        let results = hopfield_query_iterative(&noisy_query, &patterns, &config, 10);
        assert_eq!(results[0].id, "p1");
    }
}