//! Evoformer module - attention over sequences and residue pairs.
//!
//! The Evoformer processes sequence and pair representations through
//! alternating attention mechanisms to capture evolutionary relationships.

use crate::encoder::{EMBEDDING_DIM, PairEmbedding, SequenceEmbedding};

/// Number of attention heads.
pub const NUM_HEADS: usize = 8;

/// Evoformer block for processing sequence and pair representations.
#[derive(Debug, Clone)]
pub struct EvoformerBlock {
    /// Attention head dimension
    head_dim: usize,
    /// Number of attention heads
    num_heads: usize,
    /// Query projection weights
    wq: Vec<Vec<f32>>,
    /// Key projection weights
    wk: Vec<Vec<f32>>,
    /// Value projection weights
    wv: Vec<Vec<f32>>,
    /// Output projection weights
    wo: Vec<Vec<f32>>,
}

impl Default for EvoformerBlock {
    fn default() -> Self {
        Self::new(NUM_HEADS)
    }
}

impl EvoformerBlock {
    /// Create a new Evoformer block.
    #[must_use]
    pub fn new(num_heads: usize) -> Self {
        use rand::SeedableRng;
        use rand_distr::{Distribution, Normal};

        let head_dim = EMBEDDING_DIM / num_heads;
        let mut rng = rand::rngs::StdRng::seed_from_u64(123);
        let normal = Normal::new(0.0_f32, 0.02).unwrap();

        let mut init_weights = |rows: usize, cols: usize| -> Vec<Vec<f32>> {
            (0..rows)
                .map(|_| (0..cols).map(|_| normal.sample(&mut rng)).collect())
                .collect()
        };

        Self {
            head_dim,
            num_heads,
            wq: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
            wk: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
            wv: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
            wo: init_weights(EMBEDDING_DIM, EMBEDDING_DIM),
        }
    }

    /// Apply row-wise (sequence) attention.
    pub fn row_attention(&self, seq_emb: &mut SequenceEmbedding, pair_emb: &PairEmbedding) {
        let seq_len = seq_emb.sequence_length;

        // Compute Q, K, V projections
        let queries = self.project(&seq_emb.embeddings, &self.wq);
        let keys = self.project(&seq_emb.embeddings, &self.wk);
        let values = self.project(&seq_emb.embeddings, &self.wv);

        // Multi-head attention
        let scale = 1.0 / (self.head_dim as f32).sqrt();

        for i in 0..seq_len {
            let mut output = vec![0.0_f32; EMBEDDING_DIM];

            for head in 0..self.num_heads {
                let head_start = head * self.head_dim;
                let head_end = head_start + self.head_dim;

                // Compute attention scores
                let mut scores = Vec::with_capacity(seq_len);
                for j in 0..seq_len {
                    let mut score = 0.0_f32;
                    for k in head_start..head_end {
                        score += queries[i][k] * keys[j][k];
                    }
                    // Add pair bias from pair representation
                    if let Some(pair) = pair_emb.get(i, j) {
                        score += pair.iter().take(self.head_dim).sum::<f32>() * 0.1;
                    }
                    scores.push(score * scale);
                }

                // Softmax
                let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                let exp_scores: Vec<f32> = scores.iter().map(|s| (s - max_score).exp()).collect();
                let sum_exp: f32 = exp_scores.iter().sum();
                let attention: Vec<f32> = exp_scores.iter().map(|e| e / sum_exp).collect();

                // Weighted sum of values
                for j in 0..seq_len {
                    for k in head_start..head_end {
                        output[k] += attention[j] * values[j][k];
                    }
                }
            }

            // Output projection and residual
            let projected = self.project_single(&output, &self.wo);
            for (k, &p) in projected.iter().enumerate() {
                seq_emb.embeddings[i][k] += p * 0.1; // Scaled residual
            }
        }
    }

    /// Apply column-wise attention (for MSA, simplified here).
    pub fn column_attention(&self, seq_emb: &mut SequenceEmbedding) {
        // Simplified: just apply layer norm
        seq_emb.layer_norm(1e-5);
    }

    /// Apply triangle attention for pair updates.
    pub fn triangle_attention(&self, pair_emb: &mut PairEmbedding) {
        let seq_len = pair_emb.sequence_length;

        // Simplified triangle multiplication
        for i in 0..seq_len {
            for j in 0..seq_len {
                let mut update = vec![0.0_f32; pair_emb.pair_dim];

                // Outgoing edges
                for k in 0..seq_len {
                    if let (Some(ik), Some(kj)) = (pair_emb.get(i, k), pair_emb.get(k, j)) {
                        for d in 0..pair_emb.pair_dim {
                            update[d] += ik[d] * kj[d] * 0.01;
                        }
                    }
                }

                // Apply update
                for d in 0..pair_emb.pair_dim {
                    pair_emb.embeddings[i][j][d] += update[d];
                }
            }
        }
    }

    /// Project embeddings through weight matrix.
    fn project(&self, embeddings: &[Vec<f32>], weights: &[Vec<f32>]) -> Vec<Vec<f32>> {
        embeddings
            .iter()
            .map(|emb| self.project_single(emb, weights))
            .collect()
    }

    /// Project a single embedding.
    fn project_single(&self, embedding: &[f32], weights: &[Vec<f32>]) -> Vec<f32> {
        let mut output = vec![0.0_f32; weights.len()];
        for (i, row) in weights.iter().enumerate() {
            for (j, &w) in row.iter().enumerate() {
                if j < embedding.len() {
                    output[i] += embedding[j] * w;
                }
            }
        }
        output
    }
}

/// Evoformer stack with multiple blocks.
#[derive(Debug, Clone)]
pub struct Evoformer {
    /// Stack of Evoformer blocks
    blocks: Vec<EvoformerBlock>,
}

impl Evoformer {
    /// Create a new Evoformer with specified number of blocks.
    #[must_use]
    pub fn new(num_blocks: usize) -> Self {
        let blocks = (0..num_blocks).map(|_| EvoformerBlock::default()).collect();
        Self { blocks }
    }

    /// Process sequence and pair representations through all blocks.
    pub fn forward(&self, seq_emb: &mut SequenceEmbedding, pair_emb: &mut PairEmbedding) {
        for block in &self.blocks {
            block.row_attention(seq_emb, pair_emb);
            block.column_attention(seq_emb);
            block.triangle_attention(pair_emb);
            pair_emb.outer_product_update(seq_emb);
        }
    }
}

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

    #[test]
    fn test_evoformer_block() {
        let block = EvoformerBlock::default();
        assert_eq!(block.num_heads, NUM_HEADS);
        assert_eq!(block.head_dim, EMBEDDING_DIM / NUM_HEADS);
    }

    #[test]
    fn test_evoformer_forward() {
        let encoder = ProteinEncoder::new();
        let mut seq_emb = encoder.encode("ACDEF");
        let mut pair_emb = encoder.encode_pairs(5);

        let evoformer = Evoformer::new(2);
        evoformer.forward(&mut seq_emb, &mut pair_emb);

        // Check that embeddings were modified
        assert_eq!(seq_emb.sequence_length, 5);
    }
}