csm-core-lib 0.3.7

//! Echo State Network for temporal dynamics.
// Casts are intentional for reservoir math (node counts, dimension sizes)
#![allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
use crate::error::{MemoryError, Result};
use crate::hyperdim::HVec10240;
use crate::reservoir_sparse::SparseWeights;
#[cfg(target_arch = "wasm32")]
use js_sys::Date;
use rand::rngs::StdRng;
use rand::{RngExt, SeedableRng};
#[cfg(all(not(target_arch = "wasm32"), feature = "parallel"))]
use rayon::prelude::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
#[cfg(not(target_arch = "wasm32"))]
use {std::time::Instant, tracing::instrument};
#[derive(Debug, Default)]
pub struct ReservoirMetrics {
    pub steps_total: AtomicU64,
    pub step_latency_us_total: AtomicU64,
    pub step_latency_count: AtomicU64,
    pub nodes_active: AtomicU64,
}

#[derive(Debug, Clone, Default)]
pub struct ReservoirMetricsSnapshot {
    pub reservoir_steps_total: u64,
    pub avg_reservoir_step_latency_us: f64,
    pub reservoir_nodes_active: u64,
}

/// Output of a reservoir step (ADR-0078)
#[derive(Debug)]
pub struct ReservoirStepOutput<'a> {
    pub state: &'a [f32],
    pub state_norm: f64,
    pub change_norm: f64,
}
impl ReservoirMetrics {
    pub fn observe_step(&self, latency_us: u64, nodes_active: u64) {
        self.steps_total.fetch_add(1, Ordering::Relaxed);
        self.step_latency_us_total
            .fetch_add(latency_us, Ordering::Relaxed);
        self.step_latency_count.fetch_add(1, Ordering::Relaxed);
        self.nodes_active.store(nodes_active, Ordering::Relaxed);
    }
    pub fn snapshot(&self) -> ReservoirMetricsSnapshot {
        let count = self.step_latency_count.load(Ordering::Relaxed);
        let total = self.step_latency_us_total.load(Ordering::Relaxed);
        let avg = if count == 0 {
            0.0
        } else {
            total as f64 / count as f64
        };
        ReservoirMetricsSnapshot {
            reservoir_steps_total: self.steps_total.load(Ordering::Relaxed),
            avg_reservoir_step_latency_us: avg,
            reservoir_nodes_active: self.nodes_active.load(Ordering::Relaxed),
        }
    }
}
/// Sparse Echo State Network with chaotic dynamics
pub struct Reservoir {
    size: usize,
    input_size: usize,
    state: Vec<f32>,
    scratch: Vec<f32>,
    prev_state: Vec<f32>, // ADR-0064: t-1 state for inertia
    w_in: SparseWeights,
    w_res: SparseWeights,
    input_cache: Vec<f32>,
    input_projection: Vec<f32>,
    input_version: u32,
    node_versions: Vec<u32>,
    pub(crate) update_stride: usize,
    pub(crate) update_phase: usize,
    spectral_radius: f32,
    pub(crate) alpha: f32,
    pub(crate) beta: f32, // ADR-0064: inertia coefficient
    state_norm_sq: f64,   // Incremental state norm for performance
    metrics: Arc<ReservoirMetrics>,
}
impl Reservoir {
    pub const DEFAULT_SIZE: usize = 50000;
    pub const DEFAULT_RADIUS: f32 = 0.95;
    pub const DEFAULT_ALPHA: f32 = 0.3;
    const INPUT_DEGREE: usize = 4;
    const RESERVOIR_DEGREE: usize = 8;
    const RESERVOIR_LOCAL_WINDOW: usize = 512;
    const PARTIAL_UPDATE_STRIDE: usize = 32;
    pub const MAX_SIZE: usize = 100_000;

    pub fn validate_params(size: usize, input_size: usize, chaos: f32) -> Result<()> {
        if size == 0 || size > Self::MAX_SIZE {
            return Err(MemoryError::InvalidInput {
                field: "reservoir_size".into(),
                reason: format!("must be 1..={}", Self::MAX_SIZE),
            });
        }
        if input_size == 0 || input_size > Self::MAX_SIZE {
            return Err(MemoryError::InvalidInput {
                field: "reservoir_input_size".into(),
                reason: format!("must be 1..={}", Self::MAX_SIZE),
            });
        }
        if !chaos.is_finite() || chaos < 0.0 {
            return Err(MemoryError::InvalidInput {
                field: "chaos_strength".into(),
                reason: "must be finite and non-negative".into(),
            });
        }
        Ok(())
    }

    pub fn new(input_size: usize, size: usize) -> Result<Self> {
        Self::new_with_metrics(input_size, size, Arc::new(ReservoirMetrics::default()))
    }

    pub fn new_with_metrics(
        input_size: usize,
        size: usize,
        metrics: Arc<ReservoirMetrics>,
    ) -> Result<Self> {
        Self::new_seeded_with_metrics(input_size, size, rand::rng().random(), metrics)
    }

    pub fn new_seeded(input_size: usize, size: usize, seed: u64) -> Result<Self> {
        Self::new_seeded_with_metrics(
            input_size,
            size,
            seed,
            Arc::new(ReservoirMetrics::default()),
        )
    }

    pub fn new_seeded_with_metrics(
        input_size: usize,
        size: usize,
        seed: u64,
        metrics: Arc<ReservoirMetrics>,
    ) -> Result<Self> {
        Self::validate_params(size, input_size, 0.0)?;
        let mut rng = StdRng::seed_from_u64(seed);

        let w_in = SparseWeights::build(size, input_size, Self::INPUT_DEGREE, &mut rng);
        let mut w_res = SparseWeights::build_local_reservoir(
            size,
            Self::RESERVOIR_DEGREE.min(size),
            Self::RESERVOIR_LOCAL_WINDOW.min(size),
            &mut rng,
        );

        let current_radius = Self::estimate_spectral_radius(&w_res, size);
        if current_radius > 0.0 {
            let scale = Self::DEFAULT_RADIUS / current_radius;
            w_res.scale(scale);
        }

        Ok(Self {
            size,
            input_size,
            state: vec![0.0; size],
            scratch: vec![0.0; size],
            prev_state: vec![0.0; size], // ADR-0064
            w_in,
            w_res,
            input_cache: vec![0.0; input_size],
            input_projection: vec![0.0; size],
            input_version: 0,
            node_versions: vec![0; size],
            update_stride: Self::PARTIAL_UPDATE_STRIDE,
            update_phase: 0,
            spectral_radius: Self::DEFAULT_RADIUS,
            alpha: Self::DEFAULT_ALPHA,
            beta: 0.0, // ADR-0064: default no inertia
            state_norm_sq: 0.0,
            metrics,
        })
    }

    /// Single reservoir step
    #[cfg_attr(not(target_arch = "wasm32"), instrument(skip(self)))]
    pub fn step(&mut self, input: &[f32]) -> Result<ReservoirStepOutput<'_>> {
        #[cfg(not(target_arch = "wasm32"))]
        let started = Instant::now();
        #[cfg(target_arch = "wasm32")]
        let started = Date::now();

        if input.len() != self.input_size {
            return Err(MemoryError::reservoir(format!(
                "Input size mismatch: expected {}, got {}",
                self.input_size,
                input.len()
            )));
        }

        if self.input_cache != input {
            self.input_cache.copy_from_slice(input);
            self.input_version = self.input_version.wrapping_add(1);
        }

        let gamma = 1.0 - self.alpha + self.beta;
        let neg_beta = -self.beta;
        let update_phase = self.update_phase;
        let mut change_norm_sq = 0.0;
        for i in (update_phase..self.size).step_by(self.update_stride) {
            // SAFETY: all buffers (node_versions, input_projection, state, scratch, prev_state)
            // are sized to `self.size` and i is in (update_phase..self.size).
            // dot_row is safe because input.len() is verified to be self.input_size.
            unsafe {
                // Algorithmic Optimization: Lazy partial input projection.
                if *self.node_versions.get_unchecked(i) != self.input_version {
                    *self.input_projection.get_unchecked_mut(i) = self.w_in.dot_row(i, input);
                    *self.node_versions.get_unchecked_mut(i) = self.input_version;
                }

                let res_sum = self.w_res.dot_row(i, &self.state);
                let activated = fast_tanh(*self.input_projection.get_unchecked(i) + res_sum);
                let current_val = *self.state.get_unchecked(i);

                // Compute new state into scratch using a stable snapshot of `self.state`.
                // Optimized: Pre-calculated gamma (1.0 - alpha + beta) and neg_beta (-beta)
                // reduces per-node subtraction in the hot loop.
                *self.scratch.get_unchecked_mut(i) = activated.mul_add(
                    self.alpha,
                    current_val.mul_add(gamma, neg_beta * *self.prev_state.get_unchecked(i)),
                );

                let diff = *self.scratch.get_unchecked(i) - current_val;
                change_norm_sq += diff * diff;
            }
        }

        // Second pass: Commit the updates for this phase.
        // This keeps semantics synchronous within the phase (no order dependency).
        if update_phase == 0 {
            // Periodic full re-calculation to prevent drift from precision errors.
            // Optimized: Use direct multiplication instead of powi(2).
            let mut new_norm_sq = 0.0;
            for i in 0..self.size {
                // SAFETY: state, scratch, and prev_state are all sized to self.size.
                unsafe {
                    if (i % self.update_stride) == 0 {
                        let old_val = *self.state.get_unchecked(i);
                        let new_val = *self.scratch.get_unchecked(i);
                        *self.prev_state.get_unchecked_mut(i) = old_val;
                        *self.state.get_unchecked_mut(i) = new_val;
                    }
                    let val = f64::from(*self.state.get_unchecked(i));
                    new_norm_sq += val * val;
                }
            }
            self.state_norm_sq = new_norm_sq;
        } else {
            let mut delta_norm_sq = 0.0;
            for i in (update_phase..self.size).step_by(self.update_stride) {
                // SAFETY: state, scratch, and prev_state are all sized to self.size.
                unsafe {
                    let old_val = *self.state.get_unchecked(i);
                    let new_val = *self.scratch.get_unchecked(i);
                    *self.prev_state.get_unchecked_mut(i) = old_val;
                    *self.state.get_unchecked_mut(i) = new_val;
                    // Incremental norm update: subtract old square, add new square.
                    // Optimized: Use direct multiplication instead of powi(2).
                    // Using difference of squares (n^2 - o^2) = (n - o)(n + o) for numerical stability.
                    let old_f64 = f64::from(old_val);
                    let new_f64 = f64::from(new_val);
                    delta_norm_sq += (new_f64 - old_f64) * (new_f64 + old_f64);
                }
            }
            self.state_norm_sq += delta_norm_sq;
        }

        self.update_phase = (update_phase + 1) % self.update_stride;

        #[cfg(not(target_arch = "wasm32"))]
        let latency_us = u64::try_from(started.elapsed().as_micros()).unwrap_or(u64::MAX);
        #[cfg(target_arch = "wasm32")]
        let latency_us = ((Date::now() - started) * 1000.0) as u64;
        self.metrics.observe_step(latency_us, self.size as u64);

        Ok(ReservoirStepOutput {
            state: &self.state,
            state_norm: self.state_norm_sq.sqrt(),
            change_norm: (change_norm_sq as f64).sqrt(),
        })
    }

    /// Run reservoir for multiple steps
    #[cfg_attr(not(target_arch = "wasm32"), instrument(skip(self)))]
    pub fn run(&mut self, inputs: &[Vec<f32>]) -> Result<Vec<Vec<f32>>> {
        let mut states = Vec::with_capacity(inputs.len());
        for input in inputs {
            self.step(input)?;
            states.push(self.state.clone());
        }
        Ok(states)
    }

    /// Get current reservoir state
    pub fn state(&self) -> &[f32] {
        &self.state
    }

    /// Set spectral radius
    #[cfg_attr(not(target_arch = "wasm32"), instrument(skip(self)))]
    pub fn set_spectral_radius(&mut self, radius: f32) -> Result<()> {
        if !(0.9..=1.1).contains(&radius) {
            return Err(MemoryError::reservoir(
                "Spectral radius must be in [0.9, 1.1]".to_string(),
            ));
        }

        let current = Self::estimate_spectral_radius(&self.w_res, self.size);
        if current > 0.0 {
            let scale = radius / current;
            self.w_res.scale(scale);
            self.spectral_radius = radius;
        }

        Ok(())
    }

    /// Reset reservoir state
    #[cfg_attr(not(target_arch = "wasm32"), instrument(skip(self)))]
    pub fn reset(&mut self) {
        self.state.fill(0.0);
        self.scratch.fill(0.0);
        self.prev_state.fill(0.0);
        self.state_norm_sq = 0.0;
    }

    /// Project state to hypervector (parallel on non-WASM)
    #[cfg_attr(
        all(not(target_arch = "wasm32"), feature = "parallel"),
        instrument(skip(self))
    )]
    #[cfg(all(not(target_arch = "wasm32"), feature = "parallel"))]
    pub fn to_hypervector(&self) -> Result<HVec10240> {
        if self.size < HVec10240::DIMENSION {
            return Err(MemoryError::InvalidDimension {
                expected: HVec10240::DIMENSION,
                actual: self.size,
            });
        }
        let chunk_size = self.size / HVec10240::DIMENSION;
        let data: Vec<u128> = (0..80)
            .into_par_iter()
            .map(|i| {
                let mut word = 0u128;
                for j in 0..128 {
                    let bit_index = i * 128 + j;
                    let start = bit_index * chunk_size;
                    // SAFETY: bit_index < 10240 and chunk_size = size / 10240.
                    // start + chunk_size = (bit_index + 1) * chunk_size <= 10240 * (size / 10240) <= size.
                    // The range is always within bounds of self.state.
                    let sum: f32 = unsafe { self.state.get_unchecked(start..start + chunk_size) }
                        .iter()
                        .sum();
                    if sum > 0.0 {
                        word |= 1u128 << j;
                    }
                }
                word
            })
            .collect();
        let data: [u128; 80] = data.try_into().map_err(|_| {
            MemoryError::reservoir(
                "internal: par_iter produced unexpected element count".to_string(),
            )
        })?;
        Ok(HVec10240 { data })
    }

    #[cfg(any(target_arch = "wasm32", not(feature = "parallel")))]
    #[cfg_attr(not(target_arch = "wasm32"), instrument(skip(self)))]
    pub fn to_hypervector(&self) -> Result<HVec10240> {
        if self.size < HVec10240::DIMENSION {
            return Err(MemoryError::InvalidDimension {
                expected: HVec10240::DIMENSION,
                actual: self.size,
            });
        }
        let chunk_size = self.size / HVec10240::DIMENSION;
        let mut data = [0u128; 80];
        for (i, word) in data.iter_mut().enumerate() {
            for j in 0..128 {
                let bit_index = i * 128 + j;
                let start = bit_index * chunk_size;
                // SAFETY: bit_index < 10240 and chunk_size = size / 10240.
                // start + chunk_size = (bit_index + 1) * chunk_size <= 10240 * (size / 10240) <= size.
                // The range is always within bounds of self.state.
                let sum: f32 = unsafe { self.state.get_unchecked(start..start + chunk_size) }
                    .iter()
                    .sum();
                if sum > 0.0 {
                    *word |= 1u128 << j;
                }
            }
        }
        Ok(HVec10240 { data })
    }

    pub const fn size(&self) -> usize {
        self.size
    }

    pub fn metrics_snapshot(&self) -> ReservoirMetricsSnapshot {
        self.metrics.snapshot()
    }

    /// Estimate spectral radius using power iteration
    fn estimate_spectral_radius(w: &SparseWeights, size: usize) -> f32 {
        let mut v = vec![1.0f32 / size as f32; size];
        let mut y = vec![0.0f32; size];

        for _ in 0..16 {
            for (i, y_i) in y.iter_mut().enumerate() {
                // SAFETY: i is from 0..size, and y is sized to size.
                // w.dot_row is safe because v is sized to size.
                unsafe { *y_i = w.dot_row(i, &v) };
            }

            let mut norm = 0.0f32;
            for val in &y {
                norm += val * val;
            }
            norm = norm.sqrt();
            if norm.abs() < f32::EPSILON {
                return 0.0;
            }

            for i in 0..size {
                v[i] = y[i] / norm;
            }
        }

        let mut wv = vec![0.0f32; size];
        for (i, wv_i) in wv.iter_mut().enumerate() {
            // SAFETY: i is from 0..size, and wv is sized to size.
            // w.dot_row is safe because v is sized to size.
            unsafe { *wv_i = w.dot_row(i, &v) };
        }

        let mut numerator = 0.0f32;
        let mut denominator = 0.0f32;
        for i in 0..size {
            numerator += v[i] * wv[i];
            denominator += v[i] * v[i];
        }

        if denominator.abs() < f32::EPSILON {
            0.0
        } else {
            (numerator / denominator).abs()
        }
    }
}

#[inline]
fn fast_tanh(x: f32) -> f32 {
    let x2 = x * x;
    // Approximates tanh(x) as x*(27+x^2)/(27+9x^2) using FMA for speed.
    // Optimized: Redundant multiplication removed.
    x * (x2 + 27.0) / x2.mul_add(9.0, 27.0)
}
pub use crate::reservoir_chaotic::ChaoticReservoir;