//! Part 08: GPU Adapter Extended Coverage Tests
//!
//! Additional tests for `gpu/adapters/` module covering:
//! - AprF32ToGpuAdapter full conversion path
//! - AprToGpuAdapter Q4 full conversion
//! - MockExecutor integration with adapters
//! - GpuModel creation from adapters
//! - Error edge cases
//! - Backend dispatch paths

use crate::apr_transformer::{
    AprTransformer, AprTransformerConfig, AprTransformerLayer, QuantizedAprLayerQ4,
    QuantizedAprTensorQ4, QuantizedAprTransformerQ4,
};
use crate::gpu::adapters::{transpose_matrix, AprF32ToGpuAdapter, AprToGpuAdapter};
use crate::gpu::executor::MockExecutor;
use crate::gpu::scheduler::GpuGenerateConfig;

// ============================================================================
// Test Helper Functions
// ============================================================================

fn create_minimal_apr_config() -> AprTransformerConfig {
    AprTransformerConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        num_layers: 2,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        intermediate_dim: 128,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
    }
}

fn create_gqa_apr_config() -> AprTransformerConfig {
    AprTransformerConfig {
        architecture: "llama_gqa".to_string(),
        hidden_dim: 64,
        num_layers: 2,
        num_heads: 8,
        num_kv_heads: 2, // GQA: 4 Q heads per KV head
        vocab_size: 100,
        intermediate_dim: 128,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
    }
}

fn create_f32_layer(config: &AprTransformerConfig) -> AprTransformerLayer {
    let hidden_dim = config.hidden_dim;
    let num_heads = config.num_heads;
    let num_kv_heads = config.num_kv_heads;
    let intermediate_dim = config.intermediate_dim;
    let head_dim = hidden_dim / num_heads;
    let kv_dim = num_kv_heads * head_dim;
    let qkv_out_dim = hidden_dim + 2 * kv_dim;

    AprTransformerLayer {
        attn_norm_weight: vec![1.0; hidden_dim],
        attn_norm_bias: None,
        qkv_weight: vec![0.01; hidden_dim * qkv_out_dim],
        qkv_bias: None,
        attn_output_weight: vec![0.01; hidden_dim * hidden_dim],
        attn_output_bias: None,
        ffn_norm_weight: Some(vec![1.0; hidden_dim]),
        ffn_norm_bias: None,
        ffn_up_weight: vec![0.01; hidden_dim * intermediate_dim],
        ffn_up_bias: None,
        ffn_down_weight: vec![0.01; intermediate_dim * hidden_dim],
        ffn_down_bias: None,
        ffn_gate_weight: Some(vec![0.01; hidden_dim * intermediate_dim]),
        ffn_gate_bias: None,
    }
}

fn create_minimal_f32_apr() -> AprTransformer {
    let config = create_minimal_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let num_layers = config.num_layers;

    AprTransformer {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: (0..num_layers).map(|_| create_f32_layer(&config)).collect(),
        output_norm_weight: vec![1.0; hidden_dim],
        output_norm_bias: None,
        lm_head_weight: vec![0.01; vocab_size * hidden_dim],
        lm_head_bias: None,
        q4k_layers: None,
        lm_head_weight_q6k: None,
        lm_head_weight_q4k: None,
    }
}

fn create_gqa_f32_apr() -> AprTransformer {
    let config = create_gqa_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let num_layers = config.num_layers;

    AprTransformer {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: (0..num_layers).map(|_| create_f32_layer(&config)).collect(),
        output_norm_weight: vec![1.0; hidden_dim],
        output_norm_bias: None,
        lm_head_weight: vec![0.01; vocab_size * hidden_dim],
        lm_head_bias: None,
        q4k_layers: None,
        lm_head_weight_q6k: None,
        lm_head_weight_q4k: None,
    }
}

fn create_q4_layer(config: &AprTransformerConfig) -> QuantizedAprLayerQ4 {
    let hidden_dim = config.hidden_dim;
    let num_heads = config.num_heads;
    let num_kv_heads = config.num_kv_heads;
    let intermediate_dim = config.intermediate_dim;
    let head_dim = hidden_dim / num_heads;
    let kv_dim = num_kv_heads * head_dim;
    let qkv_out_dim = hidden_dim + 2 * kv_dim;

    QuantizedAprLayerQ4 {
        attn_norm_weight: vec![1.0; hidden_dim],
        qkv_weight: QuantizedAprTensorQ4::zeros(hidden_dim, qkv_out_dim),
        attn_output_weight: QuantizedAprTensorQ4::zeros(hidden_dim, hidden_dim),
        ffn_up_weight: QuantizedAprTensorQ4::zeros(hidden_dim, intermediate_dim),
        ffn_down_weight: QuantizedAprTensorQ4::zeros(intermediate_dim, hidden_dim),
        ffn_gate_weight: Some(QuantizedAprTensorQ4::zeros(hidden_dim, intermediate_dim)),
        ffn_norm_weight: Some(vec![1.0; hidden_dim]),
    }
}

fn create_q4_layer_without_gate(config: &AprTransformerConfig) -> QuantizedAprLayerQ4 {
    let hidden_dim = config.hidden_dim;
    let num_heads = config.num_heads;
    let num_kv_heads = config.num_kv_heads;
    let intermediate_dim = config.intermediate_dim;
    let head_dim = hidden_dim / num_heads;
    let kv_dim = num_kv_heads * head_dim;
    let qkv_out_dim = hidden_dim + 2 * kv_dim;

    QuantizedAprLayerQ4 {
        attn_norm_weight: vec![1.0; hidden_dim],
        qkv_weight: QuantizedAprTensorQ4::zeros(hidden_dim, qkv_out_dim),
        attn_output_weight: QuantizedAprTensorQ4::zeros(hidden_dim, hidden_dim),
        ffn_up_weight: QuantizedAprTensorQ4::zeros(hidden_dim, intermediate_dim),
        ffn_down_weight: QuantizedAprTensorQ4::zeros(intermediate_dim, hidden_dim),
        ffn_gate_weight: None,
        ffn_norm_weight: Some(vec![1.0; hidden_dim]),
    }
}

fn create_minimal_q4_apr() -> QuantizedAprTransformerQ4 {
    let config = create_minimal_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let num_layers = config.num_layers;

    QuantizedAprTransformerQ4 {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: (0..num_layers).map(|_| create_q4_layer(&config)).collect(),
        output_norm_weight: vec![1.0; hidden_dim],
        lm_head_weight: QuantizedAprTensorQ4::zeros(hidden_dim, vocab_size),
    }
}

fn create_q4_apr_without_gate() -> QuantizedAprTransformerQ4 {
    let config = create_minimal_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;
    let num_layers = config.num_layers;

    QuantizedAprTransformerQ4 {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: (0..num_layers)
            .map(|_| create_q4_layer_without_gate(&config))
            .collect(),
        output_norm_weight: vec![1.0; hidden_dim],
        lm_head_weight: QuantizedAprTensorQ4::zeros(hidden_dim, vocab_size),
    }
}

// ============================================================================
// AprF32ToGpuAdapter Comprehensive Tests
// ============================================================================

#[test]
fn test_apr_f32_to_gpu_basic_conversion() {
    let apr = create_minimal_f32_apr();
    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);

    assert!(result.is_ok(), "F32 conversion should succeed");
    let gpu_model = result.unwrap();

    // Verify config is correctly transferred
    assert_eq!(gpu_model.config.vocab_size, 100);
    assert_eq!(gpu_model.config.hidden_dim, 64);
    assert_eq!(gpu_model.config.num_heads, 4);
    assert_eq!(gpu_model.config.num_kv_heads, 4);
    assert_eq!(gpu_model.config.num_layers, 2);
    assert_eq!(gpu_model.config.intermediate_dim, 128);
}

#[test]
fn test_apr_f32_to_gpu_gqa_config_preserved() {
    let apr = create_gqa_f32_apr();
    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);

    assert!(result.is_ok());
    let gpu_model = result.unwrap();

    // GQA config should be preserved
    assert_eq!(gpu_model.config.num_heads, 8);
    assert_eq!(gpu_model.config.num_kv_heads, 2);
    assert!(gpu_model.config.is_gqa());
}

#[test]
fn test_apr_f32_to_gpu_with_optional_biases() {
    let config = create_minimal_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;

    // Create APR with all optional biases filled in
    let apr = AprTransformer {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: vec![AprTransformerLayer {
            attn_norm_weight: vec![1.0; hidden_dim],
            attn_norm_bias: Some(vec![0.1; hidden_dim]),
            qkv_weight: vec![0.01; hidden_dim * 3 * hidden_dim],
            qkv_bias: Some(vec![0.01; 3 * hidden_dim]),
            attn_output_weight: vec![0.01; hidden_dim * hidden_dim],
            attn_output_bias: Some(vec![0.02; hidden_dim]),
            ffn_norm_weight: Some(vec![1.0; hidden_dim]),
            ffn_norm_bias: Some(vec![0.1; hidden_dim]),
            ffn_up_weight: vec![0.01; hidden_dim * 128],
            ffn_up_bias: Some(vec![0.01; 128]),
            ffn_down_weight: vec![0.01; 128 * hidden_dim],
            ffn_down_bias: Some(vec![0.01; hidden_dim]),
            ffn_gate_weight: None,
            ffn_gate_bias: None,
        }],
        output_norm_weight: vec![1.0; hidden_dim],
        output_norm_bias: Some(vec![0.1; hidden_dim]),
        lm_head_weight: vec![0.01; vocab_size * hidden_dim],
        lm_head_bias: Some(vec![0.001; vocab_size]),
        q4k_layers: None,
        lm_head_weight_q6k: None,
        lm_head_weight_q4k: None,
    };

    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok(), "Should handle optional biases");
}

#[test]
fn test_apr_f32_to_gpu_without_optional_biases() {
    let config = create_minimal_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;

    // Create APR with no optional biases (all None)
    let apr = AprTransformer {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: vec![AprTransformerLayer {
            attn_norm_weight: vec![1.0; hidden_dim],
            attn_norm_bias: None,
            qkv_weight: vec![0.01; hidden_dim * 3 * hidden_dim],
            qkv_bias: None,
            attn_output_weight: vec![0.01; hidden_dim * hidden_dim],
            attn_output_bias: None,
            ffn_norm_weight: None, // Tests default creation
            ffn_norm_bias: None,
            ffn_up_weight: vec![0.01; hidden_dim * 128],
            ffn_up_bias: None,
            ffn_down_weight: vec![0.01; 128 * hidden_dim],
            ffn_down_bias: None,
            ffn_gate_weight: None,
            ffn_gate_bias: None,
        }],
        output_norm_weight: vec![1.0; hidden_dim],
        output_norm_bias: None,
        lm_head_weight: vec![0.01; vocab_size * hidden_dim],
        lm_head_bias: None,
        q4k_layers: None,
        lm_head_weight_q6k: None,
        lm_head_weight_q4k: None,
    };

    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok(), "Should handle missing optional biases");
}

#[test]
fn test_apr_f32_to_gpu_with_swiglu_gate() {
    let apr = create_minimal_f32_apr(); // Has gate weight
    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);

    assert!(result.is_ok());
    // Model should be created successfully with SwiGLU gate
}

#[test]
fn test_apr_f32_to_gpu_without_swiglu_gate() {
    let config = create_minimal_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;

    let apr = AprTransformer {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: vec![AprTransformerLayer {
            attn_norm_weight: vec![1.0; hidden_dim],
            attn_norm_bias: None,
            qkv_weight: vec![0.01; hidden_dim * 3 * hidden_dim],
            qkv_bias: None,
            attn_output_weight: vec![0.01; hidden_dim * hidden_dim],
            attn_output_bias: None,
            ffn_norm_weight: Some(vec![1.0; hidden_dim]),
            ffn_norm_bias: None,
            ffn_up_weight: vec![0.01; hidden_dim * 128],
            ffn_up_bias: None,
            ffn_down_weight: vec![0.01; 128 * hidden_dim],
            ffn_down_bias: None,
            ffn_gate_weight: None, // No gate
            ffn_gate_bias: None,
        }],
        output_norm_weight: vec![1.0; hidden_dim],
        output_norm_bias: None,
        lm_head_weight: vec![0.01; vocab_size * hidden_dim],
        lm_head_bias: None,
        q4k_layers: None,
        lm_head_weight_q6k: None,
        lm_head_weight_q4k: None,
    };

    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok());
}

// ============================================================================
// AprToGpuAdapter Q4 Comprehensive Tests
// ============================================================================

#[test]
fn test_apr_q4_to_gpu_basic_conversion() {
    let apr = create_minimal_q4_apr();
    let result = AprToGpuAdapter::to_gpu_model(&apr);

    assert!(result.is_ok(), "Q4 conversion should succeed");
    let gpu_model = result.unwrap();

    assert_eq!(gpu_model.config.vocab_size, 100);
    assert_eq!(gpu_model.config.hidden_dim, 64);
}

#[test]
fn test_apr_q4_to_gpu_without_gate() {
    let apr = create_q4_apr_without_gate();
    let result = AprToGpuAdapter::to_gpu_model(&apr);

    assert!(result.is_ok());
}

#[test]
fn test_apr_q4_config_conversion_all_fields() {
    let apr_config = AprTransformerConfig {
        architecture: "test".to_string(),
        hidden_dim: 512,
        num_layers: 8,
        num_heads: 16,
        num_kv_heads: 4,
        vocab_size: 32000,
        intermediate_dim: 2048,
        context_length: 4096,
        rope_theta: 500000.0,
        eps: 1e-6,
    };

    let gpu_config = AprToGpuAdapter::config_to_gpu(&apr_config);

    assert_eq!(gpu_config.vocab_size, 32000);
    assert_eq!(gpu_config.hidden_dim, 512);
    assert_eq!(gpu_config.num_heads, 16);
    assert_eq!(gpu_config.num_kv_heads, 4);
    assert_eq!(gpu_config.num_layers, 8);
    assert_eq!(gpu_config.intermediate_dim, 2048);
    assert_eq!(gpu_config.eps, 1e-6);
    assert_eq!(gpu_config.rope_theta, 500000.0);
}

#[test]
fn test_apr_q4_extract_qkv_mha() {
    let config = create_minimal_apr_config();
    let layer = create_q4_layer(&config);

    let result = AprToGpuAdapter::extract_qkv_weights(
        &layer,
        config.hidden_dim,
        config.num_heads,
        config.num_kv_heads,
    );
    assert!(result.is_ok());

    let head_dim = config.hidden_dim / config.num_heads;
    let kv_dim = config.num_kv_heads * head_dim;
    let qkv_out_dim = config.hidden_dim + 2 * kv_dim;
    let expected_len = config.hidden_dim * qkv_out_dim;
    assert_eq!(result.unwrap().len(), expected_len);
}

#[test]
fn test_apr_q4_extract_qkv_gqa() {
    let config = create_gqa_apr_config();
    let layer = QuantizedAprLayerQ4 {
        attn_norm_weight: vec![1.0; config.hidden_dim],
        qkv_weight: QuantizedAprTensorQ4::zeros(config.hidden_dim, config.hidden_dim + 2 * 16),
        attn_output_weight: QuantizedAprTensorQ4::zeros(config.hidden_dim, config.hidden_dim),
        ffn_up_weight: QuantizedAprTensorQ4::zeros(config.hidden_dim, config.intermediate_dim),
        ffn_down_weight: QuantizedAprTensorQ4::zeros(config.intermediate_dim, config.hidden_dim),
        ffn_gate_weight: None,
        ffn_norm_weight: Some(vec![1.0; config.hidden_dim]),
    };

    let result = AprToGpuAdapter::extract_qkv_weights(
        &layer,
        config.hidden_dim,
        config.num_heads,
        config.num_kv_heads,
    );
    assert!(result.is_ok());

    // GQA: kv_dim = 2 * 8 = 16
    let head_dim = config.hidden_dim / config.num_heads; // 64 / 8 = 8
    let kv_dim = config.num_kv_heads * head_dim; // 2 * 8 = 16
    let qkv_out_dim = config.hidden_dim + 2 * kv_dim; // 64 + 32 = 96
    let expected_len = config.hidden_dim * qkv_out_dim;
    assert_eq!(result.unwrap().len(), expected_len);
}

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

#[test]
fn test_transpose_3x3_explicit() {
    let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
    let result = transpose_matrix(&data, 3, 3);

    // Original: [[1,2,3],[4,5,6],[7,8,9]]
    // Transposed: [[1,4,7],[2,5,8],[3,6,9]]
    assert_eq!(result, vec![1.0, 4.0, 7.0, 2.0, 5.0, 8.0, 3.0, 6.0, 9.0]);
}

#[test]
fn test_transpose_2x4() {
    let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    let result = transpose_matrix(&data, 2, 4);

    // Original 2x4: [[1,2,3,4],[5,6,7,8]]
    // Transposed 4x2: [[1,5],[2,6],[3,7],[4,8]]
    assert_eq!(result, vec![1.0, 5.0, 2.0, 6.0, 3.0, 7.0, 4.0, 8.0]);
}

#[test]
fn test_transpose_4x2() {
    let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    let result = transpose_matrix(&data, 4, 2);

    // Original 4x2: [[1,2],[3,4],[5,6],[7,8]]
    // Transposed 2x4: [[1,3,5,7],[2,4,6,8]]
    assert_eq!(result, vec![1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0]);
}

#[test]
fn test_transpose_double_is_identity() {
    let original = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let transposed = transpose_matrix(&original, 2, 3);
    let double_transposed = transpose_matrix(&transposed, 3, 2);

    assert_eq!(original, double_transposed);
}

#[test]
fn test_transpose_preserves_sum() {
    let data = vec![
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
    ];
    let original_sum: f32 = data.iter().sum();

    let transposed = transpose_matrix(&data, 3, 4);
    let transposed_sum: f32 = transposed.iter().sum();

    assert!((original_sum - transposed_sum).abs() < 1e-6);
}

#[test]
fn test_transpose_symmetric_matrix() {
    // A symmetric matrix should have some invariant properties
    let data = vec![1.0, 2.0, 2.0, 3.0]; // 2x2 symmetric
    let result = transpose_matrix(&data, 2, 2);

    // For symmetric, transpose should give different result unless it's identity-like
    // [1, 2]   -> [1, 2]
    // [2, 3]      [2, 3]
    assert_eq!(result, vec![1.0, 2.0, 2.0, 3.0]);
}

// ============================================================================
// MockExecutor Integration Tests
// ============================================================================

#[test]
fn test_mock_executor_with_gpu_model_from_apr_f32() {
    let apr = create_minimal_f32_apr();
    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok());

    let mut gpu_model = result.unwrap();

    // Inject mock executor
    let mock = MockExecutor::new("apr_f32_test");
    gpu_model.with_test_executor(Box::new(mock));

    assert!(gpu_model.has_test_executor());

    // Perform a forward pass with mock
    let token_ids = vec![1, 2, 3];
    let forward_result = gpu_model.forward_gpu(&token_ids);
    assert!(forward_result.is_ok());
}

#[test]
fn test_mock_executor_with_gpu_model_from_apr_q4() {
    let apr = create_minimal_q4_apr();
    let result = AprToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok());

    let mut gpu_model = result.unwrap();

    // Inject mock executor
    let mock = MockExecutor::new("apr_q4_test");
    gpu_model.with_test_executor(Box::new(mock));

    assert!(gpu_model.has_test_executor());

    // Perform a forward pass with mock
    let token_ids = vec![5, 10, 15];
    let forward_result = gpu_model.forward_gpu(&token_ids);
    assert!(forward_result.is_ok());
}

#[test]
fn test_mock_executor_generate_with_adapted_model() {
    let apr = create_minimal_f32_apr();
    let mut gpu_model = AprF32ToGpuAdapter::to_gpu_model(&apr).unwrap();

    let mock = MockExecutor::new("generate_test");
    gpu_model.with_test_executor(Box::new(mock));

    let gen_config = GpuGenerateConfig::deterministic(3);
    let prompt = vec![1, 2];
    let result = gpu_model.generate(&prompt, &gen_config);

    assert!(result.is_ok());
    let tokens = result.unwrap();
    assert!(tokens.len() >= prompt.len());
}

#[test]
fn test_mock_executor_clear_and_restore() {
    let apr = create_minimal_f32_apr();
    let mut gpu_model = AprF32ToGpuAdapter::to_gpu_model(&apr).unwrap();

    // Initially no test executor
    assert!(!gpu_model.has_test_executor());

    // Add mock
    let mock = MockExecutor::new("clear_test");
    gpu_model.with_test_executor(Box::new(mock));
    assert!(gpu_model.has_test_executor());

    // Clear mock
    gpu_model.clear_test_executor();
    assert!(!gpu_model.has_test_executor());

    // Add another mock
    let mock2 = MockExecutor::new("restored_test");
    gpu_model.with_test_executor(Box::new(mock2));
    assert!(gpu_model.has_test_executor());
}

// ============================================================================
// GpuModelConfig Derived Tests
// ============================================================================

#[test]
fn test_gpu_model_config_from_apr_f32_dimensions() {
    let apr = create_minimal_f32_apr();
    let gpu_model = AprF32ToGpuAdapter::to_gpu_model(&apr).unwrap();

    // Verify derived dimensions
    let config = gpu_model.config();
    assert_eq!(config.head_dim(), 16); // 64 / 4
    assert_eq!(config.kv_dim(), 64); // 4 * 16 (MHA)
    assert_eq!(config.qkv_dim(), 192); // 64 + 2*64 (MHA)
    assert!(!config.is_gqa());
}

#[test]
fn test_gpu_model_config_from_apr_f32_gqa_dimensions() {
    let apr = create_gqa_f32_apr();
    let gpu_model = AprF32ToGpuAdapter::to_gpu_model(&apr).unwrap();

    let config = gpu_model.config();
    assert_eq!(config.head_dim(), 8); // 64 / 8
    assert_eq!(config.kv_dim(), 16); // 2 * 8 (GQA)
    assert_eq!(config.qkv_dim(), 96); // 64 + 2*16 (GQA)
    assert!(config.is_gqa());
}

#[test]
fn test_gpu_model_config_from_apr_q4_dimensions() {
    let apr = create_minimal_q4_apr();
    let gpu_model = AprToGpuAdapter::to_gpu_model(&apr).unwrap();

    let config = gpu_model.config();
    assert_eq!(config.head_dim(), 16);
    assert_eq!(config.kv_dim(), 64);
    assert_eq!(config.qkv_dim(), 192);
    assert!(!config.is_gqa());
}

// ============================================================================
// Multiple Layer Tests
// ============================================================================

#[test]
fn test_apr_f32_to_gpu_multiple_layers() {
    let config = create_minimal_apr_config();
    let hidden_dim = config.hidden_dim;
    let vocab_size = config.vocab_size;

    let apr = AprTransformer {
        config: config.clone(),
        token_embedding: vec![0.01; vocab_size * hidden_dim],
        layers: (0..4).map(|_| create_f32_layer(&config)).collect(), // 4 layers
        output_norm_weight: vec![1.0; hidden_dim],
        output_norm_bias: None,
        lm_head_weight: vec![0.01; vocab_size * hidden_dim],
        lm_head_bias: None,
        q4k_layers: None,
        lm_head_weight_q6k: None,
        lm_head_weight_q4k: None,
    };

    let result = AprF32ToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok());

    let gpu_model = result.unwrap();
    // Config should still say num_layers from config, but model has 4 layers
    assert_eq!(gpu_model.config.num_layers, 2); // Config says 2
}

#[test]
fn test_apr_q4_to_gpu_single_layer() {
    let config = AprTransformerConfig {
        architecture: "single".to_string(),
        hidden_dim: 64,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        intermediate_dim: 128,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-5,
    };

    let apr = QuantizedAprTransformerQ4 {
        config: config.clone(),
        token_embedding: vec![0.01; 100 * 64],
        layers: vec![create_q4_layer(&config)], // Single layer
        output_norm_weight: vec![1.0; 64],
        lm_head_weight: QuantizedAprTensorQ4::zeros(64, 100),
    };

    let result = AprToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok());

    let gpu_model = result.unwrap();
    assert_eq!(gpu_model.config.num_layers, 1);
}

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

    let apr = QuantizedAprTransformerQ4 {
        config: config.clone(),
        token_embedding: vec![0.01; 100 * 64],
        layers: vec![], // No layers
        output_norm_weight: vec![1.0; 64],
        lm_head_weight: QuantizedAprTensorQ4::zeros(64, 100),
    };

    let result = AprToGpuAdapter::to_gpu_model(&apr);
    assert!(result.is_ok());

    let gpu_model = result.unwrap();
    assert_eq!(gpu_model.config.num_layers, 0);
}

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

#[test]
fn test_dequantize_tensor_zero_expected() {
    let result = AprToGpuAdapter::dequantize_tensor(&[], 0);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), 0);
}

#[test]
fn test_dequantize_tensor_single_block() {
    // Q4_0 block: 2 bytes scale + 16 bytes data = 18 bytes for 32 values
    let mut data = vec![0u8; 18];
    data[0] = 0x00;
    data[1] = 0x3c; // f16 1.0

    let result = AprToGpuAdapter::dequantize_tensor(&data, 32);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), 32);
}

#[test]
fn test_dequantize_tensor_needs_padding() {
    // Single block produces 32 values, but we request 64
    let mut data = vec![0u8; 18];
    data[0] = 0x00;
    data[1] = 0x3c;

    let result = AprToGpuAdapter::dequantize_tensor(&data, 64);
    assert!(result.is_ok());

    let values = result.unwrap();
    assert_eq!(values.len(), 64);
    // Last 32 values should be zero (padded)
    for &v in &values[32..] {
        assert_eq!(v, 0.0);
    }
}

#[test]
fn test_dequantize_tensor_needs_truncation() {
    // Two blocks produce 64 values, but we only request 32
    let mut data = vec![0u8; 36];
    data[0] = 0x00;
    data[1] = 0x3c;
    data[18] = 0x00;
    data[19] = 0x3c;

    let result = AprToGpuAdapter::dequantize_tensor(&data, 32);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().len(), 32);
}

// ============================================================================
// RoPE Theta Tests
// ============================================================================

#[test]
fn test_rope_theta_preservation() {
    let apr_config = AprTransformerConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        num_layers: 2,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        intermediate_dim: 128,
        context_length: 256,
        rope_theta: 500000.0, // Non-standard theta
        eps: 1e-5,
    };

    let gpu_config = AprToGpuAdapter::config_to_gpu(&apr_config);
    assert_eq!(gpu_config.rope_theta, 500000.0);
}

#[test]
fn test_eps_preservation() {
    let apr_config = AprTransformerConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        num_layers: 2,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 100,
        intermediate_dim: 128,
        context_length: 256,
        rope_theta: 10000.0,
        eps: 1e-6, // Different epsilon
    };

    let gpu_config = AprToGpuAdapter::config_to_gpu(&apr_config);
    assert_eq!(gpu_config.eps, 1e-6);
}

// ============================================================================
// FFN Weight Extraction Tests
// ============================================================================

#[test]
fn test_extract_ffn_weights_dimensions() {
    let config = create_minimal_apr_config();
    let layer = create_q4_layer(&config);

    let result =
        AprToGpuAdapter::extract_ffn_weights(&layer, config.hidden_dim, config.intermediate_dim);
    assert!(result.is_ok());

    let (fc1, fc2) = result.unwrap();
    assert_eq!(fc1.len(), config.hidden_dim * config.intermediate_dim);
    assert_eq!(fc2.len(), config.intermediate_dim * config.hidden_dim);
}

#[test]
fn test_extract_out_weights_dimensions() {
    let config = create_minimal_apr_config();
    let layer = create_q4_layer(&config);

    let result = AprToGpuAdapter::extract_out_weights(&layer, config.hidden_dim);
    assert!(result.is_ok());

    let weights = result.unwrap();
    assert_eq!(weights.len(), config.hidden_dim * config.hidden_dim);
}
