// state_memory.sy -- Neurochemical State-Dependent Retrieval
// Memory encoding and retrieval based on neurochemical state similarity.
// State-dependent memory from neuroscience: memories are recalled
// more easily when the current biochemical state matches the encoding state.
//
// Every memory is stored with:
//   - Content (the data being remembered)
//   - Full 34-head state vector at encoding time
//   - ATP/ADP/AMP ratios at encoding (metabolic state)
//   - Scar intensity (emotional weight)
//   - Timestamp
//
// Retrieval computes chemical similarity between current cascade state
// and all stored memories. Closest state matches surface first.
// ATP level gates retrieval capacity: low ATP = only strong memories.
//
// This is not embedding similarity. This is not keyword search.
// This is state-dependent retrieval through neurochemical matching.
//
// References:
//   Godden & Baddeley (1975): State-dependent memory in divers
//   Eich (1980): Mood-dependent memory
//   Roediger & McDermott (1995): Context-dependent retrieval
//   McGaugh (2000): Memory consolidation and arousal
//
// Patent Pending: U.S. Provisional Application Nos. 63/939,190, 63/962,385 & 63/988,485
// Copyright 2025-2026 Nexus Concordat Inc.
//
// Aislinn McCubbins, April 2026
// Ported from Marjorie McCubbins' Python NCA

use cascade::Cascade
use energy::ATPCycling

// === CHEMICAL TAG ===
// A named neurochemical measurement at time of encoding.

struct ChemicalTag {
    name: String,
    value: Float,
}

// === ATP SNAPSHOT ===
// Metabolic state frozen at encoding time.

struct ATPSnapshot {
    atp: Float,
    adp: Float,
    amp: Float,
}

impl ATPSnapshot {
    fn new(atp: Float, adp: Float, amp: Float) -> ATPSnapshot {
        ATPSnapshot { atp: atp, adp: adp, amp: amp }
    }

    fn from_cascade(cascade: Cascade) -> ATPSnapshot {
        // Average ATP/ADP/AMP across all T2 heads
        let mut atp_sum = 0.0
        let mut adp_sum = 0.0
        let mut amp_sum = 0.0
        let n = cascade.tier_2.heads.len()
        for head in cascade.tier_2.heads {
            atp_sum = atp_sum + head.atp()
            adp_sum = adp_sum + head.adp()
            amp_sum = amp_sum + head.amp()
        }
        ATPSnapshot {
            atp: atp_sum / n,
            adp: adp_sum / n,
            amp: amp_sum / n,
        }
    }
}

// === STATE MEMORY ===
// A single memory tagged with neurochemical state.

temporal struct StateMemory {
    state access_count: Int = 0,

    const content: String,
    const state_vector: Vec<Float>,
    const atp_state: ATPSnapshot,
    const scar_intensity: Float,
    const encoding_strength: Float,
    const timestamp: Float,
    const tags: Vec<ChemicalTag>,
}

impl Evolve for StateMemory {
    fn evolve(&mut self, dt: Duration) {
        // Memories don't decay on their own in this store.
        // Decay is handled externally by the decay daemon.
        // Access count is the only mutable field, updated on retrieval.
    }
}

impl StateMemory {
    fn strength(&self) -> Float {
        // Effective strength: max of encoding strength and scar intensity
        let s = self.encoding_strength
        if self.scar_intensity > s {
            s = self.scar_intensity
        }
        s
    }

    fn is_scar(&self) -> Bool {
        self.scar_intensity > 0.5
    }

    fn eviction_score(&self) -> Float {
        // Score used to decide which memory to evict.
        // Encoding strength + access count bonus.
        // Higher score = harder to evict.
        self.encoding_strength + (self.access_count * 0.1)
    }

    fn tag_value(&self, name: String) -> Float {
        // Return the value of a named chemical tag, or 0.0 if not found.
        for tag in self.tags {
            if tag.name == name {
                return tag.value
            }
        }
        0.0
    }

    fn has_tag(&self, name: String) -> Bool {
        for tag in self.tags {
            if tag.name == name {
                return true
            }
        }
        false
    }
}

// === RETRIEVAL RESULT ===
// A memory paired with its weighted similarity score.

struct RetrievalResult {
    memory: StateMemory,
    similarity: Float,
}

// === MEMORY STATS ===

struct MemoryStats {
    count: Int,
    scar_count: Int,
    avg_encoding_strength: Float,
    min_encoding_strength: Float,
    max_encoding_strength: Float,
    avg_atp_at_encoding: Float,
    total_accesses: Int,
}

// === COSINE SIMILARITY ===
// Computes cosine similarity between two 34D state vectors.
// Returns value in [-1.0, 1.0]. Identical states = 1.0.

fn cosine_similarity(a: Vec<Float>, b: Vec<Float>) -> Float {
    let mut dot = 0.0
    let mut norm_a = 0.0
    let mut norm_b = 0.0
    let n = a.len()
    if n != b.len() { return 0.0 }

    for i in 0..n {
        dot = dot + a[i] * b[i]
        norm_a = norm_a + a[i] * a[i]
        norm_b = norm_b + b[i] * b[i]
    }

    let denom = sqrt(norm_a) * sqrt(norm_b) + 0.00000001
    dot / denom
}

// === ENCODING STRENGTH ===
// Encoding strength depends on available ATP and stimulus intensity.
//
// High ATP + high intensity = strong encoding (vivid memory)
// Low ATP + any intensity = weak encoding (fatigue impairs memory)
// High ATP + low intensity = moderate encoding (routine event)
//
// Reference: McGaugh (2000) Science

fn compute_encoding_strength(atp_avg: Float, stimulus_intensity: Float) -> Float {
    let raw = atp_avg * (0.5 + stimulus_intensity)
    if raw > 1.0 { return 1.0 }
    raw
}

// === STATE TAGGED MEMORY STORE ===
// Neurochemical state-dependent memory storage and retrieval.
//
// Encodes memories with the full cascade state vector and ATP ratios.
// Retrieves by chemical similarity, gated by current ATP level.

temporal struct StateTaggedMemoryStore {
    state memories: Vec<StateMemory>,
    const max_memories: Int,
}

impl Evolve for StateTaggedMemoryStore {
    fn evolve(&mut self, dt: Duration) {
        // Evolve all contained memories
        for mem in self.memories {
            mem.evolve(dt)
        }
    }
}

impl StateTaggedMemoryStore {
    fn new(max_memories: Int) -> StateTaggedMemoryStore {
        StateTaggedMemoryStore {
            memories: [],
            max_memories: max_memories,
        }
    }

    fn default() -> StateTaggedMemoryStore {
        StateTaggedMemoryStore::new(10000)
    }

    // === ENCODE ===
    // Encode a memory tagged with current neurochemical state.
    //
    // Args:
    //   content:             What to remember
    //   cascade:             CascadingHierarchy (reads current state)
    //   stimulus_intensity:  How intense the stimulus was (0.0-1.0)
    //   scar_intensity:      Emotional scarring weight (0.0 = normal, 1.0 = permanent)
    //   extra_tags:          Optional named chemical tags for this memory

    fn encode(&mut self,
              content: String,
              cascade: Cascade,
              stimulus_intensity: Float,
              scar_intensity: Float,
              extra_tags: Vec<ChemicalTag>) -> StateMemory {

        let state_vec = cascade.state_vector()
        let atp_state = ATPSnapshot::from_cascade(cascade)
        let strength = compute_encoding_strength(atp_state.atp, stimulus_intensity)

        // Build chemical tags from current NT state
        let mut tags: Vec<ChemicalTag> = []
        let nt_names = ["serotonin", "dopamine", "norepinephrine", "gaba",
                        "glutamate", "acetylcholine", "endorphin", "oxytocin"]
        for i in 0..nt_names.len() {
            let val = cascade.read(nt_names[i])
            tags.push(ChemicalTag { name: nt_names[i], value: val })
        }

        // Append any extra tags
        for tag in extra_tags {
            tags.push(tag)
        }

        let memory = StateMemory {
            content: content,
            state_vector: state_vec,
            atp_state: atp_state,
            scar_intensity: scar_intensity,
            encoding_strength: strength,
            timestamp: now(),
            access_count: 0,
            tags: tags,
        }

        self.memories.push(memory)

        // Evict weakest memory if over capacity
        if self.memories.len() > self.max_memories {
            self.evict_weakest()
        }

        memory
    }

    // === RETRIEVE ===
    // Retrieve memories by neurochemical state similarity.
    //
    // Current cascade state is compared to all stored memories.
    // Similarity is cosine similarity of state vectors, weighted by:
    //   - Encoding strength (strong memories are easier to find)
    //   - Scar intensity (scarred memories always surface)
    //   - Current ATP (low ATP = only strong/scarred memories retrieved)

    fn retrieve(&mut self,
                cascade: Cascade,
                top_k: Int,
                min_similarity: Float) -> Vec<RetrievalResult> {

        if self.memories.len() == 0 {
            return []
        }

        let current_state = cascade.state_vector()
        let current_atp = ATPSnapshot::from_cascade(cascade)

        // ATP gates retrieval capacity
        // High ATP: can retrieve anything
        // Low ATP: only strong/scarred memories surface
        // AMP emergency: only scars
        let retrieval_capacity = current_atp.atp

        let mut results: Vec<RetrievalResult> = []

        for mem in self.memories {
            // Cosine similarity of state vectors
            let sim = cosine_similarity(current_state, mem.state_vector)

            // Weight by encoding strength
            let mut weighted_sim = sim * mem.encoding_strength

            // Scar bonus: scarred memories always surface regardless of state
            let scar_bonus = mem.scar_intensity * 0.5
            weighted_sim = weighted_sim + scar_bonus

            // ATP retrieval gate: weak memories require high ATP to retrieve
            // Strong memories (high encoding_strength or scar) bypass this
            let memory_strength = mem.strength()

            if memory_strength < 0.5 && retrieval_capacity < 0.3 {
                // Fatigued brain can't retrieve weak memories
                weighted_sim = weighted_sim * retrieval_capacity
            } else if memory_strength < 0.3 && retrieval_capacity < 0.5 {
                // Moderate fatigue suppresses very weak memories
                weighted_sim = weighted_sim * (retrieval_capacity + 0.5) / 1.5
            }

            // Reconsolidation: accessed memories get slightly stronger
            if weighted_sim > min_similarity {
                mem.access_count = mem.access_count + 1
            }

            if weighted_sim >= min_similarity {
                results.push(RetrievalResult {
                    memory: mem,
                    similarity: weighted_sim,
                })
            }
        }

        // Sort by weighted similarity (highest first)
        results.sort_by_desc(|r| r.similarity)

        // Return top_k
        if results.len() > top_k {
            results.truncate(top_k)
        }
        results
    }

    // === RETRIEVE BY TAG ===
    // Retrieve memories that have a specific chemical tag above threshold.
    //
    // E.g., retrieve all memories encoded during high cortisol:
    //   store.retrieve_by_tag("cortisol", cascade, 10)

    fn retrieve_by_tag(&self,
                       tag_name: String,
                       cascade: Cascade,
                       top_k: Int) -> Vec<RetrievalResult> {

        let current_atp = ATPSnapshot::from_cascade(cascade)
        let retrieval_capacity = current_atp.atp

        let mut results: Vec<RetrievalResult> = []

        for mem in self.memories {
            if mem.has_tag(tag_name) {
                let tag_val = mem.tag_value(tag_name)
                let strength = mem.strength()

                // ATP gating on retrieval
                let mut effective_strength = tag_val * strength
                if retrieval_capacity < 0.3 && strength < 0.5 {
                    effective_strength = effective_strength * retrieval_capacity
                }

                results.push(RetrievalResult {
                    memory: mem,
                    similarity: effective_strength,
                })
            }
        }

        results.sort_by_desc(|r| r.similarity)
        if results.len() > top_k {
            results.truncate(top_k)
        }
        results
    }

    // === RETRIEVE BY ATP STATE ===
    // Retrieve memories encoded during specific metabolic states.
    //
    // target_atp: "energized" (ATP > 0.7), "normal" (0.3-0.7),
    //             "fatigued" (ATP < 0.3), "crisis" (AMP > 0.3)

    fn retrieve_by_atp_state(&self,
                             target_atp: String,
                             top_k: Int) -> Vec<RetrievalResult> {

        let mut results: Vec<RetrievalResult> = []

        for mem in self.memories {
            let atp = mem.atp_state.atp
            let adp = mem.atp_state.adp
            let amp = mem.atp_state.amp
            let mut match_score = 0.0

            if target_atp == "energized" && atp > 0.7 {
                match_score = atp
            } else if target_atp == "normal" && atp >= 0.3 && atp <= 0.7 {
                // Closer to 0.5 = better match
                let dev = atp - 0.5
                if dev < 0.0 { dev = 0.0 - dev }
                match_score = 1.0 - dev * 2.0
            } else if target_atp == "fatigued" && atp < 0.3 {
                match_score = 1.0 - atp
            } else if target_atp == "crisis" && amp > 0.3 {
                match_score = amp
            }

            if match_score > 0.0 {
                results.push(RetrievalResult {
                    memory: mem,
                    similarity: match_score * mem.encoding_strength,
                })
            }
        }

        results.sort_by_desc(|r| r.similarity)
        if results.len() > top_k {
            results.truncate(top_k)
        }
        results
    }

    // === EVICTION ===
    // Remove the weakest non-scarred memory when over capacity.
    // Scars are never evicted. Weakest = lowest eviction_score.

    fn evict_weakest(&mut self) {
        let mut weakest_idx = -1
        let mut weakest_score = 999999.0

        for i in 0..self.memories.len() {
            let mem = self.memories[i]

            // Never evict scars
            if mem.is_scar() {
                continue
            }

            let score = mem.eviction_score()
            if score < weakest_score {
                weakest_score = score
                weakest_idx = i
            }
        }

        if weakest_idx >= 0 {
            self.memories.remove(weakest_idx)
        }
    }

    // === STATISTICS ===

    fn stats(&self) -> MemoryStats {
        if self.memories.len() == 0 {
            return MemoryStats {
                count: 0,
                scar_count: 0,
                avg_encoding_strength: 0.0,
                min_encoding_strength: 0.0,
                max_encoding_strength: 0.0,
                avg_atp_at_encoding: 0.0,
                total_accesses: 0,
            }
        }

        let mut scar_count = 0
        let mut strength_sum = 0.0
        let mut min_strength = 999999.0
        let mut max_strength = -1.0
        let mut atp_sum = 0.0
        let mut access_sum = 0
        let n = self.memories.len()

        for mem in self.memories {
            if mem.is_scar() {
                scar_count = scar_count + 1
            }

            let s = mem.encoding_strength
            strength_sum = strength_sum + s
            if s < min_strength { min_strength = s }
            if s > max_strength { max_strength = s }

            atp_sum = atp_sum + mem.atp_state.atp
            access_sum = access_sum + mem.access_count
        }

        MemoryStats {
            count: n,
            scar_count: scar_count,
            avg_encoding_strength: strength_sum / n,
            min_encoding_strength: min_strength,
            max_encoding_strength: max_strength,
            avg_atp_at_encoding: atp_sum / n,
            total_accesses: access_sum,
        }
    }

    fn count(&self) -> Int {
        self.memories.len()
    }

    fn scar_count(&self) -> Int {
        let mut n = 0
        for mem in self.memories {
            if mem.is_scar() {
                n = n + 1
            }
        }
        n
    }
}
