// neuromodulated_loss.sy -- Chemistry-Driven Loss System
// The brain does not use a single global loss. It uses distributed local
// prediction errors, weighted dynamically by neurochemistry.
//
// Seven biological mechanisms:
//   1. NeuromodulatedLossWeighter: Chemistry-weighted dynamic loss
//   2. LocalPredictionLoss: Per-tier local prediction errors (predictive coding)
//   3. EligibilityTraceBuffer: Temporal credit assignment
//   4. HomeostaticScaler: Synaptic scaling
//   5. ConsolidationBuffer: Sleep replay analogue
//   6. HebbianCoActivationLoss: Heads that fire together wire together
//   7. NeuromodulatedLoss: Composite loss combining all components
//
// NO SIGMOID. ANYWHERE. EVER.
// All gates use H-H exponential rate functions: alpha/(alpha+beta).
//
// References:
//   Three-factor rule: Gerstner et al. (2018) Frontiers Comp. Neurosci.
//   Predictive coding: Rao & Ballard (1999); Millidge et al. (2025)
//   Neuromodulatory learning: Yu & Dayan (2005); Doya (2002)
//   Homeostatic plasticity: Turrigiano (2008) Cell
//   Sleep consolidation: Diekelmann & Born (2010) Nature Reviews Neurosci.
//   BCM sliding threshold: Bienenstock, Cooper & Munro (1982)
//
// Ported from Marjorie McCubbins' Python NCA by Aislinn McCubbins, April 2026
// Patent Pending: U.S. Provisional Application Nos. 63/939,190, 63/962,385 & 63/988,485
// Copyright 2025-2026 Nexus Concordat Inc.

use cascade::Cascade

// ============================================================
// TIER 2 NT INDICES in 34D state vector
// ============================================================
// Tier 0: 8 amino acids (indices 0-7)
// Tier 1: 6 precursors (indices 8-13)
// Tier 2: 8 NTs (indices 14-21)
// Tier 3: 12 brain regions (indices 22-33)

const NT_OFFSET: Int = 14

// Named indices within Tier 2 (0-indexed from NT_OFFSET)
const DOPAMINE: Int = 0
const NOREPINEPHRINE: Int = 1
const SEROTONIN: Int = 2
const GABA_IDX: Int = 3
const OXYTOCIN: Int = 4
const CORTISOL: Int = 5
const ENDORPHIN: Int = 6
const ACETYLCHOLINE: Int = 7


// ============================================================
// LOSS WEIGHTS -- Dynamic weights computed from chemical state
// ============================================================

struct LossWeights {
    lm_weight: Float,
    state_weight: Float,
    local_pred_weight: Float,
    hebbian_weight: Float,
    lr_modifier: Float,
    trace_decay: Float,
    cortisol_phase: String,
}

impl LossWeights {
    fn default() -> LossWeights {
        LossWeights {
            lm_weight: 1.0,
            state_weight: 0.3,
            local_pred_weight: 0.2,
            hebbian_weight: 0.1,
            lr_modifier: 1.0,
            trace_decay: 0.95,
            cortisol_phase: "acute",
        }
    }
}


// ============================================================
// H-H RATE PARAMETERS -- Learnable gating for one modulator
// ============================================================

struct HHRateParams {
    w_alpha: Float,
    b_alpha: Float,
    w_beta: Float,
    b_beta: Float,
}

impl HHRateParams {
    fn new(w_a: Float, b_a: Float, w_b: Float, b_b: Float) -> HHRateParams {
        HHRateParams {
            w_alpha: w_a,
            b_alpha: b_a,
            w_beta: w_b,
            b_beta: b_b,
        }
    }

    fn gate(&self, x: Float) -> Float {
        // H-H rate function: maps scalar to [0, 1]. No sigmoid. No tanh.
        // alpha = exp(clamp(w_alpha * x + b_alpha, -3, 3))
        // beta  = exp(clamp(w_beta  * x + b_beta,  -3, 3))
        // gate  = alpha / (alpha + beta + eps)
        let mut raw_alpha = self.w_alpha * x + self.b_alpha
        if raw_alpha < -3.0 { raw_alpha = -3.0 }
        if raw_alpha > 3.0 { raw_alpha = 3.0 }

        let mut raw_beta = self.w_beta * x + self.b_beta
        if raw_beta < -3.0 { raw_beta = -3.0 }
        if raw_beta > 3.0 { raw_beta = 3.0 }

        let alpha = exp(raw_alpha)
        let beta = exp(raw_beta)
        alpha / (alpha + beta + 0.00000001)
    }
}


// ============================================================
// NEUROMODULATED LOSS WEIGHTER
// ============================================================
// The neuromodulators ARE the loss function weights.
// When dopamine is high, reward learning dominates.
// When ACh is high, the effective learning rate increases.
// When cortisol is chronically elevated, learning degrades.
//
// Three-factor rule: delta_w = Hebbian * trace * modulator

temporal struct NeuromodulatedLossWeighter {
    state cortisol_ema: Float = 0.3,

    const base_lm_weight: Float = 1.0,
    const base_state_weight: Float = 0.3,
    const base_local_pred_weight: Float = 0.2,
    const base_hebbian_weight: Float = 0.1,
    const cortisol_chronic_threshold: Float = 0.7,
    const modulation_strength: Float = 0.5,
    const cortisol_ema_alpha: Float = 0.05,

    // H-H rate parameters for each neurochemical modulator
    const da_rates: HHRateParams,
    const ach_rates: HHRateParams,
    const ne_rates: HHRateParams,
    const serotonin_rates: HHRateParams,
    const cortisol_rates: HHRateParams,
    const gaba_rates: HHRateParams,
}

impl NeuromodulatedLossWeighter {
    fn new() -> NeuromodulatedLossWeighter {
        // Initialized so gate ~ 0.5 at resting state (x ~ 0.5)
        NeuromodulatedLossWeighter {
            cortisol_ema: 0.3,
            base_lm_weight: 1.0,
            base_state_weight: 0.3,
            base_local_pred_weight: 0.2,
            base_hebbian_weight: 0.1,
            cortisol_chronic_threshold: 0.7,
            modulation_strength: 0.5,
            cortisol_ema_alpha: 0.05,
            da_rates: HHRateParams::new(0.3, -0.15, -0.3, 0.15),
            ach_rates: HHRateParams::new(0.2, -0.1, -0.2, 0.1),
            ne_rates: HHRateParams::new(0.3, -0.15, -0.3, 0.15),
            serotonin_rates: HHRateParams::new(0.2, -0.1, -0.2, 0.1),
            cortisol_rates: HHRateParams::new(0.2, -0.1, -0.2, 0.1),
            gaba_rates: HHRateParams::new(0.15, -0.075, -0.15, 0.075),
        }
    }

    fn compute(&mut self, state_vector: Vec<Float>) -> LossWeights {
        // Extract Tier 2 NT values from 34D state vector
        let da = state_vector[NT_OFFSET + DOPAMINE]
        let ne = state_vector[NT_OFFSET + NOREPINEPHRINE]
        let serotonin = state_vector[NT_OFFSET + SEROTONIN]
        let gaba = state_vector[NT_OFFSET + GABA_IDX]
        let cortisol = state_vector[NT_OFFSET + CORTISOL]
        let ach = state_vector[NT_OFFSET + ACETYLCHOLINE]

        // Update cortisol EMA for chronic detection (slow tracking)
        self.cortisol_ema = self.cortisol_ema * (1.0 - self.cortisol_ema_alpha)
            + cortisol * self.cortisol_ema_alpha

        // Determine cortisol phase
        let is_chronic = self.cortisol_ema > self.cortisol_chronic_threshold
        let mut cortisol_phase = "acute"
        if is_chronic { cortisol_phase = "chronic" }

        let ms = self.modulation_strength

        // --- DOPAMINE: Reward signal -> weight LM loss ---
        // High DA = reward prediction error is large = learn harder
        let da_mod = self.da_rates.gate(da)
        let mut lm_weight = self.base_lm_weight * (1.0 + ms * (da_mod - 0.5))

        // --- ACh: Expected uncertainty -> learning rate modifier ---
        // High ACh = "I don't know what's happening" = learn faster
        let ach_mod = self.ach_rates.gate(ach)
        let mut lr_modifier = 1.0 + ms * (ach_mod - 0.5)

        // --- NE: Salience -> state alignment loss ---
        // High NE = "this is important" = stamp the chemical state harder
        let ne_mod = self.ne_rates.gate(ne)
        let mut state_weight = self.base_state_weight * (1.0 + ms * (ne_mod - 0.5))

        // --- SEROTONIN: Temporal patience -> eligibility trace decay ---
        // High 5-HT = patient = traces last longer (decay closer to 1.0)
        let sero_mod = self.serotonin_rates.gate(serotonin)
        let trace_decay = 0.9 + 0.09 * sero_mod   // range [0.9, 0.99]

        // --- CORTISOL: Dual-phase stress modulator ---
        let cortisol_mod = self.cortisol_rates.gate(cortisol)
        if is_chronic {
            // Chronic stress: degrade learning (reduce LR, reduce all weights)
            let mut degradation = 1.0 - ms * cortisol_mod
            if degradation < 0.3 { degradation = 0.3 }
            lr_modifier = lr_modifier * degradation
        } else {
            // Acute stress: enhance state memory (amygdala LTP)
            state_weight = state_weight * (1.0 + ms * 0.5 * cortisol_mod)
        }

        // --- GABA: Inhibition -> reduce overall plasticity ---
        // High GABA = "calm down" = less plastic
        let gaba_mod = self.gaba_rates.gate(gaba)
        let mut gaba_factor = 1.0 - ms * 0.3 * gaba_mod
        if gaba_factor < 0.5 { gaba_factor = 0.5 }
        lr_modifier = lr_modifier * gaba_factor

        // Clamp final values to safe ranges
        if lm_weight < 0.1 { lm_weight = 0.1 }
        if state_weight < 0.05 { state_weight = 0.05 }
        if lr_modifier < 0.1 { lr_modifier = 0.1 }
        if lr_modifier > 3.0 { lr_modifier = 3.0 }

        LossWeights {
            lm_weight: lm_weight,
            state_weight: state_weight,
            local_pred_weight: self.base_local_pred_weight,
            hebbian_weight: self.base_hebbian_weight,
            lr_modifier: lr_modifier,
            trace_decay: trace_decay,
            cortisol_phase: cortisol_phase,
        }
    }
}

impl Evolve for NeuromodulatedLossWeighter {
    fn evolve(&mut self, dt: Duration) {
        // Cortisol EMA decays slowly toward resting (0.3) between updates
        self.cortisol_ema = self.cortisol_ema * 0.999 + 0.3 * 0.001
    }
}


// ============================================================
// LOCAL PREDICTION LOSS -- Predictive Coding per Tier
// ============================================================
// Each tier predicts the next tier's state vector.
//
// T0 (8 amino acids) predicts T1 (6 precursors)
// T1 (6 precursors) predicts T2 (8 NTs)
// T2 (8 NTs) predicts T3 (12 brain regions)
//
// These are SMALL predictors. Linear projection + bias.
// The loss is MSE between predicted and actual state.
// This is the somato-dendritic mismatch.
//
// NO SIGMOID in prediction heads. Linear output, real-valued states.

struct TierPredictor {
    weights: Vec<Float>,
    bias: Vec<Float>,
    in_dim: Int,
    out_dim: Int,
}

impl TierPredictor {
    fn new(in_dim: Int, out_dim: Int) -> TierPredictor {
        // Initialize near-zero so predictions start weak
        let mut weights = []
        let mut bias = []
        for i in 0..(in_dim * out_dim) {
            weights = weights + [0.01]
        }
        for i in 0..out_dim {
            bias = bias + [0.0]
        }
        TierPredictor {
            weights: weights,
            bias: bias,
            in_dim: in_dim,
            out_dim: out_dim,
        }
    }

    fn predict(&self, input: Vec<Float>) -> Vec<Float> {
        // Linear projection: output[j] = sum_i(input[i] * weights[i * out_dim + j]) + bias[j]
        let mut result = []
        for j in 0..self.out_dim {
            let mut sum = self.bias[j]
            for i in 0..self.in_dim {
                sum = sum + input[i] * self.weights[i * self.out_dim + j]
            }
            result = result + [sum]
        }
        result
    }
}

struct LocalPredictionLoss {
    t0_to_t1: TierPredictor,
    t1_to_t2: TierPredictor,
    t2_to_t3: TierPredictor,
}

impl LocalPredictionLoss {
    fn new() -> LocalPredictionLoss {
        LocalPredictionLoss {
            t0_to_t1: TierPredictor::new(8, 6),
            t1_to_t2: TierPredictor::new(6, 8),
            t2_to_t3: TierPredictor::new(8, 12),
        }
    }

    fn compute(&self, state_vector: Vec<Float>) -> LocalPredictionDetails {
        // Slice state vector into tiers
        let mut t0 = []
        for i in 0..8 { t0 = t0 + [state_vector[i]] }

        let mut t1 = []
        for i in 8..14 { t1 = t1 + [state_vector[i]] }

        let mut t2 = []
        for i in 14..22 { t2 = t2 + [state_vector[i]] }

        let mut t3 = []
        for i in 22..34 { t3 = t3 + [state_vector[i]] }

        // T0 predicts T1
        let t1_pred = self.t0_to_t1.predict(t0)
        let t0_t1_loss = mse(t1_pred, t1)

        // T1 predicts T2
        let t2_pred = self.t1_to_t2.predict(t1)
        let t1_t2_loss = mse(t2_pred, t2)

        // T2 predicts T3
        let t3_pred = self.t2_to_t3.predict(t2)
        let t2_t3_loss = mse(t3_pred, t3)

        let total_loss = t0_t1_loss + t1_t2_loss + t2_t3_loss

        LocalPredictionDetails {
            t0_t1_loss: t0_t1_loss,
            t1_t2_loss: t1_t2_loss,
            t2_t3_loss: t2_t3_loss,
            total_loss: total_loss,
        }
    }
}

struct LocalPredictionDetails {
    t0_t1_loss: Float,
    t1_t2_loss: Float,
    t2_t3_loss: Float,
    total_loss: Float,
}

// Helper: mean squared error between two vectors
fn mse(predicted: Vec<Float>, actual: Vec<Float>) -> Float {
    let mut sum = 0.0
    let mut n = 0
    for i in 0..predicted.len() {
        let diff = predicted[i] - actual[i]
        sum = sum + diff * diff
        n = n + 1
    }
    if n == 0 { return 0.0 }
    sum / n
}


// ============================================================
// ELIGIBILITY TRACE BUFFER
// ============================================================
// Rolling buffer of head activations with exponential decay.
//
// When a synapse fires, it enters an "eligible" state that persists.
// If a neuromodulatory signal (reward, punishment) arrives within
// the eligibility window, the synapse updates. No modulator =
// eligibility decays, no learning.
//
// Solves temporal credit assignment: heads active N steps ago
// when reward arrives still get credit, weighted by decayed trace.
//
// The trace_decay parameter is modulated by serotonin:
// High serotonin = patient = traces last longer.

temporal struct EligibilityTraceBuffer {
    state buffer: Vec<Vec<Float>>,
    state write_idx: Int = 0,
    state filled: Int = 0,

    const buffer_size: Int = 16,
    const state_dim: Int = 34,
    const default_decay: Float = 0.95,
}

impl EligibilityTraceBuffer {
    fn new() -> EligibilityTraceBuffer {
        // Pre-fill buffer with zero vectors
        let mut buf = []
        for i in 0..16 {
            let mut row = []
            for j in 0..34 {
                row = row + [0.0]
            }
            buf = buf + [row]
        }
        EligibilityTraceBuffer {
            buffer: buf,
            write_idx: 0,
            filled: 0,
            buffer_size: 16,
            state_dim: 34,
            default_decay: 0.95,
        }
    }

    fn with_config(size: Int, dim: Int, decay: Float) -> EligibilityTraceBuffer {
        let mut buf = []
        for i in 0..size {
            let mut row = []
            for j in 0..dim {
                row = row + [0.0]
            }
            buf = buf + [row]
        }
        EligibilityTraceBuffer {
            buffer: buf,
            write_idx: 0,
            filled: 0,
            buffer_size: size,
            state_dim: dim,
            default_decay: decay,
        }
    }

    fn record(&mut self, gate_values: Vec<Float>) {
        // Record current head gate values into the trace buffer
        self.buffer[self.write_idx] = gate_values
        self.write_idx = (self.write_idx + 1) % self.buffer_size
        self.filled = self.filled + 1
        if self.filled > self.buffer_size {
            self.filled = self.buffer_size
        }
    }

    fn get_trace_weight(&self, decay: Float) -> Vec<Float> {
        // Compute eligibility-weighted importance of each head.
        // Returns [state_dim] vector where each value is the trace-weighted
        // sum of that head's recent activations.
        if self.filled == 0 {
            let mut ones = []
            for i in 0..self.state_dim {
                ones = ones + [1.0]
            }
            return ones
        }

        let mut trace = []
        for i in 0..self.state_dim {
            trace = trace + [0.0]
        }

        let mut use_decay = decay
        if use_decay <= 0.0 { use_decay = self.default_decay }

        for step in 0..self.filled {
            // Most recent entry has highest weight
            let idx = (self.write_idx - 1 - step + self.buffer_size * 2) % self.buffer_size
            // Compute decay^step via repeated multiplication
            let mut weight = 1.0
            for p in 0..step {
                weight = weight * use_decay
            }
            let entry = self.buffer[idx]
            for d in 0..self.state_dim {
                trace[d] = trace[d] + weight * entry[d]
            }
        }

        // Normalize so max trace = 1.0
        let mut max_val = 0.0
        for i in 0..self.state_dim {
            if trace[i] > max_val { max_val = trace[i] }
        }
        if max_val > 0.0 {
            for i in 0..self.state_dim {
                trace[i] = trace[i] / max_val
            }
        }

        trace
    }

    fn get_trace_default(&self) -> Vec<Float> {
        self.get_trace_weight(self.default_decay)
    }

    fn reset(&mut self) {
        // Clear the trace buffer
        for i in 0..self.buffer_size {
            for j in 0..self.state_dim {
                self.buffer[i][j] = 0.0
            }
        }
        self.write_idx = 0
        self.filled = 0
    }
}

impl Evolve for EligibilityTraceBuffer {
    fn evolve(&mut self, dt: Duration) {
        // Trace buffer is bookkeeping, no temporal evolution needed
    }
}


// ============================================================
// HOMEOSTATIC SCALER -- Synaptic Scaling
// ============================================================
// When a neuron is too quiet for too long, ALL its synapses
// scale up multiplicatively. When too active, they scale down.
// Preserves relative weight differences while preventing
// runaway excitation or total silence.

temporal struct HomeostaticScaler {
    state gate_ema: Vec<Float>,
    state initialized: Bool = false,

    const quiet_threshold: Float = 0.1,
    const overactive_threshold: Float = 0.85,
    const scale_delta: Float = 0.01,
    const ema_alpha: Float = 0.01,
}

impl HomeostaticScaler {
    fn new() -> HomeostaticScaler {
        HomeostaticScaler {
            gate_ema: [],
            initialized: false,
            quiet_threshold: 0.1,
            overactive_threshold: 0.85,
            scale_delta: 0.01,
            ema_alpha: 0.01,
        }
    }

    fn update(&mut self, gate_values: Vec<Float>) {
        // Update exponential moving average of gate activations
        if !self.initialized {
            self.gate_ema = gate_values
            self.initialized = true
        } else {
            for i in 0..gate_values.len() {
                self.gate_ema[i] = self.gate_ema[i] * (1.0 - self.ema_alpha)
                    + gate_values[i] * self.ema_alpha
            }
        }
    }

    fn get_scaling_actions(&self) -> HomeostaticActions {
        // Determine which heads need scaling up or down
        let mut scaled_up = []
        let mut scaled_down = []

        if !self.initialized {
            return HomeostaticActions {
                scaled_up: scaled_up,
                scaled_down: scaled_down,
            }
        }

        for i in 0..self.gate_ema.len() {
            if self.gate_ema[i] < self.quiet_threshold {
                // Scale UP: increase activation gate bias
                // Makes alpha_m larger relative to beta_m,
                // raises m_inf, opening the gate more
                scaled_up = scaled_up + [i]
            } else if self.gate_ema[i] > self.overactive_threshold {
                // Scale DOWN: decrease activation gate bias
                scaled_down = scaled_down + [i]
            }
        }

        HomeostaticActions {
            scaled_up: scaled_up,
            scaled_down: scaled_down,
        }
    }

    fn get_status(&self) -> HomeostaticStatus {
        if !self.initialized {
            return HomeostaticStatus {
                min_ema: 0.0,
                max_ema: 0.0,
                mean_ema: 0.0,
                quiet_count: 0,
                overactive_count: 0,
            }
        }

        let mut min_val = 999.0
        let mut max_val = -999.0
        let mut sum = 0.0
        let mut quiet = 0
        let mut over = 0

        for i in 0..self.gate_ema.len() {
            let v = self.gate_ema[i]
            if v < min_val { min_val = v }
            if v > max_val { max_val = v }
            sum = sum + v
            if v < self.quiet_threshold { quiet = quiet + 1 }
            if v > self.overactive_threshold { over = over + 1 }
        }

        let n = self.gate_ema.len()
        let mut mean = 0.0
        if n > 0 { mean = sum / n }

        HomeostaticStatus {
            min_ema: min_val,
            max_ema: max_val,
            mean_ema: mean,
            quiet_count: quiet,
            overactive_count: over,
        }
    }

    fn reset(&mut self) {
        self.gate_ema = []
        self.initialized = false
    }
}

impl Evolve for HomeostaticScaler {
    fn evolve(&mut self, dt: Duration) {
        // Homeostatic scaling is event-driven, no temporal evolution
    }
}

struct HomeostaticActions {
    scaled_up: Vec<Int>,
    scaled_down: Vec<Int>,
}

struct HomeostaticStatus {
    min_ema: Float,
    max_ema: Float,
    mean_ema: Float,
    quiet_count: Int,
    overactive_count: Int,
}


// ============================================================
// CONSOLIDATION BUFFER -- Sleep Replay Analogue
// ============================================================
// During NREM sleep, sharp-wave ripples replay compressed versions
// of important waking experiences. Not everything consolidates.
// The brain actively selects based on surprise (high prediction error).
//
// Records (example_idx, loss, chemical_state) during training.
// Marks extreme events as "scars" (one-shot, always replayed).
// At consolidation time, returns top-K by loss + all scars.

struct ReplayCandidate {
    example_idx: Int,
    loss_value: Float,
    chemical_state: Vec<Float>,
    is_scar: Bool,
}

impl ReplayCandidate {
    fn new(idx: Int, loss: Float, state: Vec<Float>) -> ReplayCandidate {
        ReplayCandidate {
            example_idx: idx,
            loss_value: loss,
            chemical_state: state,
            is_scar: false,
        }
    }
}

temporal struct ConsolidationBuffer {
    state buffer: Vec<ReplayCandidate>,
    state scars: Vec<ReplayCandidate>,

    const max_buffer_size: Int = 1000,
    const scar_loss_threshold: Float = 5.0,
    const replay_fraction: Float = 0.1,
}

impl ConsolidationBuffer {
    fn new() -> ConsolidationBuffer {
        ConsolidationBuffer {
            buffer: [],
            scars: [],
            max_buffer_size: 1000,
            scar_loss_threshold: 5.0,
            replay_fraction: 0.1,
        }
    }

    fn record(&mut self, example_idx: Int, loss: Float, chemical_state: Vec<Float>) {
        // Record a training example and its surprise level
        let mut candidate = ReplayCandidate::new(example_idx, loss, chemical_state)

        // Check for scar (extreme event)
        if loss > self.scar_loss_threshold {
            candidate.is_scar = true
            self.scars = self.scars + [candidate]
        }

        // Add to buffer (replace least surprising if full)
        if self.buffer.len() >= self.max_buffer_size {
            // Find lowest-loss entry (least surprising)
            let mut min_idx = 0
            let mut min_loss = 999999.0
            for i in 0..self.buffer.len() {
                if self.buffer[i].loss_value < min_loss {
                    min_loss = self.buffer[i].loss_value
                    min_idx = i
                }
            }
            // Only replace if new example is more surprising
            if loss > min_loss {
                self.buffer[min_idx] = candidate
            }
        } else {
            self.buffer = self.buffer + [candidate]
        }
    }

    fn get_replay_batch(&self, k: Int) -> Vec<ReplayCandidate> {
        // Get examples for consolidation replay.
        // Returns all scars + top-K most surprising non-scar examples.
        let mut use_k = k
        if use_k <= 0 {
            use_k = self.buffer.len() / 10
            if use_k < 1 { use_k = 1 }
        }

        // Simple selection: find top-K by loss (highest surprise first)
        let mut replay = []
        let mut used = []

        for pick in 0..use_k {
            let mut best_idx = -1
            let mut best_loss = -1.0
            for i in 0..self.buffer.len() {
                if !self.buffer[i].is_scar {
                    let mut already = false
                    for u in used {
                        if u == i { already = true }
                    }
                    if !already {
                        if self.buffer[i].loss_value > best_loss {
                            best_loss = self.buffer[i].loss_value
                            best_idx = i
                        }
                    }
                }
            }
            if best_idx >= 0 {
                replay = replay + [self.buffer[best_idx]]
                used = used + [best_idx]
            }
        }

        // Always include ALL scars
        for scar in self.scars {
            replay = replay + [scar]
        }

        replay
    }

    fn get_stats(&self) -> ConsolidationStats {
        let mut sum = 0.0
        let mut max_loss = 0.0
        let mut min_loss = 999999.0

        for c in self.buffer {
            sum = sum + c.loss_value
            if c.loss_value > max_loss { max_loss = c.loss_value }
            if c.loss_value < min_loss { min_loss = c.loss_value }
        }

        let n = self.buffer.len()
        let mut mean = 0.0
        if n > 0 { mean = sum / n }
        if n == 0 { min_loss = 0.0 }

        ConsolidationStats {
            buffer_size: n,
            scar_count: self.scars.len(),
            mean_loss: mean,
            max_loss: max_loss,
            min_loss: min_loss,
        }
    }

    fn clear_non_scars(&mut self) {
        // Clear buffer but keep scars (sleep cycle complete)
        self.buffer = []
        for scar in self.scars {
            self.buffer = self.buffer + [scar]
        }
    }

    fn reset(&mut self) {
        // Full reset
        self.buffer = []
        self.scars = []
    }
}

impl Evolve for ConsolidationBuffer {
    fn evolve(&mut self, dt: Duration) {
        // Consolidation is event-driven, no temporal evolution
    }
}

struct ConsolidationStats {
    buffer_size: Int,
    scar_count: Int,
    mean_loss: Float,
    max_loss: Float,
    min_loss: Float,
}


// ============================================================
// HEBBIAN CO-ACTIVATION LOSS
// ============================================================
// Heads that fire together should wire together.
//
// For co-activated heads (both above threshold), their DNT
// connection should be strengthened. Loss = -sum(gate_i * gate_j)
// for all connected head pairs. Negative because we WANT
// co-activation (gradient descent minimizes, so negative = strengthen).

struct HebbianCoActivationLoss {
    threshold: Float,
    num_heads: Int,
}

impl HebbianCoActivationLoss {
    fn new() -> HebbianCoActivationLoss {
        HebbianCoActivationLoss {
            threshold: 0.3,
            num_heads: 34,
        }
    }

    fn with_threshold(mut self, t: Float) -> HebbianCoActivationLoss {
        self.threshold = t
        self
    }

    fn compute(&self, gate_values: Vec<Float>) -> Float {
        // Compute Hebbian co-activation loss.
        // Only consider heads above threshold.
        // Outer product: co-activation matrix [34, 34]
        // Sum of co-activations (excluding self-connections)
        // Returns negative value (co-activation is good, minimize negative = maximize)
        let mut hebbian_signal = 0.0

        for i in 0..self.num_heads {
            if gate_values[i] <= self.threshold { continue }
            for j in 0..self.num_heads {
                if i == j { continue }
                if gate_values[j] <= self.threshold { continue }
                hebbian_signal = hebbian_signal + gate_values[i] * gate_values[j]
            }
        }

        // Negative: we WANT co-activation. Scale down so it doesn't dominate.
        0.0 - hebbian_signal * 0.01
    }
}


// ============================================================
// COMPOSITE NEUROMODULATED LOSS
// ============================================================
// Complete neuromodulated loss function combining all components.
//
// Drop-in replacement for bare cross-entropy in training.
// Takes model outputs + chemical state and returns a single
// scalar loss that incorporates:
//
//   1. Chemistry-weighted LM loss (dopamine-modulated)
//   2. Chemistry-weighted state alignment (NE + cortisol modulated)
//   3. Local prediction errors (per-tier predictive coding)
//   4. Hebbian co-activation reward
//   5. Eligibility-trace weighted credit assignment

temporal struct NeuromodulatedLoss {
    state weighter: NeuromodulatedLossWeighter,
    state trace: EligibilityTraceBuffer,

    const local_pred: LocalPredictionLoss,
    const hebbian: HebbianCoActivationLoss,
}

impl NeuromodulatedLoss {
    fn new() -> NeuromodulatedLoss {
        NeuromodulatedLoss {
            weighter: NeuromodulatedLossWeighter::new(),
            trace: EligibilityTraceBuffer::new(),
            local_pred: LocalPredictionLoss::new(),
            hebbian: HebbianCoActivationLoss::new(),
        }
    }

    fn compute(
        &mut self,
        lm_loss: Float,
        state_vector: Vec<Float>,
        gate_values: Vec<Float>,
        state_loss: Float,
    ) -> NeuromodulatedLossResult {
        // Get chemistry-modulated weights
        let weights = self.weighter.compute(state_vector)

        // Local prediction errors (predictive coding)
        let local_details = self.local_pred.compute(state_vector)

        // Hebbian co-activation
        let hebbian_loss = self.hebbian.compute(gate_values)

        // Record trace and get eligibility weights
        self.trace.record(gate_values)
        let trace_weights = self.trace.get_trace_weight(weights.trace_decay)

        // Mean trace for credit assignment weighting
        let mut mean_trace = 0.0
        for i in 0..trace_weights.len() {
            mean_trace = mean_trace + trace_weights[i]
        }
        if trace_weights.len() > 0 {
            mean_trace = mean_trace / trace_weights.len()
        }

        // Combine with chemistry-modulated weights
        // Higher trace = this state has been building = more credit
        let total_loss = weights.lm_weight * lm_loss * (0.7 + 0.3 * mean_trace)
            + weights.state_weight * state_loss
            + weights.local_pred_weight * local_details.total_loss
            + weights.hebbian_weight * hebbian_loss

        NeuromodulatedLossResult {
            total_loss: total_loss,
            lm_loss: lm_loss,
            state_loss: state_loss,
            local_pred_loss: local_details.total_loss,
            hebbian_loss: hebbian_loss,
            trace_mean: mean_trace,
            weights: weights,
            local_details: local_details,
        }
    }

    fn get_lr_modifier(&mut self, state_vector: Vec<Float>) -> Float {
        // Get the chemistry-modulated learning rate multiplier.
        // Call in the training loop to adjust the optimizer's learning rate.
        let weights = self.weighter.compute(state_vector)
        weights.lr_modifier
    }
}

impl Evolve for NeuromodulatedLoss {
    fn evolve(&mut self, dt: Duration) {
        self.weighter.evolve(dt)
        self.trace.evolve(dt)
    }
}

struct NeuromodulatedLossResult {
    total_loss: Float,
    lm_loss: Float,
    state_loss: Float,
    local_pred_loss: Float,
    hebbian_loss: Float,
    trace_mean: Float,
    weights: LossWeights,
    local_details: LocalPredictionDetails,
}
