holographic-memory 0.6.4

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

use std::fmt::Write as _;

use super::HmsCore;
use crate::core::entangled::EntangledHVec;
use crate::core::hopfield;
use crate::core::types::{IndexStatus, QueryExplanation, RetrievalResult, StorageHealth};
use rayon::prelude::*;

impl HmsCore {
    /// Query the memory system for the k most similar vectors.
    pub fn query(&self, query_vec: &EntangledHVec, k: u32) -> Vec<RetrievalResult> {
        if query_vec.dim != self.dimensions || k == 0 {
            return Vec::new();
        }
        // Documents memorized with deferred encoding join the vector index on first use; a
        // failure leaves them pending (and absent from this query) and the next call retries.
        // `try_read`: a query must keep progressing while compaction or a mutation holds the
        // store for writing, so a contended lock means "no catch-up this time".
        if self
            .documents
            .try_read()
            .is_some_and(|store| store.pending_vector_count() > 0)
        {
            if let Err(error) = self.encode_pending_documents() {
                tracing::warn!(%error, "deferred document encoding failed; pending chunks stay outside this query");
            }
        }

        let _transaction = self.mutation_gate.read();
        self.shards
            .read()
            .query(query_vec, k.min(10000), self.dimensions)
    }

    /// The registry of agents admitted by verified access credentials.
    #[cfg(feature = "provenance")]
    pub fn agent_registry(
        &self,
    ) -> &parking_lot::RwLock<crate::core::provenance::access::AgentRegistry> {
        &self.agents
    }

    /// [`Self::query`], refused unless `agent_did` holds a current read grant.
    /// The caller must already have authenticated control of `agent_did`.
    #[cfg(feature = "provenance")]
    pub fn query_as(
        &self,
        agent_did: &str,
        query_vec: &EntangledHVec,
        k: u32,
    ) -> anyhow::Result<Vec<RetrievalResult>> {
        self.agents
            .read()
            .authorize(agent_did, crate::core::provenance::access::Permission::Read)?;
        Ok(self.query(query_vec, k))
    }

    pub fn explain_query(&self, query_vec: &EntangledHVec, k: u32) -> QueryExplanation {
        let mut explanation = self
            .shards
            .read()
            .explain_query(query_vec, k, self.dimensions);
        let (pending, lexical) = self
            .documents
            .try_read()
            .map(|store| (store.pending_vector_count(), store.lexical_vector_count()))
            .unwrap_or((0, 0));
        if pending > 0 || lexical > 0 {
            let _ = write!(
                explanation.rationale,
                " {pending} document chunk(s) await deferred encoding (a query, \
                 encode_pending_documents or maintain_indices installs them); {lexical} \
                 lexical-only chunk(s) are outside vector search."
            );
        }
        explanation
    }

    pub fn index_status(&self) -> IndexStatus {
        let count = self.vector_count().min(u32::MAX as u64) as u32;
        let nsg = self.nsg_trained();
        let ivf = self.ivf_trained();
        let nsg_recommended = !nsg
            && count > 0
            && self.config.nsg.auto_threshold > 0
            && count as usize >= self.config.nsg.auto_threshold;
        let ivf_recommended = self.config.ivf.enabled
            && !ivf
            && count > 0
            && self.config.ivf.auto_threshold > 0
            && count as usize >= self.config.ivf.auto_threshold;
        let recommendation = if nsg_recommended && ivf_recommended {
            "train NSG and IVF indices"
        } else if nsg_recommended {
            "train the NSG index"
        } else if ivf_recommended {
            "train the IVF index"
        } else {
            "indices are current for the configured thresholds"
        };
        IndexStatus {
            vector_count: count,
            nsg_trained: nsg,
            ivf_trained: ivf,
            nsg_training_recommended: nsg_recommended,
            ivf_training_recommended: ivf_recommended,
            recommendation: recommendation.to_string(),
        }
    }

    pub fn storage_health(&self) -> StorageHealth {
        let stats = self.arena.stats();
        StorageHealth {
            format_version: stats.format_version,
            segment_count: stats.segment_count.min(u32::MAX as usize) as u32,
            used_bytes: stats.used_bytes as f64,
            capacity_bytes: stats.capacity_bytes as f64,
        }
    }

    /// Train any index whose configured threshold has been reached, then return
    /// the resulting lifecycle state. Calls are idempotent when indices are current.
    pub fn maintain_indices(&self) -> anyhow::Result<IndexStatus> {
        self.maybe_auto_shard(self.vector_count())?;
        self.encode_pending_documents()?;
        let before = self.index_status();
        if before.ivf_training_recommended {
            self.train_ivf()?;
        }
        if before.nsg_training_recommended {
            self.train_nsg()?;
        }
        let document_training_recommended = self.documents().dense_index_recommended();
        if document_training_recommended {
            self.train_document_index()?;
        }
        self.maintain_vector_index(crate::core::documents::DEFAULT_GRAPH_MIN_ROWS)?;
        Ok(self.index_status())
    }

    /// Run maintenance on a dedicated thread; training publishes only a current snapshot.
    pub fn maintain_indices_background(
        self: &std::sync::Arc<Self>,
    ) -> std::thread::JoinHandle<anyhow::Result<IndexStatus>> {
        let core = std::sync::Arc::clone(self);
        std::thread::spawn(move || core.maintain_indices())
    }

    /// Flush persistent vector data to durable storage.
    pub fn flush(&self) -> anyhow::Result<()> {
        {
            let _transaction = self.mutation_gate.write();
            self.arena.flush()?;
            self.persist_indices(&self.shards.read())?;
        }
        self.checkpoint_if_due()
    }

    /// Energy-based associative retrieval using Hopfield-Fenchel-Young dynamics.
    ///
    /// Unlike `query` (which returns a fixed top-k by similarity), this uses
    /// sparse entmax attention to naturally determine how many results are
    /// relevant. Returns at most `max_results` patterns with non-zero
    /// Hopfield attention weight.
    pub fn query_hopfield(
        &self,
        query_vec: &EntangledHVec,
        max_results: u32,
    ) -> Vec<RetrievalResult> {
        let config = &self.config.hopfield;
        self.shards.read().with_patterns(|patterns| {
            if config.max_iter > 1 {
                hopfield::hopfield_query_iterative(
                    query_vec,
                    patterns,
                    config,
                    max_results as usize,
                )
            } else {
                hopfield::hopfield_query(query_vec, patterns, config, max_results as usize)
            }
        })
    }

    /// Process multiple queries in parallel using rayon.
    pub fn query_batch(&self, queries: &[EntangledHVec], k: u32) -> Vec<Vec<RetrievalResult>> {
        queries.par_iter().map(|q| self.query(q, k)).collect()
    }

    /// Analyze components of a vector by finding its nearest neighbors.
    /// Filters by similarity threshold from QueryConfig (default 0.05).
    pub fn analyze_components(&self, vector: &EntangledHVec) -> Vec<RetrievalResult> {
        let cfg = &self.config.query;
        let neighbors = self.query(vector, cfg.component_max_neighbors);
        neighbors
            .into_iter()
            .filter(|r| r.similarity > cfg.component_similarity_threshold)
            .collect()
    }
}

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

    #[test]
    fn diagnostics_report_route_lifecycle_and_storage() -> anyhow::Result<()> {
        let directory = tempfile::tempdir()?;
        let hms = HmsCore::new(
            1024,
            Some(directory.path().to_string_lossy().into_owned()),
            None,
        )?;
        let query = EntangledHVec::new_deterministic(1024, 1);
        let explanation = hms.explain_query(&query, 5);
        assert_eq!(explanation.route, "brute_force");
        assert_eq!(hms.index_status().vector_count, 0);
        assert_eq!(
            hms.storage_health().format_version,
            crate::core::storage::STORAGE_FORMAT_VERSION
        );
        hms.flush()?;
        Ok(())
    }
}