//! Phase 34: Matmul coverage tests
//!
//! These lib tests illuminate gguf/inference/matmul.rs:
//! - embed() - Token embedding lookup
//! - embed_into() - Pre-allocated embedding lookup
//! - fused_matmul() - Quantized matrix multiplication
//!
//! Uses proptest for empirical saturation across quantization types.
//!
//! Located in lib tests to be included in `cargo test --lib` coverage

use crate::gguf::test_helpers::create_test_model_with_config;
use crate::gguf::types::{
    GGUF_TYPE_Q4_0, GGUF_TYPE_Q4_1, GGUF_TYPE_Q4_K, GGUF_TYPE_Q5_0, GGUF_TYPE_Q5_K, GGUF_TYPE_Q6_K,
    GGUF_TYPE_Q8_0,
};
use crate::gguf::{GGUFConfig, OwnedQuantizedTensor};

// =============================================================================
// Embed Tests
// =============================================================================

#[test]
fn test_phase34_embed_single_token() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let embeddings = model.embed(&[0]);

    assert_eq!(embeddings.len(), config.hidden_dim);
    assert!(embeddings.iter().all(|x| x.is_finite()));
}

#[test]
fn test_phase34_embed_multiple_tokens() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let embeddings = model.embed(&[0, 1, 2, 3, 4]);

    assert_eq!(embeddings.len(), 5 * config.hidden_dim);
    assert!(embeddings.iter().all(|x| x.is_finite()));
}

#[test]
fn test_phase34_embed_out_of_vocab() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    // Token 9999 is beyond vocab size 100
    let embeddings = model.embed(&[9999]);

    // Should return zeros for out-of-vocab tokens
    assert_eq!(embeddings.len(), config.hidden_dim);
    assert!(embeddings.iter().all(|&x| x == 0.0));
}

#[test]
fn test_phase34_embed_boundary_token() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    // Token 99 is the last valid token (vocab_size - 1)
    let embeddings = model.embed(&[99]);

    assert_eq!(embeddings.len(), config.hidden_dim);
    assert!(embeddings.iter().all(|x| x.is_finite()));
}

#[test]
fn test_phase34_embed_empty() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let embeddings = model.embed(&[]);

    assert!(embeddings.is_empty());
}

// =============================================================================
// Embed Into Tests
// =============================================================================

#[test]
fn test_phase34_embed_into_single() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let mut output = vec![0.0f32; config.hidden_dim];

    model.embed_into(0, &mut output);

    // Output should be filled with embedding values
    assert_eq!(output.len(), config.hidden_dim);
    assert!(output.iter().all(|x| x.is_finite()));
}

#[test]
fn test_phase34_embed_into_out_of_vocab() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let mut output = vec![999.0f32; config.hidden_dim]; // Pre-fill with non-zero

    model.embed_into(9999, &mut output);

    // Should be zeroed for out-of-vocab
    assert!(output.iter().all(|&x| x == 0.0));
}

// =============================================================================
// Fused Matmul Q4_0 Tests
// =============================================================================

fn create_q4_0_weight(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    // Q4_0: 18 bytes per 32 elements
    let num_elements = in_dim * out_dim;
    let num_blocks = num_elements.div_ceil(32);
    let byte_size = num_blocks * 18;

    // Create valid Q4_0 data
    let mut data = Vec::with_capacity(byte_size);
    for _ in 0..num_blocks {
        // f16 scale (2 bytes)
        let scale = half::f16::from_f32(0.1);
        data.extend_from_slice(&scale.to_le_bytes());
        // 16 bytes of quants (32 4-bit values)
        data.extend([0x88u8; 16]); // Mid-range values
    }

    OwnedQuantizedTensor {
        data,
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q4_0,
    }
}

#[test]
fn test_phase34_fused_matmul_q4_0_single_seq() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let input = vec![1.0f32; in_dim]; // Single sequence position
    let weight = create_q4_0_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q4_0 fused_matmul failed: {:?}",
        result.err()
    );

    let output = result.unwrap();
    assert_eq!(output.len(), out_dim);
    assert!(output.iter().all(|x| x.is_finite()));
}

#[test]
fn test_phase34_fused_matmul_q4_0_multi_seq() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let seq_len = 4;
    let input = vec![1.0f32; in_dim * seq_len]; // 4 sequence positions
    let weight = create_q4_0_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q4_0 multi-seq fused_matmul failed: {:?}",
        result.err()
    );

    let output = result.unwrap();
    assert_eq!(output.len(), out_dim * seq_len);
    assert!(output.iter().all(|x| x.is_finite()));
}

// =============================================================================
// Fused Matmul Q8_0 Tests
// =============================================================================

fn create_q8_0_weight(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    // Q8_0: 34 bytes per 32 elements (2 f16 scale + 32 i8 quants)
    let num_elements = in_dim * out_dim;
    let num_blocks = num_elements.div_ceil(32);
    let byte_size = num_blocks * 34;

    let mut data = Vec::with_capacity(byte_size);
    for _ in 0..num_blocks {
        // f16 scale (2 bytes)
        let scale = half::f16::from_f32(0.1);
        data.extend_from_slice(&scale.to_le_bytes());
        // 32 i8 quants
        data.extend([0i8 as u8; 32]);
    }

    OwnedQuantizedTensor {
        data,
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q8_0,
    }
}

#[test]
fn test_phase34_fused_matmul_q8_0_single_seq() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let input = vec![1.0f32; in_dim];
    let weight = create_q8_0_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q8_0 fused_matmul failed: {:?}",
        result.err()
    );

    let output = result.unwrap();
    assert_eq!(output.len(), out_dim);
}

#[test]
fn test_phase34_fused_matmul_q8_0_multi_seq() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let seq_len = 3;
    let input = vec![0.5f32; in_dim * seq_len];
    let weight = create_q8_0_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), out_dim * seq_len);
}

// =============================================================================
// Fused Matmul Q4_1 Tests
// =============================================================================

fn create_q4_1_weight(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    // Q4_1: 20 bytes per 32 elements (2 scale + 2 min + 16 quants)
    let num_elements = in_dim * out_dim;
    let num_blocks = num_elements.div_ceil(32);
    let byte_size = num_blocks * 20;

    let mut data = Vec::with_capacity(byte_size);
    for _ in 0..num_blocks {
        let scale = half::f16::from_f32(0.1);
        let min = half::f16::from_f32(0.0);
        data.extend_from_slice(&scale.to_le_bytes());
        data.extend_from_slice(&min.to_le_bytes());
        data.extend([0x00u8; 16]);
    }

    OwnedQuantizedTensor {
        data,
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q4_1,
    }
}

#[test]
fn test_phase34_fused_matmul_q4_1() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let input = vec![1.0f32; in_dim];
    let weight = create_q4_1_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q4_1 fused_matmul failed: {:?}",
        result.err()
    );
    assert_eq!(result.unwrap().len(), out_dim);
}

#[test]
fn test_phase34_fused_matmul_q4_1_multi_seq() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let seq_len = 2;
    let input = vec![1.0f32; in_dim * seq_len];
    let weight = create_q4_1_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), out_dim * seq_len);
}

// =============================================================================
// Fused Matmul Q5_0 Tests
// =============================================================================

fn create_q5_0_weight(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    // Q5_0: 22 bytes per 32 elements (2 scale + 4 high bits + 16 quants)
    let num_elements = in_dim * out_dim;
    let num_blocks = num_elements.div_ceil(32);
    let byte_size = num_blocks * 22;

    let mut data = Vec::with_capacity(byte_size);
    for _ in 0..num_blocks {
        let scale = half::f16::from_f32(0.1);
        data.extend_from_slice(&scale.to_le_bytes());
        data.extend([0u8; 4]); // high bits
        data.extend([0x00u8; 16]); // quants
    }

    OwnedQuantizedTensor {
        data,
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q5_0,
    }
}

#[test]
fn test_phase34_fused_matmul_q5_0() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let input = vec![1.0f32; in_dim];
    let weight = create_q5_0_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q5_0 fused_matmul failed: {:?}",
        result.err()
    );
    assert_eq!(result.unwrap().len(), out_dim);
}

#[test]
fn test_phase34_fused_matmul_q5_0_multi_seq() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let seq_len = 2;
    let input = vec![1.0f32; in_dim * seq_len];
    let weight = create_q5_0_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), out_dim * seq_len);
}

// =============================================================================
// Fused Matmul K-Quant Tests (Q4_K, Q5_K, Q6_K)
// =============================================================================

fn create_q4_k_weight(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    // Q4_K: 144 bytes per 256 elements
    let num_elements = in_dim * out_dim;
    let num_super_blocks = num_elements.div_ceil(256);
    let byte_size = num_super_blocks * 144;

    OwnedQuantizedTensor {
        data: vec![0u8; byte_size],
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q4_K,
    }
}

fn create_q5_k_weight(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    // Q5_K: 176 bytes per 256 elements
    let num_elements = in_dim * out_dim;
    let num_super_blocks = num_elements.div_ceil(256);
    let byte_size = num_super_blocks * 176;

    OwnedQuantizedTensor {
        data: vec![0u8; byte_size],
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q5_K,
    }
}

fn create_q6_k_weight(in_dim: usize, out_dim: usize) -> OwnedQuantizedTensor {
    // Q6_K: 210 bytes per 256 elements
    let num_elements = in_dim * out_dim;
    let num_super_blocks = num_elements.div_ceil(256);
    let byte_size = num_super_blocks * 210;

    OwnedQuantizedTensor {
        data: vec![0u8; byte_size],
        in_dim,
        out_dim,
        qtype: GGUF_TYPE_Q6_K,
    }
}

#[test]
fn test_phase34_fused_matmul_q4_k_single() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 256,
        intermediate_dim: 512,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 256;
    let out_dim = 512;
    let input = vec![1.0f32; in_dim];
    let weight = create_q4_k_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q4_K fused_matmul failed: {:?}",
        result.err()
    );
    assert_eq!(result.unwrap().len(), out_dim);
}

#[test]
fn test_phase34_fused_matmul_q4_k_multi() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 256,
        intermediate_dim: 512,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 256;
    let out_dim = 512;
    let seq_len = 2;
    let input = vec![1.0f32; in_dim * seq_len];
    let weight = create_q4_k_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), out_dim * seq_len);
}

#[test]
fn test_phase34_fused_matmul_q5_k() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 256,
        intermediate_dim: 512,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 256;
    let out_dim = 512;
    let input = vec![1.0f32; in_dim];
    let weight = create_q5_k_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q5_K fused_matmul failed: {:?}",
        result.err()
    );
}

#[test]
fn test_phase34_fused_matmul_q5_k_multi() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 256,
        intermediate_dim: 512,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 256;
    let out_dim = 512;
    let seq_len = 2;
    let input = vec![1.0f32; in_dim * seq_len];
    let weight = create_q5_k_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), out_dim * seq_len);
}

#[test]
fn test_phase34_fused_matmul_q6_k() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 256,
        intermediate_dim: 512,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 256;
    let out_dim = 512;
    let input = vec![1.0f32; in_dim];
    let weight = create_q6_k_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(
        result.is_ok(),
        "Q6_K fused_matmul failed: {:?}",
        result.err()
    );
}

#[test]
fn test_phase34_fused_matmul_q6_k_multi() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 256,
        intermediate_dim: 512,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 256;
    let out_dim = 512;
    let seq_len = 2;
    let input = vec![1.0f32; in_dim * seq_len];
    let weight = create_q6_k_weight(in_dim, out_dim);

    let result = model.fused_matmul(&input, &weight);
    assert!(result.is_ok());
}

// =============================================================================
// Unsupported QType Test
// =============================================================================

#[test]
fn test_phase34_fused_matmul_unsupported_qtype() {
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);

    let in_dim = 64;
    let out_dim = 128;
    let input = vec![1.0f32; in_dim];

    // Use a fake/unsupported qtype
    let weight = OwnedQuantizedTensor {
        data: vec![0u8; 1024],
        in_dim,
        out_dim,
        qtype: 255, // Invalid
    };

    let result = model.fused_matmul(&input, &weight);
    assert!(result.is_err());
    let err = result.unwrap_err().to_string();
    assert!(
        err.contains("255") || err.contains("Unsupported") || err.contains("supports"),
        "Error: {}",
        err
    );
}

// =============================================================================
// Edge Cases
// =============================================================================

#[test]
fn test_phase34_embed_different_dimensions() {
    for hidden_dim in [32, 64, 128, 256] {
        let config = GGUFConfig {
            architecture: "llama".to_string(),
            hidden_dim,
            intermediate_dim: hidden_dim * 2,
            num_layers: 1,
            num_heads: hidden_dim / 16,
            num_kv_heads: hidden_dim / 16,
            vocab_size: 100,
            context_length: 256,
            rope_theta: 10000.0,
            eps: 1e-5,
            rope_type: 0,
            bos_token_id: None,
        };

        let model = create_test_model_with_config(&config);
        let embeddings = model.embed(&[0, 1]);

        assert_eq!(embeddings.len(), 2 * hidden_dim);
    }
}

#[test]
fn test_phase34_fused_matmul_all_qtypes_comprehensive() {
    // Comprehensive test covering all qtypes in a single test
    let config = GGUFConfig {
        architecture: "llama".to_string(),
        hidden_dim: 256,
        intermediate_dim: 512,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let in_dim = 256;
    let out_dim = 512;
    let input = vec![1.0f32; in_dim];

    // Test each qtype
    let qtypes = [
        ("Q4_0", create_q4_0_weight(in_dim, out_dim)),
        ("Q8_0", create_q8_0_weight(in_dim, out_dim)),
        ("Q4_1", create_q4_1_weight(in_dim, out_dim)),
        ("Q5_0", create_q5_0_weight(in_dim, out_dim)),
        ("Q4_K", create_q4_k_weight(in_dim, out_dim)),
        ("Q5_K", create_q5_k_weight(in_dim, out_dim)),
        ("Q6_K", create_q6_k_weight(in_dim, out_dim)),
    ];

    for (name, weight) in qtypes {
        let result = model.fused_matmul(&input, &weight);
        assert!(
            result.is_ok(),
            "{} fused_matmul failed: {:?}",
            name,
            result.err()
        );
        let output = result.unwrap();
        assert_eq!(output.len(), out_dim, "{} output size mismatch", name);
        assert!(
            output.iter().all(|x| x.is_finite()),
            "{} produced non-finite values",
            name
        );
    }
}
