optirs-learned 0.3.1

use std::fmt::Debug;
// Multi-head attention mechanisms for transformer-based optimization
//
// This module implements the attention mechanisms used in the transformer optimizer,
// including multi-head attention, relative position bias, and rotary position embeddings.

#[allow(dead_code)]
use scirs2_core::ndarray::{s, Array, Array1, Array2, Array3};
use scirs2_core::numeric::Float;
use scirs2_core::random::{Random, Rng as SCRRng};
use std::collections::HashMap;

use super::super::TransformerOptimizerConfig;
use crate::error::{OptimError, Result};

/// Attention optimization strategies
#[derive(Debug, Clone, Copy)]
pub enum AttentionOptimization {
    /// Standard full attention
    Full,
    /// Sparse attention patterns
    Sparse,
    /// Linear attention approximation
    Linear,
    /// Local attention windows
    Local,
    /// Hierarchical attention
    Hierarchical,
    /// Adaptive attention sparsity
    Adaptive,
}

/// Multi-head attention mechanism for transformer optimizer
#[derive(Debug, Clone)]
pub struct MultiHeadAttention<T: Float + Debug + Send + Sync + 'static> {
    /// Query, Key, Value projection weights
    wq: Array2<T>,
    wk: Array2<T>,
    wv: Array2<T>,

    /// Output projection weights
    wo: Array2<T>,

    /// Number of attention heads
    numheads: usize,

    /// Head dimension
    head_dim: usize,

    /// Model dimension
    modeldim: usize,

    /// Attention optimization strategy
    optimization: AttentionOptimization,

    /// Relative position bias (if enabled)
    relative_bias: Option<RelativePositionBias<T>>,

    /// Attention scores from last forward pass
    attentionscores: Option<Array3<T>>,

    /// Attention weights from last forward pass
    attention_weights: Option<Array3<T>>,

    /// RoPE embeddings (if enabled)
    rope_embeddings: Option<RoPEEmbeddings<T>>,
}

/// Relative position bias for attention
#[derive(Debug, Clone)]
pub struct RelativePositionBias<T: Float + Debug + Send + Sync + 'static> {
    /// Bias table
    bias_table: Array2<T>,

    /// Maximum relative distance
    max_distance: usize,

    /// Cached position indices
    position_indices: Option<Array2<usize>>,
}

/// Rotary Position Embedding (RoPE)
#[derive(Debug, Clone)]
pub struct RoPEEmbeddings<T: Float + Debug + Send + Sync + 'static> {
    /// Cosine values
    cos_cached: Array2<T>,

    /// Sine values
    sin_cached: Array2<T>,

    /// Maximum sequence length
    max_seqlen: usize,

    /// Dimension
    dim: usize,
}

impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> MultiHeadAttention<T> {
    pub fn new(config: &TransformerOptimizerConfig) -> Result<Self> {
        let modeldim = config.modeldim;
        let numheads = config.numheads;
        let head_dim = modeldim / numheads;

        if !modeldim.is_multiple_of(numheads) {
            return Err(OptimError::InvalidConfig(
                "Model dimension must be divisible by number of heads".to_string(),
            ));
        }

        let mut rng = scirs2_core::random::thread_rng();

        // Initialize projection weights
        let bound = (6.0 / (2 * modeldim) as f64).sqrt();

        let mut wq = Array2::zeros((modeldim, modeldim));
        let mut wk = Array2::zeros((modeldim, modeldim));
        let mut wv = Array2::zeros((modeldim, modeldim));
        let mut wo = Array2::zeros((modeldim, modeldim));

        for elem in wq.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound).expect("unwrap failed");
        }
        for elem in wk.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound).expect("unwrap failed");
        }
        for elem in wv.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound).expect("unwrap failed");
        }
        for elem in wo.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound).expect("unwrap failed");
        }

        let relative_bias = if config.relative_position_bias {
            Some(RelativePositionBias::new(
                config.max_sequence_length,
                numheads,
            )?)
        } else {
            None
        };

        let rope_embeddings = if config.use_rope {
            Some(RoPEEmbeddings::new(config.max_sequence_length, head_dim)?)
        } else {
            None
        };

        Ok(Self {
            wq,
            wk,
            wv,
            wo,
            numheads,
            head_dim,
            modeldim,
            optimization: config.attention_optimization,
            relative_bias,
            attentionscores: None,
            attention_weights: None,
            rope_embeddings,
        })
    }

    pub fn forward(
        &mut self,
        query: &Array2<T>,
        key: &Array2<T>,
        value: &Array2<T>,
    ) -> Result<Array2<T>> {
        let (_seq_len, modeldim) = query.dim();

        if modeldim != self.modeldim {
            return Err(OptimError::InvalidConfig(format!(
                "Model dimension {} doesn't match expected {}",
                modeldim, self.modeldim
            )));
        }

        // Project to Q, K, V
        let q = self.linear_transform(query, &self.wq)?;
        let k = self.linear_transform(key, &self.wk)?;
        let v = self.linear_transform(value, &self.wv)?;

        // Reshape for multi-head attention
        let q_heads = self.reshape_for_heads(&q)?;
        let k_heads = self.reshape_for_heads(&k)?;
        let v_heads = self.reshape_for_heads(&v)?;

        // Compute attention
        let attention_output = self.compute_attention(&q_heads, &k_heads, &v_heads)?;

        // Reshape back and apply output projection
        let concat_output = self.reshape_from_heads(&attention_output)?;
        let final_output = self.linear_transform(&concat_output, &self.wo)?;

        Ok(final_output)
    }

    fn linear_transform(&self, input: &Array2<T>, weights: &Array2<T>) -> Result<Array2<T>> {
        // Matrix multiplication: input @ weights.T
        let mut output = Array2::zeros((input.nrows(), weights.nrows()));
        for i in 0..input.nrows() {
            for j in 0..weights.nrows() {
                let mut sum = T::zero();
                for k in 0..input.ncols() {
                    sum = sum + input[[i, k]] * weights[[j, k]];
                }
                output[[i, j]] = sum;
            }
        }
        Ok(output)
    }

    fn reshape_for_heads(&self, input: &Array2<T>) -> Result<Array3<T>> {
        let (seq_len, _) = input.dim();
        let mut reshaped = Array3::zeros((self.numheads, seq_len, self.head_dim));

        for h in 0..self.numheads {
            for s in 0..seq_len {
                for d in 0..self.head_dim {
                    let input_idx = h * self.head_dim + d;
                    reshaped[[h, s, d]] = input[[s, input_idx]];
                }
            }
        }

        Ok(reshaped)
    }

    fn reshape_from_heads(&self, input: &Array3<T>) -> Result<Array2<T>> {
        let (_numheads, seq_len, _head_dim) = input.dim();
        let mut reshaped = Array2::zeros((seq_len, self.modeldim));

        for h in 0..self.numheads {
            for s in 0..seq_len {
                for d in 0..self.head_dim {
                    let output_idx = h * self.head_dim + d;
                    reshaped[[s, output_idx]] = input[[h, s, d]];
                }
            }
        }

        Ok(reshaped)
    }

    fn compute_attention(
        &mut self,
        q: &Array3<T>,
        k: &Array3<T>,
        v: &Array3<T>,
    ) -> Result<Array3<T>> {
        let (_numheads, seq_len, head_dim) = q.dim();
        let scale = T::from(1.0 / (head_dim as f64).sqrt()).expect("unwrap failed");

        let mut attention_output = Array3::zeros((self.numheads, seq_len, self.head_dim));
        let mut attention_scores = Array3::zeros((self.numheads, seq_len, seq_len));

        // Compute attention for each head
        for h in 0..self.numheads {
            // Compute scaled dot-product attention
            for i in 0..seq_len {
                for j in 0..seq_len {
                    let mut score = T::zero();
                    for d in 0..self.head_dim {
                        score = score + q[[h, i, d]] * k[[h, j, d]];
                    }
                    attention_scores[[h, i, j]] = score * scale;
                }
            }

            // Apply relative position bias if enabled
            if let Some(ref bias) = self.relative_bias {
                bias.apply_bias(&mut attention_scores.slice_mut(s![h, .., ..]))?;
            }

            // Softmax over last dimension
            self.apply_softmax(&mut attention_scores.slice_mut(s![h, .., ..]))?;

            // Apply attention to values
            for i in 0..seq_len {
                for d in 0..self.head_dim {
                    let mut output = T::zero();
                    for j in 0..seq_len {
                        output = output + attention_scores[[h, i, j]] * v[[h, j, d]];
                    }
                    attention_output[[h, i, d]] = output;
                }
            }
        }

        // Cache attention scores for analysis
        self.attentionscores = Some(attention_scores);

        Ok(attention_output)
    }

    fn apply_softmax(&self, scores: &mut scirs2_core::ndarray::ArrayViewMut2<T>) -> Result<()> {
        let (rows, cols) = scores.dim();

        for i in 0..rows {
            // Find max for numerical stability
            let mut max_val = scores[[i, 0]];
            for j in 1..cols {
                if scores[[i, j]] > max_val {
                    max_val = scores[[i, j]];
                }
            }

            // Compute exponentials and sum
            let mut sum = T::zero();
            for j in 0..cols {
                let exp_val = (scores[[i, j]] - max_val).exp();
                scores[[i, j]] = exp_val;
                sum = sum + exp_val;
            }

            // Normalize
            for j in 0..cols {
                scores[[i, j]] = scores[[i, j]] / sum;
            }
        }

        Ok(())
    }

    /// Get attention patterns for analysis
    pub fn get_attention_patterns(&self) -> Option<&Array3<T>> {
        self.attentionscores.as_ref()
    }

    /// Get number of attention heads
    pub fn num_heads(&self) -> usize {
        self.numheads
    }

    /// Get head dimension
    pub fn head_dim(&self) -> usize {
        self.head_dim
    }
}

impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> RelativePositionBias<T> {
    pub fn new(max_distance: usize, num_heads: usize) -> Result<Self> {
        let mut rng = scirs2_core::random::thread_rng();
        let table_size = 2 * max_distance - 1;

        let mut bias_table = Array2::zeros((table_size, num_heads));
        let bound = 0.02;

        for elem in bias_table.iter_mut() {
            *elem = T::from((rng.random::<f64>() - 0.5) * 2.0 * bound).expect("unwrap failed");
        }

        Ok(Self {
            bias_table,
            max_distance,
            position_indices: None,
        })
    }

    pub fn apply_bias(&self, scores: &mut scirs2_core::ndarray::ArrayViewMut2<T>) -> Result<()> {
        let (seq_len, _) = scores.dim();

        // Simple bias application - in practice would be more sophisticated
        for i in 0..seq_len {
            for j in 0..seq_len {
                let rel_pos = (i as i32 - j as i32).unsigned_abs() as usize;
                let bias_idx = rel_pos.min(self.max_distance - 1);
                scores[[i, j]] = scores[[i, j]] + self.bias_table[[bias_idx, 0]];
            }
        }

        Ok(())
    }
}

impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> RoPEEmbeddings<T> {
    pub fn new(max_seqlen: usize, dim: usize) -> Result<Self> {
        if !dim.is_multiple_of(2) {
            return Err(OptimError::InvalidConfig(
                "RoPE dimension must be even".to_string(),
            ));
        }

        let mut cos_cached = Array2::zeros((max_seqlen, dim));
        let mut sin_cached = Array2::zeros((max_seqlen, dim));

        let base = 10000.0;

        for pos in 0..max_seqlen {
            for i in (0..dim).step_by(2) {
                let theta = pos as f64 / base.powf((i as f64) / (dim as f64));
                cos_cached[[pos, i]] = T::from(theta.cos()).expect("unwrap failed");
                cos_cached[[pos, i + 1]] = T::from(theta.cos()).expect("unwrap failed");
                sin_cached[[pos, i]] = T::from(theta.sin()).expect("unwrap failed");
                sin_cached[[pos, i + 1]] = T::from(theta.sin()).expect("unwrap failed");
            }
        }

        Ok(Self {
            cos_cached,
            sin_cached,
            max_seqlen,
            dim,
        })
    }

    pub fn apply_rope(&self, x: &mut Array2<T>, positions: &[usize]) -> Result<()> {
        for (seq_idx, &pos) in positions.iter().enumerate() {
            if pos >= self.max_seqlen {
                return Err(OptimError::InvalidConfig(
                    "Position exceeds maximum sequence length".to_string(),
                ));
            }

            for i in (0..self.dim).step_by(2) {
                let x_i = x[[seq_idx, i]];
                let x_i_plus_1 = x[[seq_idx, i + 1]];

                let cos_val = self.cos_cached[[pos, i]];
                let sin_val = self.sin_cached[[pos, i]];

                x[[seq_idx, i]] = x_i * cos_val - x_i_plus_1 * sin_val;
                x[[seq_idx, i + 1]] = x_i * sin_val + x_i_plus_1 * cos_val;
            }
        }

        Ok(())
    }
}