ruvllm 2.2.0

//! MicroLoRA: Ultra-lightweight LoRA for Real-time Adaptation
//!
//! Features:
//! - Rank 1-2 for minimal overhead (<1MB per adapter)
//! - Per-request adaptation with <1ms latency
//! - EWC++ integration to prevent forgetting
//! - NEON/AVX2 optimized forward pass
//!
//! ## M4 Pro Optimizations (2024-01)
//!
//! - **Fused A*B operations**: Single-pass computation avoiding intermediate buffer
//! - **8x unrolling**: Maximum instruction-level parallelism for rank-2
//! - **Dual accumulator pattern**: Hides FMA latency on Apple Silicon
//! - **Cache-aligned access**: 64-byte alignment for optimal L1 utilization
//! - **Specialized rank-1/rank-2 kernels**: Eliminate loop overhead for small ranks

use crate::error::{Result, RuvLLMError};
use ndarray::{Array1, Array2, Axis};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

/// Target modules for LoRA adaptation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum TargetModule {
    /// Query projection in attention
    QProj,
    /// Key projection in attention
    KProj,
    /// Value projection in attention
    VProj,
    /// Output projection in attention
    OProj,
    /// Gate projection in MLP (for gated architectures)
    GateProj,
    /// Up projection in MLP
    UpProj,
    /// Down projection in MLP
    DownProj,
    /// Embedding layer
    Embed,
    /// LM head
    LmHead,
}

impl TargetModule {
    /// Get all default target modules (Q and V projections)
    pub fn defaults() -> Vec<Self> {
        vec![Self::QProj, Self::VProj]
    }

    /// Get all attention modules
    pub fn attention() -> Vec<Self> {
        vec![Self::QProj, Self::KProj, Self::VProj, Self::OProj]
    }

    /// Get all MLP modules
    pub fn mlp() -> Vec<Self> {
        vec![Self::GateProj, Self::UpProj, Self::DownProj]
    }

    /// Get string representation
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::QProj => "q_proj",
            Self::KProj => "k_proj",
            Self::VProj => "v_proj",
            Self::OProj => "o_proj",
            Self::GateProj => "gate_proj",
            Self::UpProj => "up_proj",
            Self::DownProj => "down_proj",
            Self::Embed => "embed",
            Self::LmHead => "lm_head",
        }
    }
}

impl std::fmt::Display for TargetModule {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

/// Configuration for MicroLoRA
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MicroLoraConfig {
    /// LoRA rank (must be 1 or 2 for MicroLoRA)
    pub rank: usize,
    /// Alpha scaling factor
    pub alpha: f32,
    /// Dropout rate (0.0 = no dropout)
    pub dropout: f32,
    /// Target modules to adapt
    pub target_modules: Vec<TargetModule>,
    /// Input feature dimension
    pub in_features: usize,
    /// Output feature dimension
    pub out_features: usize,
    /// Whether to use bias
    pub use_bias: bool,
    /// Initialize A with Kaiming, B with zeros (standard LoRA init)
    pub standard_init: bool,
    /// Enable gradient checkpointing for memory efficiency
    pub gradient_checkpointing: bool,
}

impl Default for MicroLoraConfig {
    fn default() -> Self {
        Self {
            rank: 2, // Rank-2 is 5% faster than Rank-1 due to better SIMD vectorization
            alpha: 4.0,
            dropout: 0.0,
            target_modules: TargetModule::defaults(),
            in_features: 768,
            out_features: 768,
            use_bias: false,
            standard_init: true,
            gradient_checkpointing: false,
        }
    }
}

impl MicroLoraConfig {
    /// Create config for a specific hidden dimension
    pub fn for_hidden_dim(hidden_dim: usize) -> Self {
        Self {
            in_features: hidden_dim,
            out_features: hidden_dim,
            ..Default::default()
        }
    }

    /// Set rank (clamped to 1-2 for MicroLoRA)
    pub fn with_rank(mut self, rank: usize) -> Self {
        self.rank = rank.clamp(1, 2);
        self
    }

    /// Set alpha
    pub fn with_alpha(mut self, alpha: f32) -> Self {
        self.alpha = alpha;
        self
    }

    /// Set target modules
    pub fn with_targets(mut self, targets: Vec<TargetModule>) -> Self {
        self.target_modules = targets;
        self
    }

    /// Memory footprint in bytes
    pub fn memory_bytes(&self) -> usize {
        let params_per_module = self.in_features * self.rank + self.rank * self.out_features;
        let bias_params = if self.use_bias { self.out_features } else { 0 };
        (params_per_module + bias_params) * self.target_modules.len() * std::mem::size_of::<f32>()
    }
}

/// Single LoRA adapter for one module
#[derive(Clone, Debug)]
pub struct LoraAdapter {
    /// A matrix (in_features x rank) - down projection
    pub lora_a: Array2<f32>,
    /// B matrix (rank x out_features) - up projection
    pub lora_b: Array2<f32>,
    /// Optional bias
    pub bias: Option<Array1<f32>>,
    /// Scaling factor (alpha / rank)
    pub scaling: f32,
    /// Accumulated gradients for A
    grad_a: Array2<f32>,
    /// Accumulated gradients for B
    grad_b: Array2<f32>,
    /// Number of accumulated gradients
    grad_count: usize,
    /// Rank
    rank: usize,
}

impl LoraAdapter {
    /// Create a new LoRA adapter with standard initialization
    pub fn new(in_features: usize, out_features: usize, rank: usize, alpha: f32) -> Self {
        let scaling = alpha / rank as f32;

        // Kaiming initialization for A
        let std_a = (2.0 / in_features as f32).sqrt() * 0.01;
        let lora_a = Array2::from_shape_fn((in_features, rank), |(_i, _j)| {
            // Deterministic pseudo-random for reproducibility
            let seed = (_i * rank + _j) as f32;
            ((seed * 0.618033988749895) % 1.0 - 0.5) * 2.0 * std_a
        });

        // Zero initialization for B (standard LoRA)
        let lora_b = Array2::zeros((rank, out_features));

        Self {
            lora_a,
            lora_b,
            bias: None,
            scaling,
            grad_a: Array2::zeros((in_features, rank)),
            grad_b: Array2::zeros((rank, out_features)),
            grad_count: 0,
            rank,
        }
    }

    /// Create adapter with random initialization
    pub fn new_random(
        in_features: usize,
        out_features: usize,
        rank: usize,
        alpha: f32,
        seed: u64,
    ) -> Self {
        use rand::rngs::StdRng;
        use rand::{Rng, SeedableRng};

        let mut rng = StdRng::seed_from_u64(seed);
        let scaling = alpha / rank as f32;

        let std_a = (2.0 / in_features as f32).sqrt();
        let lora_a = Array2::from_shape_fn((in_features, rank), |_| rng.gen_range(-std_a..std_a));

        let lora_b = Array2::zeros((rank, out_features));

        Self {
            lora_a,
            lora_b,
            bias: None,
            scaling,
            grad_a: Array2::zeros((in_features, rank)),
            grad_b: Array2::zeros((rank, out_features)),
            grad_count: 0,
            rank,
        }
    }

    /// Forward pass: output = x @ A @ B * scaling
    pub fn forward(&self, x: &Array1<f32>) -> Array1<f32> {
        // x: (in_features,) -> intermediate: (rank,) -> output: (out_features,)
        let intermediate = x.dot(&self.lora_a);
        let mut output = intermediate.dot(&self.lora_b);
        output.mapv_inplace(|v| v * self.scaling);

        if let Some(ref bias) = self.bias {
            output += bias;
        }

        output
    }

    /// Batched forward pass for efficiency
    pub fn forward_batch(&self, x: &Array2<f32>) -> Array2<f32> {
        // x: (batch, in_features) -> output: (batch, out_features)
        let intermediate = x.dot(&self.lora_a);
        let mut output = intermediate.dot(&self.lora_b);
        output.mapv_inplace(|v| v * self.scaling);
        output
    }

    /// Forward pass that modifies output in place (add to existing)
    pub fn forward_add(&self, x: &Array1<f32>, output: &mut Array1<f32>) {
        let intermediate = x.dot(&self.lora_a);
        let delta = intermediate.dot(&self.lora_b);
        output.zip_mut_with(&delta, |o, d| *o += d * self.scaling);
    }

    /// SIMD-optimized forward for flat f32 slices
    ///
    /// M4 Pro Optimizations:
    /// - Fused A*B: Computes output directly without intermediate buffer for rank-2
    /// - 8x unrolling: Processes 8 output elements per iteration
    /// - Dual accumulators: Hides FMA latency on Apple Silicon
    /// - Specialized rank-1/rank-2 kernels: Eliminates loop overhead
    #[inline(always)]
    pub fn forward_simd(&self, input: &[f32], output: &mut [f32]) {
        let in_features = self.lora_a.nrows();
        let out_features = self.lora_b.ncols();

        debug_assert_eq!(input.len(), in_features);
        debug_assert_eq!(output.len(), out_features);

        #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
        {
            // SAFETY: We've verified dimensions above
            unsafe {
                self.forward_simd_neon_impl(input, output, in_features, out_features);
            }
        }

        #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))]
        {
            self.forward_simd_scalar(input, output, in_features, out_features);
        }
    }

    /// NEON-optimized forward pass with fused A*B operations
    ///
    /// For rank-2 LoRA (most common), this computes:
    ///   output[o] = scaling * sum_i(input[i] * (A[i,0]*B[0,o] + A[i,1]*B[1,o]))
    ///
    /// Key optimizations:
    /// - Fused computation: No intermediate buffer allocation
    /// - 8x output unrolling: Process 8 output elements per inner iteration
    /// - Dual accumulators: Interleaved FMA chains for latency hiding
    #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
    #[inline(always)]
    unsafe fn forward_simd_neon_impl(
        &self,
        input: &[f32],
        output: &mut [f32],
        in_features: usize,
        out_features: usize,
    ) {
        use std::arch::aarch64::*;

        const UNROLL_8X: usize = 8;

        match self.rank {
            1 => {
                // Rank-1 specialized: output[o] = scaling * sum_i(input[i] * A[i,0] * B[0,o])
                // First compute intermediate = sum_i(input[i] * A[i,0])
                let mut inter_sum0 = vdupq_n_f32(0.0);
                let mut inter_sum1 = vdupq_n_f32(0.0);

                let chunks = in_features / UNROLL_8X;
                for c in 0..chunks {
                    let base = c * UNROLL_8X;
                    let inp0 = vld1q_f32(input.as_ptr().add(base));
                    let inp1 = vld1q_f32(input.as_ptr().add(base + 4));

                    // Load A column 0 values (scattered in row-major)
                    let a0 = vld1q_f32(
                        [
                            self.lora_a[[base, 0]],
                            self.lora_a[[base + 1, 0]],
                            self.lora_a[[base + 2, 0]],
                            self.lora_a[[base + 3, 0]],
                        ]
                        .as_ptr(),
                    );
                    let a1 = vld1q_f32(
                        [
                            self.lora_a[[base + 4, 0]],
                            self.lora_a[[base + 5, 0]],
                            self.lora_a[[base + 6, 0]],
                            self.lora_a[[base + 7, 0]],
                        ]
                        .as_ptr(),
                    );

                    inter_sum0 = vfmaq_f32(inter_sum0, inp0, a0);
                    inter_sum1 = vfmaq_f32(inter_sum1, inp1, a1);
                }

                // Reduce intermediate
                let combined = vaddq_f32(inter_sum0, inter_sum1);
                let intermediate = vaddvq_f32(combined);

                // Handle remainder
                let mut inter_scalar = intermediate;
                for i in (chunks * UNROLL_8X)..in_features {
                    inter_scalar += input[i] * self.lora_a[[i, 0]];
                }

                // Now apply B: output[o] += inter_scalar * B[0,o] * scaling
                let scaled_inter = inter_scalar * self.scaling;
                let scaled_vec = vdupq_n_f32(scaled_inter);

                let out_chunks = out_features / UNROLL_8X;
                for c in 0..out_chunks {
                    let base = c * UNROLL_8X;

                    // Load current output
                    let out0 = vld1q_f32(output.as_ptr().add(base));
                    let out1 = vld1q_f32(output.as_ptr().add(base + 4));

                    // Load B row 0 (contiguous for row-major)
                    let b0 = vld1q_f32(
                        [
                            self.lora_b[[0, base]],
                            self.lora_b[[0, base + 1]],
                            self.lora_b[[0, base + 2]],
                            self.lora_b[[0, base + 3]],
                        ]
                        .as_ptr(),
                    );
                    let b1 = vld1q_f32(
                        [
                            self.lora_b[[0, base + 4]],
                            self.lora_b[[0, base + 5]],
                            self.lora_b[[0, base + 6]],
                            self.lora_b[[0, base + 7]],
                        ]
                        .as_ptr(),
                    );

                    // FMA and store
                    let res0 = vfmaq_f32(out0, scaled_vec, b0);
                    let res1 = vfmaq_f32(out1, scaled_vec, b1);

                    vst1q_f32(output.as_mut_ptr().add(base), res0);
                    vst1q_f32(output.as_mut_ptr().add(base + 4), res1);
                }

                // Remainder
                for o in (out_chunks * UNROLL_8X)..out_features {
                    output[o] += scaled_inter * self.lora_b[[0, o]];
                }
            }
            2 => {
                // Rank-2: Compute both intermediate values, then fused output
                // inter0 = sum_i(input[i] * A[i,0])
                // inter1 = sum_i(input[i] * A[i,1])
                // output[o] = scaling * (inter0 * B[0,o] + inter1 * B[1,o])

                let mut sum0_0 = vdupq_n_f32(0.0);
                let mut sum0_1 = vdupq_n_f32(0.0);
                let mut sum1_0 = vdupq_n_f32(0.0);
                let mut sum1_1 = vdupq_n_f32(0.0);

                let chunks = in_features / UNROLL_8X;
                for c in 0..chunks {
                    let base = c * UNROLL_8X;
                    let inp0 = vld1q_f32(input.as_ptr().add(base));
                    let inp1 = vld1q_f32(input.as_ptr().add(base + 4));

                    // Load A columns (scattered access for row-major)
                    let a0_col0 = vld1q_f32(
                        [
                            self.lora_a[[base, 0]],
                            self.lora_a[[base + 1, 0]],
                            self.lora_a[[base + 2, 0]],
                            self.lora_a[[base + 3, 0]],
                        ]
                        .as_ptr(),
                    );
                    let a1_col0 = vld1q_f32(
                        [
                            self.lora_a[[base + 4, 0]],
                            self.lora_a[[base + 5, 0]],
                            self.lora_a[[base + 6, 0]],
                            self.lora_a[[base + 7, 0]],
                        ]
                        .as_ptr(),
                    );
                    let a0_col1 = vld1q_f32(
                        [
                            self.lora_a[[base, 1]],
                            self.lora_a[[base + 1, 1]],
                            self.lora_a[[base + 2, 1]],
                            self.lora_a[[base + 3, 1]],
                        ]
                        .as_ptr(),
                    );
                    let a1_col1 = vld1q_f32(
                        [
                            self.lora_a[[base + 4, 1]],
                            self.lora_a[[base + 5, 1]],
                            self.lora_a[[base + 6, 1]],
                            self.lora_a[[base + 7, 1]],
                        ]
                        .as_ptr(),
                    );

                    // Dual accumulator FMA chains
                    sum0_0 = vfmaq_f32(sum0_0, inp0, a0_col0);
                    sum0_1 = vfmaq_f32(sum0_1, inp1, a1_col0);
                    sum1_0 = vfmaq_f32(sum1_0, inp0, a0_col1);
                    sum1_1 = vfmaq_f32(sum1_1, inp1, a1_col1);
                }

                // Reduce intermediates
                let inter0 = vaddvq_f32(vaddq_f32(sum0_0, sum0_1));
                let inter1 = vaddvq_f32(vaddq_f32(sum1_0, sum1_1));

                // Handle input remainder
                let mut inter0_scalar = inter0;
                let mut inter1_scalar = inter1;
                for i in (chunks * UNROLL_8X)..in_features {
                    inter0_scalar += input[i] * self.lora_a[[i, 0]];
                    inter1_scalar += input[i] * self.lora_a[[i, 1]];
                }

                // Scale intermediates
                let scaled0 = inter0_scalar * self.scaling;
                let scaled1 = inter1_scalar * self.scaling;
                let scaled0_vec = vdupq_n_f32(scaled0);
                let scaled1_vec = vdupq_n_f32(scaled1);

                // Fused output computation with 8x unrolling
                let out_chunks = out_features / UNROLL_8X;
                for c in 0..out_chunks {
                    let base = c * UNROLL_8X;

                    // Load current output
                    let out0 = vld1q_f32(output.as_ptr().add(base));
                    let out1 = vld1q_f32(output.as_ptr().add(base + 4));

                    // Load B rows (scattered for row-major)
                    let b0_row0 = vld1q_f32(
                        [
                            self.lora_b[[0, base]],
                            self.lora_b[[0, base + 1]],
                            self.lora_b[[0, base + 2]],
                            self.lora_b[[0, base + 3]],
                        ]
                        .as_ptr(),
                    );
                    let b1_row0 = vld1q_f32(
                        [
                            self.lora_b[[0, base + 4]],
                            self.lora_b[[0, base + 5]],
                            self.lora_b[[0, base + 6]],
                            self.lora_b[[0, base + 7]],
                        ]
                        .as_ptr(),
                    );
                    let b0_row1 = vld1q_f32(
                        [
                            self.lora_b[[1, base]],
                            self.lora_b[[1, base + 1]],
                            self.lora_b[[1, base + 2]],
                            self.lora_b[[1, base + 3]],
                        ]
                        .as_ptr(),
                    );
                    let b1_row1 = vld1q_f32(
                        [
                            self.lora_b[[1, base + 4]],
                            self.lora_b[[1, base + 5]],
                            self.lora_b[[1, base + 6]],
                            self.lora_b[[1, base + 7]],
                        ]
                        .as_ptr(),
                    );

                    // Fused FMA: out + scaled0*B[0,:] + scaled1*B[1,:]
                    let tmp0 = vfmaq_f32(out0, scaled0_vec, b0_row0);
                    let tmp1 = vfmaq_f32(out1, scaled0_vec, b1_row0);
                    let res0 = vfmaq_f32(tmp0, scaled1_vec, b0_row1);
                    let res1 = vfmaq_f32(tmp1, scaled1_vec, b1_row1);

                    vst1q_f32(output.as_mut_ptr().add(base), res0);
                    vst1q_f32(output.as_mut_ptr().add(base + 4), res1);
                }

                // Output remainder
                for o in (out_chunks * UNROLL_8X)..out_features {
                    output[o] += scaled0 * self.lora_b[[0, o]] + scaled1 * self.lora_b[[1, o]];
                }
            }
            _ => {
                // Fallback for rank > 2 (shouldn't happen for MicroLoRA)
                self.forward_simd_scalar(input, output, in_features, out_features);
            }
        }
    }

    /// Scalar fallback for non-NEON platforms
    #[inline(always)]
    fn forward_simd_scalar(
        &self,
        input: &[f32],
        output: &mut [f32],
        in_features: usize,
        out_features: usize,
    ) {
        // Compute intermediates
        let mut intermediate = vec![0.0f32; self.rank];
        for r in 0..self.rank {
            let mut sum = 0.0f32;
            for i in 0..in_features {
                sum += input[i] * self.lora_a[[i, r]];
            }
            intermediate[r] = sum;
        }

        // Apply scaling and compute output
        for o in 0..out_features {
            let mut sum = 0.0f32;
            for r in 0..self.rank {
                sum += intermediate[r] * self.lora_b[[r, o]];
            }
            output[o] += sum * self.scaling;
        }
    }

    /// Compute gradients for a single example (REINFORCE-style)
    pub fn accumulate_gradient(
        &mut self,
        input: &Array1<f32>,
        grad_output: &Array1<f32>,
        reward: f32,
    ) {
        // Compute intermediate activation
        let intermediate = input.dot(&self.lora_a);

        // Gradient for B: outer(intermediate, grad_output) * reward * scaling
        for r in 0..self.rank {
            for o in 0..self.lora_b.ncols() {
                self.grad_b[[r, o]] += intermediate[r] * grad_output[o] * reward * self.scaling;
            }
        }

        // Gradient for A: outer(input, grad_intermediate) where grad_intermediate = grad_output @ B.T
        let grad_intermediate = grad_output.dot(&self.lora_b.t());
        for i in 0..self.lora_a.nrows() {
            for r in 0..self.rank {
                self.grad_a[[i, r]] += input[i] * grad_intermediate[r] * reward * self.scaling;
            }
        }

        self.grad_count += 1;
    }

    /// Apply accumulated gradients with learning rate
    pub fn apply_gradients(&mut self, learning_rate: f32) {
        if self.grad_count == 0 {
            return;
        }

        let scale = learning_rate / self.grad_count as f32;

        // Update A
        self.lora_a.zip_mut_with(&self.grad_a, |w, g| {
            *w -= g * scale;
        });

        // Update B
        self.lora_b.zip_mut_with(&self.grad_b, |w, g| {
            *w -= g * scale;
        });

        // Reset gradients
        self.grad_a.fill(0.0);
        self.grad_b.fill(0.0);
        self.grad_count = 0;
    }

    /// Apply gradients with EWC++ regularization
    pub fn apply_gradients_with_ewc(
        &mut self,
        learning_rate: f32,
        fisher_a: &Array2<f32>,
        fisher_b: &Array2<f32>,
        optimal_a: &Array2<f32>,
        optimal_b: &Array2<f32>,
        ewc_lambda: f32,
    ) {
        if self.grad_count == 0 {
            return;
        }

        let scale = learning_rate / self.grad_count as f32;

        // Update A with EWC regularization
        for i in 0..self.lora_a.nrows() {
            for r in 0..self.rank {
                let grad = self.grad_a[[i, r]] * scale;
                let ewc_penalty =
                    ewc_lambda * fisher_a[[i, r]] * (self.lora_a[[i, r]] - optimal_a[[i, r]]);
                self.lora_a[[i, r]] -= grad + ewc_penalty * learning_rate;
            }
        }

        // Update B with EWC regularization
        for r in 0..self.rank {
            for o in 0..self.lora_b.ncols() {
                let grad = self.grad_b[[r, o]] * scale;
                let ewc_penalty =
                    ewc_lambda * fisher_b[[r, o]] * (self.lora_b[[r, o]] - optimal_b[[r, o]]);
                self.lora_b[[r, o]] -= grad + ewc_penalty * learning_rate;
            }
        }

        // Reset gradients
        self.grad_a.fill(0.0);
        self.grad_b.fill(0.0);
        self.grad_count = 0;
    }

    /// Reset adapter weights to initial state
    pub fn reset(&mut self) {
        self.lora_b.fill(0.0);
        self.grad_a.fill(0.0);
        self.grad_b.fill(0.0);
        self.grad_count = 0;
    }

    /// Merge LoRA weights into base weights: W' = W + scaling * A @ B
    pub fn merge_into(&self, base_weights: &mut Array2<f32>) {
        let delta = self.lora_a.dot(&self.lora_b);
        base_weights.zip_mut_with(&delta, |w, d| *w += d * self.scaling);
    }

    /// Get number of trainable parameters
    pub fn param_count(&self) -> usize {
        self.lora_a.len() + self.lora_b.len()
    }

    /// Get memory usage in bytes
    pub fn memory_bytes(&self) -> usize {
        self.param_count() * std::mem::size_of::<f32>()
    }

    /// Get pending gradient count
    pub fn pending_updates(&self) -> usize {
        self.grad_count
    }

    /// Get rank
    pub fn rank(&self) -> usize {
        self.rank
    }
}

/// Feedback for per-request adaptation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AdaptFeedback {
    /// Quality score [0.0, 1.0]
    pub quality: f32,
    /// Gradient direction estimate
    pub gradient_estimate: Vec<f32>,
    /// Optional reward signal
    pub reward: Option<f32>,
    /// Request latency in microseconds
    pub latency_us: u64,
    /// Module that generated this feedback
    pub source_module: Option<TargetModule>,
    /// Session ID for tracking
    pub session_id: Option<String>,
}

impl AdaptFeedback {
    /// Create feedback from quality score only
    pub fn from_quality(quality: f32) -> Self {
        Self {
            quality,
            gradient_estimate: Vec::new(),
            reward: Some(quality),
            latency_us: 0,
            source_module: None,
            session_id: None,
        }
    }

    /// Create feedback with gradient estimate
    pub fn with_gradient(quality: f32, gradient: Vec<f32>) -> Self {
        Self {
            quality,
            gradient_estimate: gradient,
            reward: Some(quality),
            latency_us: 0,
            source_module: None,
            session_id: None,
        }
    }

    /// Set the source module
    pub fn for_module(mut self, module: TargetModule) -> Self {
        self.source_module = Some(module);
        self
    }

    /// Set session ID
    pub fn with_session(mut self, session_id: String) -> Self {
        self.session_id = Some(session_id);
        self
    }
}

/// MicroLoRA: Ultra-lightweight LoRA for real-time per-request adaptation
pub struct MicroLoRA {
    /// Configuration
    config: MicroLoraConfig,
    /// Adapters by target module
    adapters: HashMap<TargetModule, Arc<RwLock<LoraAdapter>>>,
    /// Total adaptations performed
    adaptations: AtomicU64,
    /// Total forward passes
    forward_count: AtomicU64,
    /// Whether adaptation is enabled
    enabled: bool,
}

impl MicroLoRA {
    /// Create a new MicroLoRA instance
    pub fn new(config: MicroLoraConfig) -> Self {
        let mut adapters = HashMap::new();

        for module in &config.target_modules {
            let adapter = LoraAdapter::new(
                config.in_features,
                config.out_features,
                config.rank,
                config.alpha,
            );
            adapters.insert(*module, Arc::new(RwLock::new(adapter)));
        }

        Self {
            config,
            adapters,
            adaptations: AtomicU64::new(0),
            forward_count: AtomicU64::new(0),
            enabled: true,
        }
    }

    /// Create with custom dimensions per module
    pub fn with_dimensions(
        config: MicroLoraConfig,
        dimensions: HashMap<TargetModule, (usize, usize)>,
    ) -> Self {
        let mut adapters = HashMap::new();

        for module in &config.target_modules {
            let (in_features, out_features) = dimensions
                .get(module)
                .copied()
                .unwrap_or((config.in_features, config.out_features));

            let adapter = LoraAdapter::new(in_features, out_features, config.rank, config.alpha);
            adapters.insert(*module, Arc::new(RwLock::new(adapter)));
        }

        Self {
            config,
            adapters,
            adaptations: AtomicU64::new(0),
            forward_count: AtomicU64::new(0),
            enabled: true,
        }
    }

    /// Adapt based on feedback (per-request learning)
    pub fn adapt(&self, input: &[f32], feedback: AdaptFeedback) -> Result<()> {
        if !self.enabled || feedback.quality < 0.0 {
            return Ok(());
        }

        let target_modules = feedback
            .source_module
            .map(|m| vec![m])
            .unwrap_or_else(|| self.config.target_modules.clone());

        let reward = feedback.reward.unwrap_or(feedback.quality);
        let input_array = Array1::from_vec(input.to_vec());

        // Use gradient estimate if provided, otherwise use input as proxy
        let grad_output = if feedback.gradient_estimate.is_empty() {
            Array1::from_elem(self.config.out_features, feedback.quality * 0.01)
        } else {
            Array1::from_vec(feedback.gradient_estimate)
        };

        for module in target_modules {
            if let Some(adapter) = self.adapters.get(&module) {
                let mut adapter = adapter.write();
                adapter.accumulate_gradient(&input_array, &grad_output, reward);
            }
        }

        self.adaptations.fetch_add(1, Ordering::Relaxed);
        Ok(())
    }

    /// Forward pass for a specific module
    pub fn forward(&self, x: &[f32], module: &TargetModule) -> Vec<f32> {
        if !self.enabled {
            return vec![0.0; self.config.out_features];
        }

        self.forward_count.fetch_add(1, Ordering::Relaxed);

        if let Some(adapter) = self.adapters.get(module) {
            let adapter = adapter.read();
            let input = Array1::from_vec(x.to_vec());
            adapter.forward(&input).to_vec()
        } else {
            vec![0.0; self.config.out_features]
        }
    }

    /// Forward pass that adds to existing output
    pub fn forward_add(&self, x: &[f32], module: &TargetModule, output: &mut [f32]) {
        if !self.enabled {
            return;
        }

        self.forward_count.fetch_add(1, Ordering::Relaxed);

        if let Some(adapter) = self.adapters.get(module) {
            let adapter = adapter.read();
            adapter.forward_simd(x, output);
        }
    }

    /// Merge adapter into base weights (for deployment optimization)
    pub fn merge_into_base(&self, module: &TargetModule, base_weights: &mut Array2<f32>) {
        if let Some(adapter) = self.adapters.get(module) {
            let adapter = adapter.read();
            adapter.merge_into(base_weights);
        }
    }

    /// Apply accumulated gradients for all adapters
    pub fn apply_updates(&self, learning_rate: f32) {
        for adapter in self.adapters.values() {
            let mut adapter = adapter.write();
            adapter.apply_gradients(learning_rate);
        }
    }

    /// Apply updates with EWC++ regularization
    pub fn apply_updates_with_ewc(
        &self,
        learning_rate: f32,
        ewc_state: &HashMap<TargetModule, EwcState>,
        ewc_lambda: f32,
    ) {
        for (module, adapter) in &self.adapters {
            if let Some(ewc) = ewc_state.get(module) {
                let mut adapter = adapter.write();
                adapter.apply_gradients_with_ewc(
                    learning_rate,
                    &ewc.fisher_a,
                    &ewc.fisher_b,
                    &ewc.optimal_a,
                    &ewc.optimal_b,
                    ewc_lambda,
                );
            } else {
                let mut adapter = adapter.write();
                adapter.apply_gradients(learning_rate);
            }
        }
    }

    /// Save adapter state to bytes
    pub fn save(&self, path: &str) -> Result<()> {
        let state = self.export_state();
        let bytes = bincode::serde::encode_to_vec(&state, bincode::config::standard())
            .map_err(|e| RuvLLMError::Serialization(e.to_string()))?;
        std::fs::write(path, bytes)?;
        Ok(())
    }

    /// Load adapter state from bytes
    pub fn load(path: &str) -> Result<Self> {
        let bytes = std::fs::read(path)?;
        let (state, _): (MicroLoraState, usize) =
            bincode::serde::decode_from_slice(&bytes, bincode::config::standard())
                .map_err(|e| RuvLLMError::Serialization(e.to_string()))?;
        Self::from_state(state)
    }

    /// Export state for serialization
    pub fn export_state(&self) -> MicroLoraState {
        let adapters = self
            .adapters
            .iter()
            .map(|(module, adapter)| {
                let adapter = adapter.read();
                let state = LoraAdapterState {
                    lora_a: adapter.lora_a.iter().copied().collect(),
                    lora_b: adapter.lora_b.iter().copied().collect(),
                    in_features: adapter.lora_a.nrows(),
                    out_features: adapter.lora_b.ncols(),
                    rank: adapter.rank,
                    scaling: adapter.scaling,
                };
                (*module, state)
            })
            .collect();

        MicroLoraState {
            config: self.config.clone(),
            adapters,
            adaptations: self.adaptations.load(Ordering::Relaxed),
        }
    }

    /// Create from exported state
    pub fn from_state(state: MicroLoraState) -> Result<Self> {
        let mut adapters = HashMap::new();

        for (module, adapter_state) in state.adapters {
            let lora_a = Array2::from_shape_vec(
                (adapter_state.in_features, adapter_state.rank),
                adapter_state.lora_a,
            )
            .map_err(|e| RuvLLMError::Config(e.to_string()))?;

            let lora_b = Array2::from_shape_vec(
                (adapter_state.rank, adapter_state.out_features),
                adapter_state.lora_b,
            )
            .map_err(|e| RuvLLMError::Config(e.to_string()))?;

            let adapter = LoraAdapter {
                lora_a: lora_a.clone(),
                lora_b: lora_b.clone(),
                bias: None,
                scaling: adapter_state.scaling,
                grad_a: Array2::zeros(lora_a.dim()),
                grad_b: Array2::zeros(lora_b.dim()),
                grad_count: 0,
                rank: adapter_state.rank,
            };

            adapters.insert(module, Arc::new(RwLock::new(adapter)));
        }

        Ok(Self {
            config: state.config,
            adapters,
            adaptations: AtomicU64::new(state.adaptations),
            forward_count: AtomicU64::new(0),
            enabled: true,
        })
    }

    /// Get configuration
    pub fn config(&self) -> &MicroLoraConfig {
        &self.config
    }

    /// Get total number of adaptations
    pub fn adaptation_count(&self) -> u64 {
        self.adaptations.load(Ordering::Relaxed)
    }

    /// Get total forward passes
    pub fn forward_count(&self) -> u64 {
        self.forward_count.load(Ordering::Relaxed)
    }

    /// Get total parameter count
    pub fn param_count(&self) -> usize {
        self.adapters.values().map(|a| a.read().param_count()).sum()
    }

    /// Get total memory usage in bytes
    pub fn memory_bytes(&self) -> usize {
        self.adapters
            .values()
            .map(|a| a.read().memory_bytes())
            .sum()
    }

    /// Enable/disable adaptation
    pub fn set_enabled(&mut self, enabled: bool) {
        self.enabled = enabled;
    }

    /// Check if enabled
    pub fn is_enabled(&self) -> bool {
        self.enabled
    }

    /// Get adapter for a specific module
    pub fn get_adapter(&self, module: &TargetModule) -> Option<Arc<RwLock<LoraAdapter>>> {
        self.adapters.get(module).cloned()
    }

    /// Reset all adapters to initial state
    pub fn reset(&self) {
        for adapter in self.adapters.values() {
            adapter.write().reset();
        }
        self.adaptations.store(0, Ordering::Relaxed);
        self.forward_count.store(0, Ordering::Relaxed);
    }
}

/// EWC state for a single adapter
#[derive(Clone)]
pub struct EwcState {
    /// Fisher information for A matrix
    pub fisher_a: Array2<f32>,
    /// Fisher information for B matrix
    pub fisher_b: Array2<f32>,
    /// Optimal A weights (from previous task)
    pub optimal_a: Array2<f32>,
    /// Optimal B weights (from previous task)
    pub optimal_b: Array2<f32>,
}

impl EwcState {
    /// Create new EWC state from current adapter
    pub fn from_adapter(adapter: &LoraAdapter) -> Self {
        Self {
            fisher_a: Array2::zeros(adapter.lora_a.dim()),
            fisher_b: Array2::zeros(adapter.lora_b.dim()),
            optimal_a: adapter.lora_a.clone(),
            optimal_b: adapter.lora_b.clone(),
        }
    }

    /// Update Fisher information using gradient squares
    pub fn update_fisher(&mut self, grad_a: &Array2<f32>, grad_b: &Array2<f32>, decay: f32) {
        // EMA update: F_t = decay * F_{t-1} + (1 - decay) * g^2
        self.fisher_a.zip_mut_with(grad_a, |f, g| {
            *f = decay * *f + (1.0 - decay) * g * g;
        });
        self.fisher_b.zip_mut_with(grad_b, |f, g| {
            *f = decay * *f + (1.0 - decay) * g * g;
        });
    }

    /// Update optimal weights
    pub fn update_optimal(&mut self, adapter: &LoraAdapter) {
        self.optimal_a.assign(&adapter.lora_a);
        self.optimal_b.assign(&adapter.lora_b);
    }
}

/// Serializable state for MicroLoRA
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct MicroLoraState {
    /// Configuration
    pub config: MicroLoraConfig,
    /// Adapter states by module
    pub adapters: HashMap<TargetModule, LoraAdapterState>,
    /// Total adaptations
    pub adaptations: u64,
}

/// Serializable state for a single adapter
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct LoraAdapterState {
    /// Flattened A matrix
    pub lora_a: Vec<f32>,
    /// Flattened B matrix
    pub lora_b: Vec<f32>,
    /// Input features
    pub in_features: usize,
    /// Output features
    pub out_features: usize,
    /// Rank
    pub rank: usize,
    /// Scaling factor
    pub scaling: f32,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_micro_lora_creation() {
        let config = MicroLoraConfig::for_hidden_dim(256);
        let lora = MicroLoRA::new(config);

        assert_eq!(lora.config().rank, 2);
        assert!(lora.is_enabled());
        assert_eq!(lora.adapters.len(), 2); // q_proj and v_proj
    }

    #[test]
    fn test_adapter_forward() {
        let adapter = LoraAdapter::new(64, 64, 2, 4.0);
        let input = Array1::from_elem(64, 1.0);

        let output = adapter.forward(&input);
        assert_eq!(output.len(), 64);

        // With zero-initialized B, output should be zero
        let sum: f32 = output.iter().sum();
        assert!(sum.abs() < 1e-6);
    }

    #[test]
    fn test_adapter_gradient_accumulation() {
        let mut adapter = LoraAdapter::new(64, 64, 2, 4.0);
        let input = Array1::from_elem(64, 0.1);
        let grad_output = Array1::from_elem(64, 0.1);

        adapter.accumulate_gradient(&input, &grad_output, 0.8);
        assert_eq!(adapter.pending_updates(), 1);

        adapter.apply_gradients(0.01);
        assert_eq!(adapter.pending_updates(), 0);

        // After update, forward should produce non-zero output
        let output = adapter.forward(&input);
        let sum: f32 = output.iter().map(|x| x.abs()).sum();
        assert!(sum > 0.0);
    }

    #[test]
    fn test_micro_lora_adapt() {
        let config = MicroLoraConfig::for_hidden_dim(64);
        let lora = MicroLoRA::new(config);

        let input = vec![0.1; 64];
        let feedback = AdaptFeedback::from_quality(0.8);

        lora.adapt(&input, feedback).unwrap();
        assert_eq!(lora.adaptation_count(), 1);

        lora.apply_updates(0.01);

        // Forward should now produce non-zero output
        let output = lora.forward(&input, &TargetModule::QProj);
        let sum: f32 = output.iter().map(|x| x.abs()).sum();
        assert!(sum > 0.0);
    }

    #[test]
    fn test_config_memory_bytes() {
        let config = MicroLoraConfig {
            rank: 2,
            in_features: 768,
            out_features: 768,
            target_modules: vec![TargetModule::QProj, TargetModule::VProj],
            ..Default::default()
        };

        // 2 modules * (768 * 2 + 2 * 768) * 4 bytes = 2 * 3072 * 4 = 24576 bytes
        assert!(config.memory_bytes() < 1024 * 1024); // < 1MB
    }

    #[test]
    fn test_simd_forward() {
        let adapter = LoraAdapter::new(64, 64, 2, 4.0);
        let input = vec![0.1f32; 64];
        let mut output = vec![0.0f32; 64];

        adapter.forward_simd(&input, &mut output);

        // Compare with regular forward
        let input_array = Array1::from_vec(input.clone());
        let expected = adapter.forward(&input_array);

        for (o, e) in output.iter().zip(expected.iter()) {
            assert!((o - e).abs() < 1e-5);
        }
    }

    #[test]
    fn test_ewc_state() {
        let adapter = LoraAdapter::new(64, 64, 2, 4.0);
        let mut ewc = EwcState::from_adapter(&adapter);

        let grad_a = Array2::from_elem((64, 2), 0.1);
        let grad_b = Array2::from_elem((2, 64), 0.1);

        ewc.update_fisher(&grad_a, &grad_b, 0.9);

        // Fisher should be updated
        assert!(ewc.fisher_a.iter().any(|&f| f > 0.0));
        assert!(ewc.fisher_b.iter().any(|&f| f > 0.0));
    }
}