//! Muninn - The Memory Consolidator.
//!
//! Implements Elastic Weight Consolidation (EWC) for density matrices.

use crate::fisher::Galdr;
use mimirs_core::{DensityMemory, MemoryId, MimirError};
use std::collections::HashMap;

/// Configuration for consolidation sessions.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ConsolidationConfig {
    /// Purity threshold to trigger consolidation.
    pub collapse_threshold: f32,
    /// EWC regularization strength λ.
    pub ewc_lambda: f32,
}

impl Default for ConsolidationConfig {
    fn default() -> Self {
        Self {
            collapse_threshold: 0.3,
            ewc_lambda: 1.0,
        }
    }
}

/// Muninn - The Memory Consolidator (L3).
///
/// Responsibilities:
/// - Symbolic Abstraction: Extracting invariant laws from episodic clusters.
/// - Reasoning Compression: Distilling long thought chains into principles.
/// - Fisher-weighted mixing: Protecting structural criticality during merge.
pub struct Muninn {
    _config: ConsolidationConfig,
    /// Persistent Fisher states for durable memories.
    _durable_fisher: HashMap<MemoryId, Galdr>,
    /// Generative distillation replay buffer (L5).
    pub replay: crate::replay::GenerativeDistillationBuffer,
}

impl Muninn {
    /// Opna - Creates a new Muninn consolidator.
    pub fn new(config: ConsolidationConfig) -> Self {
        Self {
            _config: config,
            _durable_fisher: HashMap::new(),
            replay: crate::replay::GenerativeDistillationBuffer::new(1000),
        }
    }

    /// Abstrahera - Extract a symbolic abstraction from a cluster of related memories.
    /// Named after the process of drawing out the essence (Abstraction).
    pub fn abstrahera(&self, cluster: &[mimirs_core::Memory]) -> Result<mimirs_core::Memory, MimirError> {
        if cluster.is_empty() {
            return Err(MimirError::Validation("Empty cluster for abstraction".to_string()));
        }

        // Synthesize a more meaningful content label
        let mut words: std::collections::HashSet<String> = std::collections::HashSet::new();
        for m in cluster {
            for word in m.content.split_whitespace().take(10) {
                if word.len() > 4 {
                    words.insert(word.to_lowercase());
                }
            }
        }
        
        let keywords: Vec<String> = words.into_iter().take(5).collect();
        let content = if keywords.is_empty() {
            "Generalized abstraction of episodic cluster".to_string()
        } else {
            format!("Abstracted concept: {}", keywords.join(", "))
        };

        Ok(mimirs_core::Memory {
            id: mimirs_core::MemoryId::new(),
            content,
            metadata: std::collections::HashMap::new(),
            scope: mimirs_core::MemoryScope::Agent,
            verifiability: mimirs_core::VerifiabilityStage::Verified,
            memory_class: mimirs_core::MemoryClass::Semantic,
            rho: None, // Will be synthesized from cluster density overlap
            qrc_state: None,
            scramble_score: None,
        })
    }

    /// Consolidate - Mixes two memories into one using Fisher-weighted mixing.
    /// Protected directions in the Hilbert space are preserved via Galdr importance.
    pub fn consolidate(
        &self,
        base: &DensityMemory,
        new: &DensityMemory,
        importance: &Galdr,
    ) -> Result<DensityMemory, MimirError> {
        if base.dimension() != new.dimension() {
            return Err(MimirError::Validation(
                "Dimension mismatch in consolidation".to_string(),
            ));
        }

        // Fisher-weighted mixing:
        // Protect directions where importance (F) is high.
        let weights = importance.weights();
        let sum_w: f64 = weights.iter().sum();
        let n = weights.len() as f64;
        let avg_importance = if n > 0.0 { sum_w / n } else { 0.5 };

        // DERIVATION: Elastic mixing weight w = 0.5 + 0.45 * (1 - exp(-avg_F)).
        // Ensures w ∈ [0.5, 0.95], providing stability (favoring base) proportional
        // to mean Fisher importance, preventing catastrophic forgetting.
        let w = 0.5 + (0.45 * (1.0 - (-avg_importance).exp()));
        debug_assert!(w.is_finite());
        debug_assert!((0.5..=0.95).contains(&w));
        
        base.mix(new, w as f32)
    }

    /// Pushes a memory into the generative replay buffer.
    pub fn push_to_replay(&mut self, memory: &DensityMemory) {
        self.replay.push(memory.clone());
    }

    /// Samples a batch from the replay buffer for distillation.
    pub fn sample_replay(&self, n: usize) -> Vec<DensityMemory> {
        self.replay.sample_batch(n)
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use mimirs_core::{DensityMemory, ScrambleScore};

    #[test]
    fn test_muninn_consolidation() {
        let muninn = Muninn::new(ConsolidationConfig::default());
        let base = DensityMemory::from_pure(&[1.0, 0.0]).unwrap();
        let new = DensityMemory::from_pure(&[0.0, 1.0]).unwrap();
        let importance = Galdr::estimate(&base, ScrambleScore::new(0.5));
        
        let result = muninn.consolidate(&base, &new, &importance).unwrap();
        assert_eq!(result.dimension(), 2);
        assert!(result.purity() < 1.0); // Mixed state should have lower purity
    }
}
