consciousness/poc-memory/src/search.rs

// Memory search: composable algorithm pipeline.
//
// Each algorithm is a stage: takes seeds Vec<(String, f64)>, returns
// new/modified seeds. Stages compose left-to-right in a pipeline.
//
// Available algorithms:
//   spread    — spreading activation through graph edges
//   spectral  — nearest neighbors in spectral embedding space
//   manifold  — extrapolation along direction defined by seeds (TODO)
//
// Seed extraction (matching query terms to node keys) is shared
// infrastructure, not an algorithm stage.

use crate::store::StoreView;
use crate::graph::Graph;
use crate::spectral;

use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::fmt;

pub struct SearchResult {
    pub key: String,
    pub activation: f64,
    pub is_direct: bool,
    pub snippet: Option<String>,
}

/// A parsed algorithm stage with its parameters.
#[derive(Clone, Debug)]
pub struct AlgoStage {
    pub algo: Algorithm,
    pub params: HashMap<String, String>,
}

#[derive(Clone, Debug)]
pub enum Algorithm {
    Spread,
    Spectral,
    Manifold,
}

impl fmt::Display for Algorithm {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Algorithm::Spread => write!(f, "spread"),
            Algorithm::Spectral => write!(f, "spectral"),
            Algorithm::Manifold => write!(f, "manifold"),
        }
    }
}

impl AlgoStage {
    /// Parse "spread,max_hops=4,edge_decay=0.5" into an AlgoStage.
    pub fn parse(s: &str) -> Result<Self, String> {
        let mut parts = s.split(',');
        let name = parts.next().unwrap_or("");
        let algo = match name {
            "spread" => Algorithm::Spread,
            "spectral" => Algorithm::Spectral,
            "manifold" => Algorithm::Manifold,
            _ => return Err(format!("unknown algorithm: {}", name)),
        };
        let mut params = HashMap::new();
        for part in parts {
            if let Some((k, v)) = part.split_once('=') {
                params.insert(k.to_string(), v.to_string());
            } else {
                return Err(format!("bad param (expected key=val): {}", part));
            }
        }
        Ok(AlgoStage { algo, params })
    }

    fn param_f64(&self, key: &str, default: f64) -> f64 {
        self.params.get(key)
            .and_then(|v| v.parse().ok())
            .unwrap_or(default)
    }

    fn param_u32(&self, key: &str, default: u32) -> u32 {
        self.params.get(key)
            .and_then(|v| v.parse().ok())
            .unwrap_or(default)
    }

    fn param_usize(&self, key: &str, default: usize) -> usize {
        self.params.get(key)
            .and_then(|v| v.parse().ok())
            .unwrap_or(default)
    }
}

/// Extract seeds from weighted terms by matching against node keys.
///
/// Returns (seeds, direct_hits) where direct_hits tracks which keys
/// were matched directly (vs found by an algorithm stage).
pub fn match_seeds(
    terms: &BTreeMap<String, f64>,
    store: &impl StoreView,
) -> (Vec<(String, f64)>, HashSet<String>) {
    let mut seeds: Vec<(String, f64)> = Vec::new();
    let mut direct_hits: HashSet<String> = HashSet::new();

    let mut key_map: HashMap<String, (String, f64)> = HashMap::new();
    store.for_each_node(|key, _content, weight| {
        key_map.insert(key.to_lowercase(), (key.to_owned(), weight as f64));
    });

    for (term, &term_weight) in terms {
        if let Some((orig_key, node_weight)) = key_map.get(term) {
            let score = term_weight * node_weight;
            seeds.push((orig_key.clone(), score));
            direct_hits.insert(orig_key.clone());
        }
    }

    (seeds, direct_hits)
}

/// Run a pipeline of algorithm stages.
pub fn run_pipeline(
    stages: &[AlgoStage],
    seeds: Vec<(String, f64)>,
    graph: &Graph,
    store: &impl StoreView,
    debug: bool,
    max_results: usize,
) -> Vec<(String, f64)> {
    let mut current = seeds;

    for stage in stages {
        if debug {
            println!("\n[search] === {} ({} seeds in) ===", stage.algo, current.len());
        }

        current = match stage.algo {
            Algorithm::Spread => run_spread(&current, graph, store, stage, debug),
            Algorithm::Spectral => run_spectral(&current, graph, stage, debug),
            Algorithm::Manifold => {
                if debug { println!("  (manifold not yet implemented, passing through)"); }
                current
            }
        };

        if debug {
            println!("[search] {} → {} results", stage.algo, current.len());
            for (i, (key, score)) in current.iter().enumerate().take(15) {
                let cutoff = if i + 1 == max_results { " <-- cutoff" } else { "" };
                println!("  [{:.4}] {}{}", score, key, cutoff);
            }
            if current.len() > 15 {
                println!("  ... ({} more)", current.len() - 15);
            }
        }
    }

    current.truncate(max_results);
    current
}

/// Spreading activation: propagate scores through graph edges.
///
/// Tunable params: max_hops (default from store), edge_decay (default from store),
/// min_activation (default from store).
fn run_spread(
    seeds: &[(String, f64)],
    graph: &Graph,
    store: &impl StoreView,
    stage: &AlgoStage,
    _debug: bool,
) -> Vec<(String, f64)> {
    let store_params = store.params();
    let max_hops = stage.param_u32("max_hops", store_params.max_hops);
    let edge_decay = stage.param_f64("edge_decay", store_params.edge_decay);
    let min_activation = stage.param_f64("min_activation", store_params.min_activation);

    spreading_activation(seeds, graph, store, max_hops, edge_decay, min_activation)
}

/// Spectral projection: find nearest neighbors in spectral embedding space.
///
/// Tunable params: k (default 20, number of neighbors to find).
fn run_spectral(
    seeds: &[(String, f64)],
    graph: &Graph,
    stage: &AlgoStage,
    debug: bool,
) -> Vec<(String, f64)> {
    let k = stage.param_usize("k", 20);

    let emb = match spectral::load_embedding() {
        Ok(e) => e,
        Err(e) => {
            if debug { println!("  no spectral embedding: {}", e); }
            return seeds.to_vec();
        }
    };

    let weighted_seeds: Vec<(&str, f64)> = seeds.iter()
        .map(|(k, w)| (k.as_str(), *w))
        .collect();
    let projected = spectral::nearest_to_seeds_weighted(
        &emb, &weighted_seeds, Some(graph), k,
    );

    if debug {
        for (key, dist) in &projected {
            let score = 1.0 / (1.0 + dist);
            println!("  dist={:.6}  score={:.4}  {}", dist, score, key);
        }
    }

    // Merge: keep original seeds, add spectral results as new seeds
    let seed_set: HashSet<&str> = seeds.iter().map(|(k, _)| k.as_str()).collect();
    let mut result = seeds.to_vec();
    for (key, dist) in projected {
        if !seed_set.contains(key.as_str()) {
            let score = 1.0 / (1.0 + dist);
            result.push((key, score));
        }
    }
    result
}

fn spreading_activation(
    seeds: &[(String, f64)],
    graph: &Graph,
    store: &impl StoreView,
    max_hops: u32,
    edge_decay: f64,
    min_activation: f64,
) -> Vec<(String, f64)> {
    let mut activation: HashMap<String, f64> = HashMap::new();
    let mut queue: VecDeque<(String, f64, u32)> = VecDeque::new();

    for (key, act) in seeds {
        let current = activation.entry(key.clone()).or_insert(0.0);
        if *act > *current {
            *current = *act;
            queue.push_back((key.clone(), *act, 0));
        }
    }

    while let Some((key, act, depth)) = queue.pop_front() {
        if depth >= max_hops { continue; }

        for (neighbor, strength) in graph.neighbors(&key) {
            let neighbor_weight = store.node_weight(neighbor.as_str());
            let propagated = act * edge_decay * neighbor_weight * strength as f64;
            if propagated < min_activation { continue; }

            let current = activation.entry(neighbor.clone()).or_insert(0.0);
            if propagated > *current {
                *current = propagated;
                queue.push_back((neighbor.clone(), propagated, depth + 1));
            }
        }
    }

    let mut results: Vec<_> = activation.into_iter().collect();
    results.sort_by(|a, b| b.1.total_cmp(&a.1));
    results
}

/// Search with weighted terms: exact key matching + spectral projection.
///
/// Terms are matched against node keys. Matching nodes become seeds,
/// scored by term_weight × node_weight. Seeds are then projected into
/// spectral space to find nearby nodes, with link weights modulating distance.
pub fn search_weighted(
    terms: &BTreeMap<String, f64>,
    store: &impl StoreView,
) -> Vec<SearchResult> {
    search_weighted_inner(terms, store, false, 5)
}

/// Like search_weighted but with debug output and configurable result count.
pub fn search_weighted_debug(
    terms: &BTreeMap<String, f64>,
    store: &impl StoreView,
    max_results: usize,
) -> Vec<SearchResult> {
    search_weighted_inner(terms, store, true, max_results)
}

fn search_weighted_inner(
    terms: &BTreeMap<String, f64>,
    store: &impl StoreView,
    debug: bool,
    max_results: usize,
) -> Vec<SearchResult> {
    let graph = crate::graph::build_graph_fast(store);
    let (seeds, direct_hits) = match_seeds(terms, store);

    if seeds.is_empty() {
        return Vec::new();
    }

    if debug {
        println!("\n[search] === SEEDS ({}) ===", seeds.len());
        let mut sorted_seeds = seeds.clone();
        sorted_seeds.sort_by(|a, b| b.1.total_cmp(&a.1));
        for (key, score) in sorted_seeds.iter().take(20) {
            println!("  {:.4}  {}", score, key);
        }
    }

    // Default pipeline: spectral → spread (legacy behavior)
    let pipeline = vec![
        AlgoStage { algo: Algorithm::Spectral, params: HashMap::new() },
        AlgoStage { algo: Algorithm::Spread, params: HashMap::new() },
    ];

    let raw_results = run_pipeline(&pipeline, seeds, &graph, store, debug, max_results);

    raw_results.into_iter()
        .take(max_results)
        .map(|(key, activation)| {
            let is_direct = direct_hits.contains(&key);
            SearchResult { key, activation, is_direct, snippet: None }
        }).collect()
}

/// Search with equal-weight terms (for interactive use).
pub fn search(query: &str, store: &impl StoreView) -> Vec<SearchResult> {
    let terms: BTreeMap<String, f64> = query.split_whitespace()
        .map(|t| (t.to_lowercase(), 1.0))
        .collect();
    search_weighted(&terms, store)
}

/// Extract meaningful search terms from natural language.
/// Strips common English stop words, returns up to max_terms words.
pub fn extract_query_terms(text: &str, max_terms: usize) -> String {
    const STOP_WORDS: &[&str] = &[
        "the", "a", "an", "is", "are", "was", "were", "do", "does", "did",
        "have", "has", "had", "will", "would", "could", "should", "can",
        "may", "might", "shall", "been", "being", "to", "of", "in", "for",
        "on", "with", "at", "by", "from", "as", "but", "or", "and", "not",
        "no", "if", "then", "than", "that", "this", "it", "its", "my",
        "your", "our", "we", "you", "i", "me", "he", "she", "they", "them",
        "what", "how", "why", "when", "where", "about", "just", "let",
        "want", "tell", "show", "think", "know", "see", "look", "make",
        "get", "go", "some", "any", "all", "very", "really", "also", "too",
        "so", "up", "out", "here", "there",
    ];

    text.to_lowercase()
        .split(|c: char| !c.is_alphanumeric())
        .filter(|w| !w.is_empty() && w.len() > 2 && !STOP_WORDS.contains(w))
        .take(max_terms)
        .collect::<Vec<_>>()
        .join(" ")
}

/// Format search results as text lines (for hook consumption).
pub fn format_results(results: &[SearchResult]) -> String {
    let mut out = String::new();
    for (i, r) in results.iter().enumerate().take(5) {
        let marker = if r.is_direct { "→" } else { " " };
        out.push_str(&format!("{}{:2}. [{:.2}/{:.2}] {}",
            marker, i + 1, r.activation, r.activation, r.key));
        out.push('\n');
        if let Some(ref snippet) = r.snippet {
            out.push_str(&format!("         {}\n", snippet));
        }
    }
    out
}