use crate::generate::GenerationConfig;
use crate::layers::*;

#[test]
fn test_kvcache_creation() {
    let cache = KVCache::new(4, 512, 64).expect("test");
    assert_eq!(cache.num_layers(), 4);
    assert_eq!(cache.max_seq_len(), 512);
    assert_eq!(cache.head_dim(), 64);
    assert_eq!(cache.current_pos(), 0);
    assert!(!cache.is_full());
}

#[test]
fn test_kvcache_zero_params_error() {
    assert!(KVCache::new(0, 512, 64).is_err());
    assert!(KVCache::new(4, 0, 64).is_err());
    assert!(KVCache::new(4, 512, 0).is_err());
}

#[test]
fn test_kvcache_update_and_retrieve() {
    let mut cache = KVCache::new(2, 10, 4).expect("test");

    // Add first position
    let key = Tensor::from_vec(vec![4], vec![1.0, 2.0, 3.0, 4.0]).expect("test");
    let value = Tensor::from_vec(vec![4], vec![5.0, 6.0, 7.0, 8.0]).expect("test");

    cache.update(0, &key, &value).expect("test");
    cache.advance();

    // Retrieve and verify
    let cached_key = cache.get_key(0).expect("test");
    let cached_value = cache.get_value(0).expect("test");

    assert_eq!(cached_key.shape(), &[1, 4]);
    assert_eq!(cached_value.shape(), &[1, 4]);

    for i in 0..4 {
        assert!((cached_key.data()[i] - key.data()[i]).abs() < 1e-6);
        assert!((cached_value.data()[i] - value.data()[i]).abs() < 1e-6);
    }
}

#[test]
fn test_kvcache_multiple_positions() {
    let mut cache = KVCache::new(1, 10, 2).expect("test");

    // Add multiple positions
    for pos in 0..3 {
        #[allow(clippy::cast_precision_loss)]
        let base = pos as f32;
        let key = Tensor::from_vec(vec![2], vec![base, base + 0.5]).expect("test");
        let value = Tensor::from_vec(vec![2], vec![base + 1.0, base + 1.5]).expect("test");

        cache.update(0, &key, &value).expect("test");
        cache.advance();
    }

    assert_eq!(cache.current_pos(), 3);

    // Retrieve all positions
    let cached_key = cache.get_key(0).expect("test");
    let cached_value = cache.get_value(0).expect("test");

    assert_eq!(cached_key.shape(), &[3, 2]);
    assert_eq!(cached_value.shape(), &[3, 2]);

    // Verify first position
    assert!((cached_key.data()[0] - 0.0).abs() < 1e-6);
    assert!((cached_key.data()[1] - 0.5).abs() < 1e-6);
    // Verify second position
    assert!((cached_key.data()[2] - 1.0).abs() < 1e-6);
    assert!((cached_key.data()[3] - 1.5).abs() < 1e-6);
}

#[test]
fn test_kvcache_multiple_layers() {
    let mut cache = KVCache::new(2, 10, 4).expect("test");

    let key0 = Tensor::from_vec(vec![4], vec![1.0; 4]).expect("test");
    let value0 = Tensor::from_vec(vec![4], vec![2.0; 4]).expect("test");
    let key1 = Tensor::from_vec(vec![4], vec![3.0; 4]).expect("test");
    let value1 = Tensor::from_vec(vec![4], vec![4.0; 4]).expect("test");

    cache.update(0, &key0, &value0).expect("test");
    cache.update(1, &key1, &value1).expect("test");
    cache.advance();

    // Verify layer 0
    let layer0_key = cache.get_key(0).expect("test");
    assert!((layer0_key.data()[0] - 1.0).abs() < 1e-6);

    // Verify layer 1
    let layer1_key = cache.get_key(1).expect("test");
    assert!((layer1_key.data()[0] - 3.0).abs() < 1e-6);
}

#[test]
fn test_kvcache_layer_out_of_bounds_error() {
    let mut cache = KVCache::new(2, 10, 4).expect("test");
    let key = Tensor::from_vec(vec![4], vec![1.0; 4]).expect("test");
    let value = Tensor::from_vec(vec![4], vec![2.0; 4]).expect("test");

    // Update layer 2 (out of bounds)
    assert!(cache.update(2, &key, &value).is_err());

    // Get layer 2 (out of bounds)
    assert!(cache.get_key(2).is_err());
    assert!(cache.get_value(2).is_err());
}

#[test]
fn test_kvcache_size_mismatch_error() {
    let mut cache = KVCache::new(1, 10, 4).expect("test");

    // Wrong key size
    let key = Tensor::from_vec(vec![3], vec![1.0; 3]).expect("test");
    let value = Tensor::from_vec(vec![4], vec![2.0; 4]).expect("test");
    assert!(cache.update(0, &key, &value).is_err());

    // Wrong value size
    let key = Tensor::from_vec(vec![4], vec![1.0; 4]).expect("test");
    let value = Tensor::from_vec(vec![3], vec![2.0; 3]).expect("test");
    assert!(cache.update(0, &key, &value).is_err());
}

#[test]
fn test_kvcache_full_error() {
    let mut cache = KVCache::new(1, 2, 4).expect("test");
    let key = Tensor::from_vec(vec![4], vec![1.0; 4]).expect("test");
    let value = Tensor::from_vec(vec![4], vec![2.0; 4]).expect("test");

    // Fill cache
    cache.update(0, &key, &value).expect("test");
    cache.advance();
    cache.update(0, &key, &value).expect("test");
    cache.advance();

    assert!(cache.is_full());

    // Try to add more
    assert!(cache.update(0, &key, &value).is_err());
}

#[test]
fn test_kvcache_clear() {
    let mut cache = KVCache::new(1, 10, 4).expect("test");
    let key = Tensor::from_vec(vec![4], vec![1.0; 4]).expect("test");
    let value = Tensor::from_vec(vec![4], vec![2.0; 4]).expect("test");

    cache.update(0, &key, &value).expect("test");
    cache.advance();
    assert_eq!(cache.current_pos(), 1);

    cache.clear();
    assert_eq!(cache.current_pos(), 0);
    assert!(!cache.is_full());
}

#[test]
fn test_kvcache_empty_retrieval() {
    let cache = KVCache::new(1, 10, 4).expect("test");

    // Retrieve from empty cache
    let cached_key = cache.get_key(0).expect("test");
    let cached_value = cache.get_value(0).expect("test");

    // Should return [1, 4] tensor with zeros
    assert_eq!(cached_key.shape(), &[1, 4]);
    assert_eq!(cached_value.shape(), &[1, 4]);
    for &val in cached_key.data() {
        assert!((val - 0.0).abs() < 1e-6);
    }
}

// TransformerBlock tests

#[test]
fn test_transformer_block_creation() {
    let block = TransformerBlock::new(64, 4, 256, 1e-5).expect("test");
    assert_eq!(block.hidden_dim(), 64);
}

#[test]
fn test_transformer_block_zero_params_error() {
    // Zero hidden_dim
    assert!(TransformerBlock::new(0, 4, 256, 1e-5).is_err());
    // Zero num_heads
    assert!(TransformerBlock::new(64, 0, 256, 1e-5).is_err());
    // Zero intermediate_dim
    assert!(TransformerBlock::new(64, 4, 0, 1e-5).is_err());
}

#[test]
fn test_transformer_block_head_divisibility_error() {
    // 63 not divisible by 4
    assert!(TransformerBlock::new(63, 4, 256, 1e-5).is_err());
}

#[test]
fn test_transformer_block_forward_shape() {
    // Use num_heads=1 so head_dim=hidden_dim (simplified single-head attention)
    let block = TransformerBlock::new(8, 1, 32, 1e-5).expect("test");

    // Input: [2, 8] (seq_len=2, hidden_dim=8)
    let input = Tensor::from_vec(vec![2, 8], vec![0.1; 16]).expect("test");
    let output = block.forward(&input).expect("test");

    // Output should have same shape
    assert_eq!(output.shape(), &[2, 8]);
    assert_eq!(output.data().len(), 16);
}

#[test]
fn test_transformer_block_forward_valid_output() {
    let block = TransformerBlock::new(4, 1, 16, 1e-5).expect("test");

    // Input must be 2D [seq_len, hidden_dim]
    let input = Tensor::from_vec(vec![1, 4], vec![1.0, 2.0, 3.0, 4.0]).expect("test");
    let output = block.forward(&input).expect("test");

    // Output should be finite
    for &val in output.data() {
        assert!(val.is_finite(), "Output contains non-finite values");
    }
}

#[test]
fn test_transformer_block_residual_connection() {
    let block = TransformerBlock::new(4, 1, 16, 1e-5).expect("test");

    // Input must be 2D [seq_len, hidden_dim]
    // With zero input, output should be non-zero due to residual + processing
    let input = Tensor::from_vec(vec![1, 4], vec![0.0, 0.0, 0.0, 0.0]).expect("test");
    let output = block.forward(&input).expect("test");

    // Even with zero input, layer norm and attention should produce non-zero output
    // (though it might be small due to normalization)
    assert_eq!(output.shape(), &[1, 4]);
}

#[test]
fn test_transformer_block_shape_mismatch_error() {
    let block = TransformerBlock::new(8, 1, 32, 1e-5).expect("test");

    // Wrong hidden_dim (input has 4, block expects 8)
    let input = Tensor::from_vec(vec![4], vec![1.0; 4]).expect("test");
    let result = block.forward(&input);
    assert!(result.is_err());
}

#[test]
fn test_transformer_block_mutable_access() {
    let mut block = TransformerBlock::new(4, 1, 16, 1e-5).expect("test");

    // Verify we can access mutable references
    let _attn_norm = block.attn_norm_mut();
    let _attention = block.attention_mut();
    let _ffn_norm = block.ffn_norm_mut();
    let _ffn = block.ffn_mut();
}

// Embedding tests

#[test]
fn test_embedding_creation() {
    let embed = Embedding::new(1000, 64).expect("test");
    assert_eq!(embed.vocab_size(), 1000);
    assert_eq!(embed.embed_dim(), 64);
}

#[test]
fn test_embedding_zero_params_error() {
    assert!(Embedding::new(0, 64).is_err());
    assert!(Embedding::new(1000, 0).is_err());
}

#[test]
fn test_embedding_forward_shape() {
    let embed = Embedding::new(100, 8).expect("test");

    let token_ids = vec![0, 1, 2];
    let output = embed.forward(&token_ids).expect("test");

    assert_eq!(output.shape(), &[3, 8]);
    assert_eq!(output.data().len(), 24);
}

#[test]
fn test_embedding_forward_lookup() {
    let mut embed = Embedding::new(10, 4).expect("test");

    // Set specific embedding for token 5
    let offset = 5 * 4;
    embed.weights_mut()[offset] = 1.0;
    embed.weights_mut()[offset + 1] = 2.0;
    embed.weights_mut()[offset + 2] = 3.0;
    embed.weights_mut()[offset + 3] = 4.0;

    let output = embed.forward(&[5]).expect("test");
    assert_eq!(output.shape(), &[1, 4]);
    assert!((output.data()[0] - 1.0).abs() < 1e-6);
    assert!((output.data()[1] - 2.0).abs() < 1e-6);
    assert!((output.data()[2] - 3.0).abs() < 1e-6);
    assert!((output.data()[3] - 4.0).abs() < 1e-6);
}

#[test]
fn test_embedding_out_of_bounds_error() {
    let embed = Embedding::new(10, 4).expect("test");
    assert!(embed.forward(&[10]).is_err()); // ID 10 is out of bounds
    assert!(embed.forward(&[100]).is_err());
}

#[test]
fn test_embedding_empty_input_error() {
    let embed = Embedding::new(10, 4).expect("test");
    assert!(embed.forward(&[]).is_err());
}

// Model tests

#[test]
fn test_model_creation() {
    let config = ModelConfig {
        vocab_size: 100,
        hidden_dim: 8,
        num_heads: 1,
        num_layers: 2,
        intermediate_dim: 32,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");
    assert_eq!(model.config().vocab_size, 100);
    assert_eq!(model.config().num_layers, 2);
}

#[test]
fn test_model_forward_shape() {
    let config = ModelConfig {
        vocab_size: 50,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    let token_ids = vec![0, 1, 2];
    let output = model.forward(&token_ids).expect("test");

    // Output should be [seq_len, vocab_size]
    assert_eq!(output.shape(), &[3, 50]);
}

#[test]
fn test_model_forward_valid_output() {
    let config = ModelConfig {
        vocab_size: 20,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    let output = model.forward(&[0, 1]).expect("test");

    // Output should be finite
    for &val in output.data() {
        assert!(val.is_finite(), "Output contains non-finite values");
    }
}

#[test]
fn test_model_num_parameters() {
    let config = ModelConfig {
        vocab_size: 100,
        hidden_dim: 8,
        num_heads: 1,
        num_layers: 2,
        intermediate_dim: 32,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    let params = model.num_parameters();
    assert!(params > 0);
    // Should be at least embedding + lm_head
    assert!(params >= 100 * 8 + 8 * 100);
}

#[test]
fn test_model_mutable_access() {
    let config = ModelConfig {
        vocab_size: 50,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let mut model = Model::new(config).expect("test");

    // Verify we can access mutable references
    let _embed = model.embedding_mut();
    let _blocks = model.blocks_mut();
    let _norm = model.final_norm_mut();
    let _head = model.lm_head_mut();
}

#[test]
fn test_model_generate_basic() {
    let config = ModelConfig {
        vocab_size: 20,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    let gen_config = GenerationConfig::greedy().with_max_tokens(5);
    let tokens = model.generate(&[0], &gen_config).expect("test");

    // Should have prompt + up to 5 generated tokens
    assert!(tokens.len() <= 6);
    assert!(!tokens.is_empty());
    // First token should be prompt
    assert_eq!(tokens[0], 0);
}

#[test]
fn test_model_generate_respects_max_tokens() {
    let config = ModelConfig {
        vocab_size: 10,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    let gen_config = GenerationConfig::greedy().with_max_tokens(3);
    let tokens = model.generate(&[0, 1], &gen_config).expect("test");

    // Should have 2 prompt + 3 generated = 5 max
    assert!(tokens.len() <= 5);
}

#[test]
fn test_model_generate_with_eos() {
    let config = ModelConfig {
        vocab_size: 10,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    // Set EOS token
    let gen_config = GenerationConfig::greedy()
        .with_max_tokens(100)
        .with_eos_token_id(5);

    let tokens = model.generate(&[0], &gen_config).expect("test");

    // Should stop before max_tokens if EOS is generated
    // (may or may not hit EOS depending on model weights)
    assert!(tokens.len() <= 101);
}

#[test]
fn test_model_generate_empty_prompt_error() {
    let config = ModelConfig {
        vocab_size: 10,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    let gen_config = GenerationConfig::greedy();
    let result = model.generate(&[], &gen_config);
    assert!(result.is_err());
}

#[test]
fn test_model_generate_deterministic_with_seed() {
    let config = ModelConfig {
        vocab_size: 20,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    // Same seed should give same results
    let gen_config = GenerationConfig::greedy()
        .with_max_tokens(5)
        .with_seed(12345);

    let tokens1 = model.generate(&[0], &gen_config).expect("test");
    let tokens2 = model.generate(&[0], &gen_config).expect("test");

    assert_eq!(tokens1, tokens2);
}

#[test]
fn test_model_generate_top_k() {
    let config = ModelConfig {
        vocab_size: 20,
        hidden_dim: 4,
        num_heads: 1,
        num_layers: 1,
        intermediate_dim: 16,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    let gen_config = GenerationConfig::top_k(5).with_max_tokens(3).with_seed(42);

    let tokens = model.generate(&[0], &gen_config).expect("test");

    // Should generate valid tokens
    assert!(tokens.len() <= 4);
    for &token in &tokens {
        assert!(token < 20);
    }
}

// MultiHeadAttention tests

#[test]
fn test_multi_head_attention_creation_mha() {
    // Standard Multi-Head Attention (num_kv_heads = num_heads)
    let mha = MultiHeadAttention::mha(64, 8).expect("test");
    assert_eq!(mha.num_heads(), 8);
    assert_eq!(mha.num_kv_heads(), 8);
    assert_eq!(mha.head_dim(), 8); // 64 / 8
    assert_eq!(mha.hidden_dim(), 64);
    assert!(mha.is_mha());
    assert!(!mha.is_mqa());
    assert!(!mha.is_gqa());
}

#[test]
fn test_multi_head_attention_creation_mqa() {
    // Multi-Query Attention (num_kv_heads = 1)
    let mqa = MultiHeadAttention::mqa(64, 8).expect("test");
    assert_eq!(mqa.num_heads(), 8);
    assert_eq!(mqa.num_kv_heads(), 1);
    assert_eq!(mqa.head_dim(), 8);
    assert_eq!(mqa.hidden_dim(), 64);
    assert!(mqa.is_mqa());
    assert!(!mqa.is_mha());
    assert!(!mqa.is_gqa());
}

#[test]
fn test_multi_head_attention_creation_gqa() {
    // Grouped-Query Attention (1 < num_kv_heads < num_heads)
    let gqa = MultiHeadAttention::gqa(64, 8, 2).expect("test");
    assert_eq!(gqa.num_heads(), 8);
    assert_eq!(gqa.num_kv_heads(), 2);
    assert_eq!(gqa.head_dim(), 8);
    assert_eq!(gqa.hidden_dim(), 64);
    assert!(gqa.is_gqa());
    assert!(!gqa.is_mha());
    assert!(!gqa.is_mqa());
}

#[test]
fn test_multi_head_attention_zero_hidden_dim_error() {
    let result = MultiHeadAttention::new(0, 8, 8);
    assert!(result.is_err());
}

#[test]
fn test_multi_head_attention_zero_num_heads_error() {
    let result = MultiHeadAttention::new(64, 0, 1);
    assert!(result.is_err());
}

#[test]
fn test_multi_head_attention_zero_num_kv_heads_error() {
    let result = MultiHeadAttention::new(64, 8, 0);
    assert!(result.is_err());
}

#[test]
fn test_multi_head_attention_kv_heads_too_large_error() {
    // num_kv_heads cannot be greater than num_heads
    let result = MultiHeadAttention::new(64, 8, 16);
    assert!(result.is_err());
}

#[test]
fn test_multi_head_attention_indivisible_error() {
    // 65 is not divisible by 8
    let result = MultiHeadAttention::new(65, 8, 8);
    assert!(result.is_err());
}

#[test]
fn test_multi_head_attention_heads_not_divisible_error() {
    // num_heads must be divisible by num_kv_heads
    let result = MultiHeadAttention::new(64, 8, 3);
    assert!(result.is_err());
}

#[test]
fn test_multi_head_attention_mha_forward() {
    // Standard MHA with 2 heads
    let mha = MultiHeadAttention::mha(8, 2).expect("test");

    // Input: [seq_len=2, hidden_dim=8]
    let input = Tensor::from_vec(
        vec![2, 8],
        vec![
            1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // token 1
            0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // token 2
        ],
    )
    .expect("test");

    let output = mha.forward(&input).expect("test");

    // Output should have same shape as input
    assert_eq!(output.shape(), &[2, 8]);
}

#[test]
fn test_multi_head_attention_mqa_forward() {
    // Multi-Query Attention with 2 heads (shared K/V)
    let mqa = MultiHeadAttention::mqa(8, 2).expect("test");

    // Input: [seq_len=2, hidden_dim=8]
    let input = Tensor::from_vec(
        vec![2, 8],
        vec![
            1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // token 1
            0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // token 2
        ],
    )
    .expect("test");

    let output = mqa.forward(&input).expect("test");

    // Output should have same shape as input
    assert_eq!(output.shape(), &[2, 8]);
}

#[test]
fn test_multi_head_attention_shape_validation() {
    let mha = MultiHeadAttention::mha(8, 2).expect("test");

    // Wrong number of dimensions (1D instead of 2D)
    let input_1d = Tensor::from_vec(vec![8], vec![1.0; 8]).expect("test");
    let result = mha.forward(&input_1d);
    assert!(result.is_err());

    // Wrong hidden dimension
    let input_wrong_dim = Tensor::from_vec(vec![2, 16], vec![1.0; 32]).expect("test");
    let result = mha.forward(&input_wrong_dim);
    assert!(result.is_err());
}

#[test]
fn test_multi_head_attention_mha_vs_mqa_shape_consistency() {
    // Both MHA and MQA should produce same output shape
    let mha = MultiHeadAttention::mha(16, 4).expect("test");
    let mqa = MultiHeadAttention::mqa(16, 4).expect("test");

    let input = Tensor::from_vec(vec![3, 16], vec![0.5; 48]).expect("test");

    let multi_head_output = mha.forward(&input).expect("test");
    let multi_query_output = mqa.forward(&input).expect("test");

    // Both should have same output shape
    assert_eq!(multi_head_output.shape(), &[3, 16]);
    assert_eq!(multi_query_output.shape(), &[3, 16]);
    assert_eq!(multi_head_output.shape(), multi_query_output.shape());
}

#[test]
fn test_multi_head_attention_single_head() {
    // Edge case: single head (equivalent to single attention)
    let mha = MultiHeadAttention::mha(8, 1).expect("test");

    let input = Tensor::from_vec(vec![2, 8], vec![0.5; 16]).expect("test");
    let output = mha.forward(&input).expect("test");

    assert_eq!(output.shape(), &[2, 8]);
}

#[test]
fn test_multi_head_attention_mqa_kv_sharing() {
    // MQA should work with larger number of heads
    let mqa = MultiHeadAttention::mqa(32, 8).expect("test");

    let input = Tensor::from_vec(vec![4, 32], vec![0.1; 128]).expect("test");
    let output = mqa.forward(&input).expect("test");

    assert_eq!(output.shape(), &[4, 32]);
}

#[test]
fn test_multi_head_attention_long_sequence() {
    // Test with longer sequence
    let mha = MultiHeadAttention::mha(16, 4).expect("test");

    // Sequence length = 10
    let input = Tensor::from_vec(vec![10, 16], vec![0.3; 160]).expect("test");
    let output = mha.forward(&input).expect("test");

    assert_eq!(output.shape(), &[10, 16]);
}

#[test]
fn test_multi_head_attention_mqa_memory_efficiency() {
    // MQA should still work correctly with shared K/V
    // This tests that the shared K/V logic is correct
    let mqa = MultiHeadAttention::mqa(64, 16).expect("test");

    // Small batch
    let input = Tensor::from_vec(vec![2, 64], vec![0.2; 128]).expect("test");
    let output = mqa.forward(&input).expect("test");

    assert_eq!(output.shape(), &[2, 64]);
    assert_eq!(output.data().len(), 128); // 2 * 64
}

#[test]
fn test_multi_head_attention_gqa_forward() {
    // Grouped-Query Attention with 8 heads, 2 KV heads (4 heads per group)
    let gqa = MultiHeadAttention::gqa(32, 8, 2).expect("test");

    // Input: [seq_len=3, hidden_dim=32]
    let input = Tensor::from_vec(vec![3, 32], vec![0.1; 96]).expect("test");

    let output = gqa.forward(&input).expect("test");

    // Output should have same shape as input
    assert_eq!(output.shape(), &[3, 32]);
}

#[test]
fn test_multi_head_attention_gqa_shape_consistency() {
    // MHA, MQA, and GQA should all produce same output shape
    let mha = MultiHeadAttention::mha(64, 8).expect("test");
    let mqa = MultiHeadAttention::mqa(64, 8).expect("test");
    let gqa = MultiHeadAttention::gqa(64, 8, 2).expect("test");

    let input = Tensor::from_vec(vec![4, 64], vec![0.5; 256]).expect("test");

    let multi_head_out = mha.forward(&input).expect("test");
    let multi_query_out = mqa.forward(&input).expect("test");
    let grouped_query_out = gqa.forward(&input).expect("test");

    // All should have same output shape
    assert_eq!(multi_head_out.shape(), &[4, 64]);
    assert_eq!(multi_query_out.shape(), &[4, 64]);
    assert_eq!(grouped_query_out.shape(), &[4, 64]);
    assert_eq!(multi_head_out.shape(), multi_query_out.shape());
    assert_eq!(multi_head_out.shape(), grouped_query_out.shape());
}

#[test]
fn test_multi_head_attention_gqa_different_group_sizes() {
    // Test GQA with different group sizes
    // 16 heads, 4 KV heads (4 heads per group)
    let gqa1 = MultiHeadAttention::gqa(128, 16, 4).expect("test");
    let input = Tensor::from_vec(vec![2, 128], vec![0.3; 256]).expect("test");
    let output1 = gqa1.forward(&input).expect("test");
    assert_eq!(output1.shape(), &[2, 128]);

    // 16 heads, 8 KV heads (2 heads per group)
    let gqa2 = MultiHeadAttention::gqa(128, 16, 8).expect("test");
    let output2 = gqa2.forward(&input).expect("test");
    assert_eq!(output2.shape(), &[2, 128]);
}

// ============================================================================
// Phase 3 Acceptance Tests (Refs REALIZAR-PERF-SPEC-001)
// ============================================================================

/// Phase 3 acceptance test: verify tok/s meets spec target
///
/// Per spec Phase 3 acceptance criteria:
/// ```rust,ignore
/// assert!(benchmark_tokens_per_second() >= 25.0);
/// ```
///
/// Note: This test uses a small test model to verify generation
/// throughput meets the baseline. Real phi-2 benchmarking requires
/// the actual model file and full optimization integration.
#[test]
#[ignore = "performance benchmark - run explicitly with --include-ignored"]
fn test_phase3_acceptance_tokens_per_second() {
    use crate::generate::GenerationConfig;
    use std::time::Instant;

    // Create baseline model configuration
    // The optimized components (Flash Attention v2, FusedLayerNormLinear)
    // show significant speedups individually - see companion tests:
    // - Flash Attention v2 SIMD: ~10x faster than parallel for small sequences
    // - FusedLayerNormLinear parallel: ~3.6x faster for large batches
    //
    // This test verifies the generation loop meets baseline throughput.
    // Full phi-2 integration requires wiring up optimized components.
    let config = ModelConfig {
        vocab_size: 100, // Small vocab for fast softmax
        hidden_dim: 64,  // Smaller hidden dimension
        num_heads: 4,    // Multiple heads
        num_layers: 2,   // Two transformer layers
        intermediate_dim: 128,
        eps: 1e-5,
    };
    let model = Model::new(config).expect("test");

    // Warmup run
    let prompt = vec![1, 5, 10];
    let gen_config = GenerationConfig::greedy().with_max_tokens(5);
    let _ = model.generate(&prompt, &gen_config).expect("test");

    // Benchmark: generate 20 tokens 10 times
    let tokens_per_run = 20;
    let num_runs = 10;
    let gen_config = GenerationConfig::greedy().with_max_tokens(tokens_per_run);

    let start = Instant::now();
    for _ in 0..num_runs {
        let _ = model.generate(&prompt, &gen_config).expect("test");
    }
    let elapsed = start.elapsed();

    let total_tokens = tokens_per_run * num_runs;
    let tok_per_sec = total_tokens as f64 / elapsed.as_secs_f64();

    // Phase 3 acceptance: ≥25 tok/s
    // With optimized components, this should be achievable.
    // The individual component tests show:
    // - Flash Attention v2: 87µs/iter
    // - FusedLayerNormLinear parallel: 2.9ms/iter for 32x256->512
    assert!(
        tok_per_sec >= 25.0,
        "Phase 3 acceptance FAILED: {:.1} tok/s < 25.0 tok/s target. \
         Note: Full optimization requires integrating Flash Attention v2 \
         and FusedLayerNormLinear into Model::forward()",
        tok_per_sec
    );

    // Report performance
    eprintln!(
        "Phase 3 acceptance PASSED: {:.1} tok/s (target: ≥25.0 tok/s)",
        tok_per_sec
    );
}

/// Test Flash Attention v2 + parallel performance improvement
#[test]
fn test_phase3_flash_attention_v2_performance() {
    use std::time::Instant;

    let head_dim = 64;
    let seq_len = 32;

    // Attention::new takes head_dim only
    let attn = Attention::new(head_dim).expect("test");

    // Create QKV tensors
    let q = Tensor::from_vec(vec![seq_len, head_dim], vec![0.1; seq_len * head_dim]).expect("test");
    let k = Tensor::from_vec(vec![seq_len, head_dim], vec![0.2; seq_len * head_dim]).expect("test");
    let v = Tensor::from_vec(vec![seq_len, head_dim], vec![0.3; seq_len * head_dim]).expect("test");

    // Warmup
    let _ = attn.flash_forward_v2(&q, &k, &v, 8).expect("test");
    let _ = attn.flash_forward_parallel(&q, &k, &v, 8).expect("test");

    // Benchmark Flash Attention v2 (SIMD)
    let iterations = 100;
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn.flash_forward_v2(&q, &k, &v, 8).expect("test");
    }
    let v2_time = start.elapsed();

    // Benchmark Flash Attention parallel
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn.flash_forward_parallel(&q, &k, &v, 8).expect("test");
    }
    let parallel_time = start.elapsed();

    // Report performance (informational)
    let v2_us = v2_time.as_micros() as f64 / iterations as f64;
    let parallel_us = parallel_time.as_micros() as f64 / iterations as f64;

    eprintln!(
        "Flash Attention v2: {:.2}us/iter, Parallel: {:.2}us/iter",
        v2_us, parallel_us
    );

    // Both implementations should complete without error
    // Performance comparison is informational
    assert!(v2_us > 0.0, "v2 should have measurable time");
    assert!(parallel_us > 0.0, "parallel should have measurable time");
}

/// Test FusedLayerNormLinear performance improvement
#[test]
fn test_phase3_fused_layernorm_linear_performance() {
    use std::time::Instant;

    let feature_dim = 256;
    let out_features = 512;
    let batch_size = 32;

    // FusedLayerNormLinear::new initializes with default weights
    // (norm_weight=1.0, norm_bias=0.0, linear_weight=0.0, linear_bias=0.0)
    // which is fine for performance testing
    let fused = FusedLayerNormLinear::new(feature_dim, out_features, 1e-5).expect("test");

    // Create input batch
    let input = Tensor::from_vec(
        vec![batch_size, feature_dim],
        vec![0.5; batch_size * feature_dim],
    )
    .expect("test");

    // Warmup
    let _ = fused.forward(&input).expect("test");
    let _ = fused.forward_parallel(&input).expect("test");

    // Benchmark fused forward
    let iterations = 100;
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = fused.forward(&input).expect("test");
    }
    let fused_time = start.elapsed();

    // Benchmark parallel fused forward
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = fused.forward_parallel(&input).expect("test");
    }
    let parallel_time = start.elapsed();

    // Report performance
    let fused_us = fused_time.as_micros() as f64 / iterations as f64;
    let parallel_us = parallel_time.as_micros() as f64 / iterations as f64;

    eprintln!(
        "FusedLayerNormLinear: {:.2}us/iter, Parallel: {:.2}us/iter",
        fused_us, parallel_us
    );

    // Verify performance is measurable
    assert!(fused_us > 0.0, "fused should have measurable time");
    assert!(parallel_us > 0.0, "parallel should have measurable time");
}

// =========================================================================
// BENCH-SPRINT-002: QuantizedLinear Tests (Q4_K Integration)
// Per benchmark-model-runners-spec.md v2.0: Inline dequantization for 8x
// memory bandwidth reduction vs f32.
// =========================================================================

/// RED: Test QuantizedLinear creation from Q4_K weight bytes
#[test]
fn test_quantized_linear_creation() {
    // Q4_K format: 144 bytes per super-block of 256 values
    // For in_features=256, out_features=4, we need 4 rows * 144 bytes = 576 bytes
    let in_features = 256;
    let out_features = 4;
    let bytes_per_row = 144; // One super-block per row
    let weight_bytes = vec![0u8; out_features * bytes_per_row];
    let bias = vec![0.0f32; out_features];

    let layer = QuantizedLinear::new(in_features, out_features, weight_bytes, bias);
    assert!(
        layer.is_ok(),
        "Should create QuantizedLinear from Q4_K bytes"
    );

    let layer = layer.expect("test");
    assert_eq!(layer.in_features(), in_features);
    assert_eq!(layer.out_features(), out_features);
}

/// RED: Test QuantizedLinear forward pass produces correct output
#[test]
fn test_quantized_linear_forward() {
    // Create test Q4_K weights (zeros for simplicity)
    let in_features = 256;
    let out_features = 4;
    let bytes_per_row = 144;
    let weight_bytes = vec![0u8; out_features * bytes_per_row];
    let bias = vec![1.0f32; out_features]; // Non-zero bias

    let layer = QuantizedLinear::new(in_features, out_features, weight_bytes, bias)
        .expect("Should create layer");

    // Input activations
    let input = Tensor::from_vec(vec![in_features], vec![1.0f32; in_features])
        .expect("Should create input");

    // Forward pass
    let output = layer.forward(&input).expect("Forward should work");

    // Output should have shape [out_features]
    assert_eq!(output.shape(), &[out_features]);

    // With zero weights and bias=1.0, output should be [1.0, 1.0, 1.0, 1.0]
    for &val in output.data() {
        assert!(
            (val - 1.0).abs() < 1e-5,
            "Output should equal bias with zero weights"
        );
    }
}

/// RED: Test QuantizedLinear forward with batch input
#[test]
fn test_quantized_linear_batch_forward() {
    let in_features = 256;
    let out_features = 4;
    let batch_size = 8;
    let bytes_per_row = 144;
    let weight_bytes = vec![0u8; out_features * bytes_per_row];
    let bias = vec![2.0f32; out_features];

    let layer = QuantizedLinear::new(in_features, out_features, weight_bytes, bias)
        .expect("Should create layer");

    // Batch input [batch_size, in_features]
    let input = Tensor::from_vec(
        vec![batch_size, in_features],
        vec![1.0f32; batch_size * in_features],
    )
    .expect("Should create batch input");

    let output = layer.forward(&input).expect("Batch forward should work");

    // Output should have shape [batch_size, out_features]
    assert_eq!(output.shape(), &[batch_size, out_features]);
}

/// RED: Test QuantizedLinear memory usage is ~8x less than Linear
#[test]
fn test_quantized_linear_memory_efficiency() {
    let in_features = 4096; // Realistic embedding dim
    let out_features = 4096;

    // f32 Linear: 4096 * 4096 * 4 bytes = 64MB
    let f32_bytes = in_features * out_features * std::mem::size_of::<f32>();

    // Q4_K: 4096/256 = 16 super-blocks per row, 16 * 144 = 2304 bytes/row
    // Total: 4096 * 2304 = ~9.4MB (6.8x reduction, close to theoretical 8x)
    let super_blocks_per_row = in_features.div_ceil(256);
    let q4k_bytes = out_features * super_blocks_per_row * 144;

    let ratio = f32_bytes as f64 / q4k_bytes as f64;

    // Q4_K should be at least 6x smaller than f32 (accounting for scale/min overhead)
    assert!(
        ratio > 6.0,
        "Q4_K should be >6x smaller than f32: ratio={}",
        ratio
    );
    eprintln!(
        "Memory efficiency: f32={} bytes, Q4_K={} bytes, ratio={:.2}x",
        f32_bytes, q4k_bytes, ratio
    );
}

// ========================
// SlidingWindowAttention Tests
// ========================

#[test]
fn test_sliding_window_attention_new() {
    let swa = SlidingWindowAttention::new(64, 4096).expect("test");
    assert_eq!(swa.head_dim(), 64);
    assert_eq!(swa.window_size(), 4096);
    assert!((swa.scale() - 0.125).abs() < 1e-6); // 1/sqrt(64) = 0.125
}

#[test]
fn test_sliding_window_attention_new_errors() {
    // Zero head_dim should error
    assert!(SlidingWindowAttention::new(0, 4096).is_err());
    // Zero window_size should error
    assert!(SlidingWindowAttention::new(64, 0).is_err());
}

#[test]
fn test_sliding_window_attention_forward_basic() {
    let swa = SlidingWindowAttention::new(4, 3).expect("test");
    // Small test: 5 positions, window size 3
    // Query: 5x4, Key: 5x4, Value: 5x4
    let query_data: Vec<f32> = (0..20).map(|i| i as f32 * 0.1).collect();
    let key_data: Vec<f32> = (0..20).map(|i| i as f32 * 0.1).collect();
    let value_data: Vec<f32> = (0..20).map(|i| (i % 4) as f32).collect();

    let query = Tensor::from_vec(vec![5, 4], query_data).expect("test");
    let key = Tensor::from_vec(vec![5, 4], key_data).expect("test");
    let value = Tensor::from_vec(vec![5, 4], value_data).expect("test");

    let output = swa.forward(&query, &key, &value).expect("test");
    assert_eq!(output.size(), 20); // 5 positions * 4 head_dim
}

#[test]
fn test_sliding_window_attention_causal_masking() {
    // Test that position i can only attend to positions <= i
    let swa = SlidingWindowAttention::new(2, 10).expect("test"); // Large window, so only causal matters
                                                                 // Query: 3x2, Key: 3x2, Value: 3x2
    let query = Tensor::from_vec(vec![3, 2], vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0]).expect("test");
    let key = Tensor::from_vec(vec![3, 2], vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0]).expect("test");
    let value = Tensor::from_vec(vec![3, 2], vec![1.0, 0.0, 0.0, 1.0, 0.5, 0.5]).expect("test");

    let output = swa.forward(&query, &key, &value).expect("test");
    assert_eq!(output.size(), 6);

    // Position 0 can only attend to itself
    // Position 1 can attend to positions 0,1
    // Position 2 can attend to positions 0,1,2
    // All positions produce valid outputs (not zeros)
    let data = output.data();
    assert!(data[0].abs() > 0.0 || data[1].abs() > 0.0);
}

#[test]
fn test_sliding_window_attention_window_boundary() {
    // Window size 2: each position can attend to at most 2 keys
    let swa = SlidingWindowAttention::new(2, 2).expect("test");
    // 5 positions, window=2
    let query = Tensor::from_vec(vec![5, 2], vec![1.0; 10]).expect("test");
    let key = Tensor::from_vec(vec![5, 2], vec![1.0; 10]).expect("test");
    let value_data: Vec<f32> = (0..10).map(|i| i as f32).collect();
    let value = Tensor::from_vec(vec![5, 2], value_data).expect("test");

    let output = swa.forward(&query, &key, &value).expect("test");
    assert_eq!(output.size(), 10);

    // Position 0: attends to [0] (only 1 key available due to causality)
    // Position 1: attends to [0,1] (2 keys)
    // Position 2: attends to [1,2] (window slides, excludes 0)
    // Position 3: attends to [2,3]
    // Position 4: attends to [3,4]
}

#[test]
fn test_sliding_window_attention_effective_context() {
    let swa = SlidingWindowAttention::new(64, 4).expect("test");

    // Position 0, seq_len 10: can attend to min(1, 4) = 1
    assert_eq!(swa.effective_context(0, 10), 1);

    // Position 3, seq_len 10: can attend to min(4, 4) = 4
    assert_eq!(swa.effective_context(3, 10), 4);

    // Position 7, seq_len 10: can attend to 4 (window kicks in)
    assert_eq!(swa.effective_context(7, 10), 4);

    // Position 2, seq_len 3: can attend to min(3, 4) = 3
    assert_eq!(swa.effective_context(2, 3), 3);
}

#[test]
fn test_sliding_window_attention_memory_ratio() {
    let swa = SlidingWindowAttention::new(64, 4096).expect("test");

    // For short sequences, ratio ~= 1.0
    let ratio_short = swa.memory_ratio(1000);
    assert!(
        ratio_short > 0.9,
        "Short sequences should use ~full attention"
    );

    // For long sequences, ratio approaches window_size / seq_len
    let ratio_long = swa.memory_ratio(100_000);
    let expected = 4096.0 / 100_000.0;
    assert!(
        (ratio_long - expected).abs() < 0.01,
        "Long sequences should use ~window_size/seq_len memory: got {}, expected {}",
        ratio_long,
        expected
    );
}

#[test]
fn test_sliding_window_attention_error_mismatched_kv() {
    let swa = SlidingWindowAttention::new(4, 3).expect("test");
    let query = Tensor::from_vec(vec![2, 4], vec![1.0; 8]).expect("test");
    let key = Tensor::from_vec(vec![3, 4], vec![1.0; 12]).expect("test");
    let value = Tensor::from_vec(vec![2, 4], vec![1.0; 8]).expect("test"); // Different from K

    // K and V must have same seq_len
    let result = swa.forward(&query, &key, &value);
    assert!(result.is_err());
}

#[test]
fn test_sliding_window_attention_error_bad_head_dim() {
    let swa = SlidingWindowAttention::new(4, 3).expect("test");
    // Key has wrong head_dim (3 instead of 4)
    let query = Tensor::from_vec(vec![2, 4], vec![1.0; 8]).expect("test");
    let key = Tensor::from_vec(vec![2, 3], vec![1.0; 6]).expect("test");
    let value = Tensor::from_vec(vec![2, 4], vec![1.0; 8]).expect("test");

    let result = swa.forward(&query, &key, &value);
    assert!(result.is_err());
}

#[test]
fn test_sliding_window_attention_bidirectional() {
    let swa = SlidingWindowAttention::new(2, 4).expect("test");
    // 5 positions, bidirectional window
    let query = Tensor::from_vec(vec![5, 2], vec![1.0; 10]).expect("test");
    let key = Tensor::from_vec(vec![5, 2], vec![1.0; 10]).expect("test");
    let value_data: Vec<f32> = (0..10).map(|i| i as f32).collect();
    let value = Tensor::from_vec(vec![5, 2], value_data).expect("test");

    let output_causal = swa.forward(&query, &key, &value).expect("test");
    let output_bidir = swa
        .forward_with_mask(&query, &key, &value, false)
        .expect("test");

    // Bidirectional can attend to more positions, so outputs may differ
    assert_eq!(output_causal.size(), output_bidir.size());
    // Both should produce valid outputs
    assert!(output_causal.data().iter().any(|&x| x.abs() > 0.0));
    assert!(output_bidir.data().iter().any(|&x| x.abs() > 0.0));
}

#[test]
fn test_sliding_window_attention_forward_with_mask_causal() {
    let swa = SlidingWindowAttention::new(2, 3).expect("test");
    let query = Tensor::from_vec(vec![3, 2], vec![1.0; 6]).expect("test");
    let key = Tensor::from_vec(vec![3, 2], vec![1.0; 6]).expect("test");
    let value = Tensor::from_vec(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("test");

    // forward_with_mask(causal=true) should match forward()
    let output_forward = swa.forward(&query, &key, &value).expect("test");
    let output_mask = swa
        .forward_with_mask(&query, &key, &value, true)
        .expect("test");

    for (a, b) in output_forward.data().iter().zip(output_mask.data().iter()) {
        assert!(
            (a - b).abs() < 1e-6,
            "Causal outputs should match: {} vs {}",
            a,
            b
        );
    }
}

#[test]
fn test_sliding_window_attention_single_token() {
    let swa = SlidingWindowAttention::new(4, 3).expect("test");
    // Single token input
    let query = Tensor::from_vec(vec![1, 4], vec![1.0, 2.0, 3.0, 4.0]).expect("test");
    let key = Tensor::from_vec(vec![1, 4], vec![1.0, 2.0, 3.0, 4.0]).expect("test");
    let value = Tensor::from_vec(vec![1, 4], vec![0.5, 0.5, 0.5, 0.5]).expect("test");

    let output = swa.forward(&query, &key, &value).expect("test");
    assert_eq!(output.size(), 4);
    // Self-attention on single token returns the value
    let data = output.data();
    for &v in data {
        assert!((v - 0.5).abs() < 1e-6);
    }
}

// ========================================================================
// IMP-003: Fused QKV + Attention Tests (EXTREME TDD - RED Phase)
// Per spec: performance-parity-ollama-llamacpp-gpu-inference-llms.md
// ========================================================================

#[test]
fn test_fused_qkv_attention_basic() {
    // IMP-003: Fused attention should match separate Q/K/V computation
    let fused = FusedQKVAttention::new(4, 64).expect("test");
    let input = Tensor::from_vec(vec![8, 64], vec![0.1; 8 * 64]).expect("test");

    let output = fused.forward(&input).expect("test");
    assert_eq!(output.shape(), &[8, 64]);
}

#[test]
fn test_fused_qkv_attention_correctness() {
    // Verify fused output matches separate computation within 4 ULPs
    let head_dim = 16;
    let hidden_dim = 64;
    let seq_len = 4;

    let fused = FusedQKVAttention::new(head_dim, hidden_dim).expect("test");
    let input = Tensor::from_vec(
        vec![seq_len, hidden_dim],
        (0..(seq_len * hidden_dim))
            .map(|i| (i as f32 * 0.01).sin())
            .collect(),
    )
    .expect("test");

    let output = fused.forward(&input).expect("test");

    // Output should have same shape as input
    assert_eq!(output.shape(), input.shape());

    // Values should be finite (no NaN/Inf)
    for &val in output.data() {
        assert!(val.is_finite(), "Output contains non-finite value: {}", val);
    }
}

#[test]
fn test_fused_qkv_attention_single_token() {
    // Single token case - important for autoregressive generation
    let fused = FusedQKVAttention::new(8, 32).expect("test");
    let input = Tensor::from_vec(vec![1, 32], vec![0.5; 32]).expect("test");

    let output = fused.forward(&input).expect("test");
    assert_eq!(output.shape(), &[1, 32]);
}

#[test]
fn test_fused_qkv_attention_error_zero_head_dim() {
    let result = FusedQKVAttention::new(0, 64);
    assert!(result.is_err());
}

#[test]
fn test_fused_qkv_attention_error_zero_hidden_dim() {
    let result = FusedQKVAttention::new(8, 0);
    assert!(result.is_err());
}

#[test]
fn test_fused_qkv_attention_error_mismatched_input() {
    let fused = FusedQKVAttention::new(8, 64).expect("test");
    // Input with wrong hidden dim
    let input = Tensor::from_vec(vec![4, 32], vec![0.1; 4 * 32]).expect("test");

    let result = fused.forward(&input);
    assert!(result.is_err());
}

#[test]
fn test_fused_qkv_attention_numerical_stability() {
    // Test with extreme values - should not produce NaN/Inf
    let fused = FusedQKVAttention::new(8, 32).expect("test");

    // Large values that could overflow naive softmax
    let input = Tensor::from_vec(vec![4, 32], vec![100.0; 4 * 32]).expect("test");

    let output = fused.forward(&input).expect("test");

    for &val in output.data() {
        assert!(
            val.is_finite(),
            "Large inputs caused non-finite output: {}",
            val
        );
    }

    // Small values that could underflow
    let input_small = Tensor::from_vec(vec![4, 32], vec![1e-10; 4 * 32]).expect("test");

    let output_small = fused.forward(&input_small).expect("test");

    for &val in output_small.data() {
        assert!(
            val.is_finite(),
            "Small inputs caused non-finite output: {}",
            val
        );
    }
}

#[test]
fn test_fused_qkv_attention_causal_mask() {
    // Causal attention: position i can only attend to positions <= i
    let fused = FusedQKVAttention::new(4, 16).expect("test");
    let input =
        Tensor::from_vec(vec![4, 16], (0..64).map(|i| (i as f32) * 0.1).collect()).expect("test");

    let output = fused.forward(&input).expect("test");

    // Each output position should only depend on prior positions
    // This is implicitly verified by the implementation using causal mask
    assert_eq!(output.shape(), &[4, 16]);
}

// ========================================================================
// QA Checklist Section A: Correctness Tests (QA-001 to QA-010)
// Per spec: performance-parity-ollama-llamacpp-gpu-inference-llms.md §5
// ========================================================================

/// QA-003: Attention scores match reference implementation within tolerance
#[test]
fn test_qa_003_attention_scores_correctness() {
    let head_dim = 4;
    let attention = Attention::new(head_dim).expect("test");

    // Create simple Q, K, V tensors for verification
    let q = Tensor::from_vec(
        vec![2, head_dim],
        vec![1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
    )
    .expect("test");
    let k = q.clone();
    let v = Tensor::from_vec(
        vec![2, head_dim],
        vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
    )
    .expect("test");

    let output = attention.forward(&q, &k, &v).expect("test");

    // Output should have correct shape
    assert_eq!(output.shape(), &[2, head_dim]);

    // Attention with identical Q and K should weight values appropriately
    // Position 0 can only attend to position 0 (causal)
    // Position 1 can attend to both positions
    let data = output.data();
    for &val in data {
        assert!(val.is_finite(), "QA-003: Attention output should be finite");
    }
}

/// QA-004: RoPE embeddings produce correct rotations
#[test]
fn test_qa_004_rope_embeddings_correctness() {
    let rope = RoPE::new(64, 10000.0).expect("test");

    // Apply RoPE at position 0 - should be identity-like
    let input = Tensor::from_vec(vec![1, 64], vec![1.0; 64]).expect("test");
    let output_pos0 = rope.forward(&input, 0).expect("test");

    // Apply at position 1 - should be rotated
    let output_pos1 = rope.forward(&input, 1).expect("test");

    // Outputs at different positions should differ
    let data0 = output_pos0.data();
    let data1 = output_pos1.data();

    let mut differs = false;
    for (a, b) in data0.iter().zip(data1.iter()) {
        if (a - b).abs() > 1e-6 {
            differs = true;
            break;
        }
    }
    assert!(
        differs,
        "QA-004: RoPE should produce different outputs at different positions"
    );

    // All outputs should be finite
    for &val in data0 {
        assert!(val.is_finite(), "QA-004: RoPE output should be finite");
    }
}

/// QA-005: Softmax outputs sum to 1.0 within tolerance
#[test]
fn test_qa_005_softmax_sum_to_one() {
    // Various input sizes
    for size in [4, 16, 64, 256] {
        let input = Tensor::from_vec(
            vec![size],
            (0..size).map(|i| (i as f32 * 0.1).sin()).collect(),
        )
        .expect("test");

        let output = softmax(&input).expect("test");
        let sum: f32 = output.data().iter().sum();

        assert!(
            (sum - 1.0).abs() < 1e-5,
            "QA-005: Softmax sum should be 1.0, got {} for size {}",
            sum,
            size
        );

        // All values should be positive
        for &val in output.data() {
            assert!(val >= 0.0, "QA-005: Softmax outputs should be non-negative");
            assert!(val <= 1.0, "QA-005: Softmax outputs should be <= 1.0");
        }
    }
}

/// QA-006: Layer norm outputs have unit variance within tolerance
#[test]
fn test_qa_006_layer_norm_unit_variance() {
    let hidden_dim = 64;
    let layer_norm = LayerNorm::new(hidden_dim, 1e-5).expect("test");

    // Create input with known statistics
    let input = Tensor::from_vec(
        vec![1, hidden_dim],
        (0..hidden_dim).map(|i| i as f32).collect(),
    )
    .expect("test");

    let output = layer_norm.forward(&input).expect("test");
    let data = output.data();

    // Calculate variance of output
    let mean: f32 = data.iter().sum::<f32>() / (hidden_dim as f32);
    let variance: f32 = data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / (hidden_dim as f32);

    // Mean should be near 0 (before gamma/beta adjustment)
    // Variance should be near 1 (normalized)
    assert!(
        mean.abs() < 0.1,
        "QA-006: Layer norm mean should be near 0, got {}",
        mean
    );

    // Note: variance may differ due to gamma/beta, but should be reasonable
    assert!(
        variance > 0.0 && variance < 10.0,
        "QA-006: Layer norm variance should be bounded, got {}",
        variance
    );
}

/// QA-007: GELU activation matches expected behavior
#[test]
fn test_qa_007_gelu_activation_correctness() {
    // GELU(0) ≈ 0
    let input_zero = Tensor::from_vec(vec![1], vec![0.0]).expect("test");
    let output_zero = gelu(&input_zero).expect("test");
    assert!(
        output_zero.data()[0].abs() < 1e-5,
        "QA-007: GELU(0) should be ~0, got {}",
        output_zero.data()[0]
    );

    // GELU(x) > 0 for x > 0
    let input_pos = Tensor::from_vec(vec![1], vec![1.0]).expect("test");
    let output_pos = gelu(&input_pos).expect("test");
    assert!(
        output_pos.data()[0] > 0.0,
        "QA-007: GELU(1.0) should be positive"
    );

    // GELU is approximately linear for large x
    let input_large = Tensor::from_vec(vec![1], vec![10.0]).expect("test");
    let output_large = gelu(&input_large).expect("test");
    assert!(
        (output_large.data()[0] - 10.0).abs() < 1.0,
        "QA-007: GELU(10) should be ~10"
    );

    // GELU(x) < 0 for small negative x but bounded
    let input_neg = Tensor::from_vec(vec![1], vec![-0.5]).expect("test");
    let output_neg = gelu(&input_neg).expect("test");
    assert!(
        output_neg.data()[0] < 0.0 && output_neg.data()[0] > -1.0,
        "QA-007: GELU(-0.5) should be small negative"
    );
}

/// QA-009: KV cache produces identical results to recomputation
#[test]
fn test_qa_009_kv_cache_correctness() {
    use crate::inference::KVCache;

    let num_layers = 2;
    let hidden_dim = 64;
    let max_seq_len = 32;

    let mut cache = KVCache::new(num_layers, hidden_dim, max_seq_len);

    // Store K and V values for layer 0
    let k_data: Vec<f32> = (0..hidden_dim).map(|i| i as f32 * 0.1).collect();
    let v_data: Vec<f32> = (0..hidden_dim).map(|i| i as f32 * 0.2).collect();

    cache.store(0, &k_data, &v_data);
    cache.advance();

    // Store more values
    let k_data2: Vec<f32> = (0..hidden_dim).map(|i| i as f32 * 0.3).collect();
    let v_data2: Vec<f32> = (0..hidden_dim).map(|i| i as f32 * 0.4).collect();
    cache.store(0, &k_data2, &v_data2);
    cache.advance();

    // Retrieve and verify
    let k_out = cache.get_k(0);
    let v_out = cache.get_v(0);

    // Should have 2 positions worth of data
    assert_eq!(
        k_out.len(),
        2 * hidden_dim,
        "QA-009: K cache should contain 2 positions"
    );
    assert_eq!(
        v_out.len(),
        2 * hidden_dim,
        "QA-009: V cache should contain 2 positions"
    );

    // First position values should match first stored data
    for i in 0..hidden_dim {
        assert!(
            (k_out[i] - k_data[i]).abs() < 1e-6,
            "QA-009: K cache position 0 should match stored value at index {}",
            i
        );
        assert!(
            (v_out[i] - v_data[i]).abs() < 1e-6,
            "QA-009: V cache position 0 should match stored value at index {}",
            i
        );
    }

    // Second position values should match second stored data
    for i in 0..hidden_dim {
        assert!(
            (k_out[hidden_dim + i] - k_data2[i]).abs() < 1e-6,
            "QA-009: K cache position 1 should match stored value at index {}",
            i
        );
    }
}

/// QA-010: Quantized inference matches F32 within acceptable tolerance
#[test]
fn test_qa_010_quantized_vs_f32_tolerance() {
    use crate::quantize::{dequantize_q4_k, dequantize_q8_0};

    // Q8_0 block format: 2 bytes scale (f16) + 32 bytes quants = 34 bytes
    // Note: Q8_0 block size is 34 bytes per GGML/GGUF spec
    let mut q8_data = vec![0u8; 34]; // 1 block = 34 bytes
                                     // scale = 1.0 (f16 = 0x3C00)
    q8_data[0..2].copy_from_slice(&0x3C00_u16.to_le_bytes());
    // quants = 0..31 (signed i8, stored as u8)
    for i in 0..32 {
        q8_data[2 + i] = i as u8; // quants start at offset 2
    }

    let dequant = dequantize_q8_0(&q8_data).expect("test");
    assert_eq!(
        dequant.len(),
        32,
        "QA-010: Q8_0 should produce 32 values per block"
    );

    // All values should be finite
    for &val in &dequant {
        assert!(
            val.is_finite(),
            "QA-010: Q8_0 dequantized values should be finite"
        );
    }

    // Q4_K should be within reasonable tolerance
    let mut q4k_data = vec![0u8; 144]; // 1 super-block
                                       // d = 1.0, dmin = 0.0
    q4k_data[0..2].copy_from_slice(&0x3C00_u16.to_le_bytes());
    q4k_data[2..4].copy_from_slice(&0x0000_u16.to_le_bytes());

    let q4k_dequant = dequantize_q4_k(&q4k_data).expect("test");
    assert_eq!(
        q4k_dequant.len(),
        256,
        "QA-010: Q4_K should produce 256 values per super-block"
    );

    // All values should be finite
    for &val in &q4k_dequant {
        assert!(
            val.is_finite(),
            "QA-010: Dequantized values should be finite"
        );
    }
}

// ========================================================================
// QA Checklist Section B: Performance Tests (QA-011 to QA-020)
// Per spec: performance-parity-ollama-llamacpp-gpu-inference-llms.md §5
// ========================================================================
