//! GGUF Part 04: IMP-112 (HybridScheduler Caching) + IMP-111 (Flash Attention) +
//!               IMP-113 (Batched GPU Kernels) + IMP-114 (Flattened GEMM)
//!
//! Extracted from gguf_monolith.rs (PMAT-802)
//!
//! ## Test Coverage
//!
//! - IMP-112a-d: HybridScheduler caching, cached vs uncached, multiple ops, attention
//! - PARITY-114: CUDA GEMM correctness verification
//! - IMP-111a-d: Online softmax, tiled attention, causal mask, tile sizes
//! - IMP-113a-d: Batched GEMM single dispatch, attention correctness, dispatch count, softmax
//! - IMP-114a-d: Flattened batched GEMM, loop matching, attention, large batches

use crate::gguf::test_helpers::create_test_model_with_config;
use crate::gguf::{GGUFConfig, OwnedQuantizedModelCached};

// =========================================================================
// IMP-112: HybridScheduler Caching
// =========================================================================

#[test]
#[cfg(feature = "gpu")]
fn test_imp_112a_cached_scheduler_initialization() {
    // IMP-112a: Verify cached scheduler initializes lazily and is reused
    // This tests that OwnedQuantizedModelCached provides scheduler caching
    let config = GGUFConfig {
        architecture: "test".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);

    // Create cached wrapper
    let cached_model = OwnedQuantizedModelCached::new(model);

    // First call should initialize scheduler
    let tokens = vec![1u32, 5, 10];
    let result1 = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("IMP-112a: First cached forward should succeed");

    // Verify output shape
    assert_eq!(
        result1.len(),
        tokens.len() * config.vocab_size,
        "IMP-112a: Should return correct output shape"
    );

    // Second call should reuse scheduler (much faster)
    let result2 = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("IMP-112a: Second cached forward should succeed");

    // Results should be identical (same scheduler, same computation)
    assert_eq!(result1.len(), result2.len());
    for i in 0..result1.len() {
        let diff = (result1[i] - result2[i]).abs();
        assert!(
            diff < 1e-6,
            "IMP-112a: Results should be identical on repeated calls, pos {}: diff={}",
            i,
            diff
        );
    }
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_112b_cached_matches_uncached() {
    // IMP-112b: Verify cached scheduler produces identical results to uncached
    //
    // PARITY-114 RESOLVED: CUDA GEMM grid launch dimensions were swapped (M<->N).
    // The fix swaps Grid X and Grid Y in all GEMM launch configurations:
    // - Grid X = (n + 31) / 32 for columns (N dimension)
    // - Grid Y = (m + 31) / 32 for rows (M dimension)
    //
    // Both cached (CUDA) and uncached (wgpu) paths now produce matching results.
    let config = GGUFConfig {
        architecture: "test".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 cached_model = OwnedQuantizedModelCached::new(model.clone());

    let tokens = vec![1u32, 5, 10, 20];

    // Uncached forward (creates new scheduler each time)
    let uncached_result = model
        .forward_batch_gpu(&tokens)
        .expect("Uncached forward should succeed");

    // Cached forward (reuses scheduler)
    let cached_result = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("Cached forward should succeed");

    // Results should match
    assert_eq!(uncached_result.len(), cached_result.len());
    for i in 0..uncached_result.len() {
        let diff = (uncached_result[i] - cached_result[i]).abs();
        assert!(
            diff < 1e-4,
            "IMP-112b: Cached should match uncached, pos {}: uncached={}, cached={}, diff={}",
            i,
            uncached_result[i],
            cached_result[i],
            diff
        );
    }
}

/// PARITY-114 RESOLVED: CUDA path correctness verification
///
/// This test verifies the CUDA GEMM fix is correct by checking that
/// CUDA produces values in a reasonable range (not ~10x smaller as before).
///
/// Root cause was swapped grid dimensions in CUDA launch config:
/// - Grid X was (m+31)/32 but should be (n+31)/32 for columns
/// - Grid Y was (n+31)/32 but should be (m+31)/32 for rows
#[test]
#[cfg(feature = "cuda")]
fn test_parity_114_cuda_gemm_correctness() {
    let config = GGUFConfig {
        architecture: "test".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 cached_model = OwnedQuantizedModelCached::new(model);

    let tokens = vec![1u32, 5, 10, 20];

    // Cached forward uses CUDA when available
    let result = cached_model
        .forward_batch_gpu_cached(&tokens)
        .expect("PARITY-114: CUDA forward should succeed");

    // CUDA produces finite values
    assert_eq!(result.len(), tokens.len() * config.vocab_size);
    assert!(
        result.iter().all(|x| x.is_finite()),
        "PARITY-114: CUDA should produce finite values"
    );

    // PARITY-114 RESOLVED: Values should no longer be ~10x smaller
    // The fix swaps Grid X and Grid Y in CUDA launch configurations
    let max_val = result.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let min_val = result.iter().cloned().fold(f32::INFINITY, f32::min);
    eprintln!(
        "PARITY-114 RESOLVED: CUDA output range: [{:.4}, {:.4}]",
        min_val, max_val
    );
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_112c_multiple_operations_same_scheduler() {
    // IMP-112c: Verify multiple different operations share the same scheduler
    let config = GGUFConfig {
        architecture: "test".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 cached_model = OwnedQuantizedModelCached::new(model);

    // Multiple forward passes with different inputs
    let tokens1 = vec![1u32, 2, 3];
    let tokens2 = vec![10u32, 20, 30, 40];
    let tokens3 = vec![5u32];

    let result1 = cached_model
        .forward_batch_gpu_cached(&tokens1)
        .expect("IMP-112c: Forward 1 should succeed");
    let result2 = cached_model
        .forward_batch_gpu_cached(&tokens2)
        .expect("IMP-112c: Forward 2 should succeed");
    let result3 = cached_model
        .forward_batch_gpu_cached(&tokens3)
        .expect("IMP-112c: Forward 3 should succeed");

    // Verify shapes
    assert_eq!(result1.len(), 3 * config.vocab_size);
    assert_eq!(result2.len(), 4 * config.vocab_size);
    assert_eq!(result3.len(), config.vocab_size);

    // All results should be finite
    assert!(result1.iter().all(|x| x.is_finite()));
    assert!(result2.iter().all(|x| x.is_finite()));
    assert!(result3.iter().all(|x| x.is_finite()));
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_112d_cached_attention_matches_uncached() {
    // IMP-112d: Verify cached parallel attention matches uncached
    let config = GGUFConfig {
        architecture: "test".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 cached_model = OwnedQuantizedModelCached::new(model.clone());

    let seq_len = 8;
    let hidden_dim = config.hidden_dim;

    // Create Q, K, V tensors
    let q: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    // Uncached attention
    let uncached_attn = model
        .parallel_multihead_attention_gpu(&q, &k, &v, seq_len)
        .expect("Uncached attention should succeed");

    // Cached attention
    let cached_attn = cached_model
        .parallel_multihead_attention_gpu_cached(&q, &k, &v, seq_len)
        .expect("Cached attention should succeed");

    // Results should match
    assert_eq!(uncached_attn.len(), cached_attn.len());
    for i in 0..uncached_attn.len() {
        let diff = (uncached_attn[i] - cached_attn[i]).abs();
        assert!(
            diff < 1e-4,
            "IMP-112d: Cached attention should match uncached, pos {}: diff={}",
            i,
            diff
        );
    }
}

// =========================================================================
// IMP-111: Flash Attention-style Tiled Computation
// =========================================================================

#[test]
fn test_imp_111a_online_softmax_correctness() {
    // IMP-111a: Verify online softmax matches standard softmax
    // Online softmax processes data in tiles, tracking running max and sum
    let config = GGUFConfig {
        architecture: "test".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);

    // Test data: attention scores for one row
    let scores: Vec<f32> = (0..16).map(|i| ((i % 7) as f32 - 3.0) * 0.5).collect();

    // Standard softmax (reference)
    let standard = model.standard_softmax(&scores);

    // Online softmax (tiled, O(1) memory per tile)
    let tile_size = 4;
    let online = model
        .online_softmax(&scores, tile_size)
        .expect("IMP-111a: Online softmax should succeed");

    // Results should match within numerical tolerance
    assert_eq!(standard.len(), online.len());
    for i in 0..standard.len() {
        let diff = (standard[i] - online[i]).abs();
        assert!(
            diff < 1e-5,
            "IMP-111a: Online softmax differs at {}: standard={}, online={}, diff={}",
            i,
            standard[i],
            online[i],
            diff
        );
    }

    // Verify both sum to 1
    let std_sum: f32 = standard.iter().sum();
    let online_sum: f32 = online.iter().sum();
    assert!(
        (std_sum - 1.0).abs() < 1e-5,
        "Standard softmax should sum to 1"
    );
    assert!(
        (online_sum - 1.0).abs() < 1e-5,
        "Online softmax should sum to 1"
    );
}

#[test]
fn test_imp_111b_tiled_attention_matches_standard() {
    // IMP-111b: Verify tiled attention produces same output as standard
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let seq_len = 8;
    let head_dim = config.hidden_dim / config.num_heads; // 8

    // Create Q, K, V for single head
    let q: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    let scale = 1.0 / (head_dim as f32).sqrt();

    // Standard attention (materializes full attention matrix)
    let standard_output = model
        .standard_single_head_attention(&q, &k, &v, seq_len, head_dim, scale)
        .expect("Standard attention should succeed");

    // Tiled attention (O(1) memory for softmax per tile)
    let tile_size = 4;
    let tiled_output = model
        .tiled_single_head_attention(&q, &k, &v, seq_len, head_dim, scale, tile_size)
        .expect("IMP-111b: Tiled attention should succeed");

    // Results should match
    assert_eq!(standard_output.len(), tiled_output.len());
    for i in 0..standard_output.len() {
        let diff = (standard_output[i] - tiled_output[i]).abs();
        assert!(
            diff < 1e-4,
            "IMP-111b: Tiled attention differs at {}: standard={}, tiled={}, diff={}",
            i,
            standard_output[i],
            tiled_output[i],
            diff
        );
    }
}

#[test]
fn test_imp_111c_tiled_causal_attention() {
    // IMP-111c: Verify tiled attention respects causal mask
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let seq_len = 8;
    let head_dim = config.hidden_dim / config.num_heads;

    // Create deterministic Q, K, V
    let q: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| (i as f32 * 0.1) % 1.0)
        .collect();
    let k: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i + 5) as f32 * 0.1) % 1.0)
        .collect();
    let v: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i + 10) as f32 * 0.1) % 1.0)
        .collect();

    let scale = 1.0 / (head_dim as f32).sqrt();
    let tile_size = 4;

    // Tiled causal attention
    let tiled_output = model
        .tiled_causal_attention(&q, &k, &v, seq_len, head_dim, scale, tile_size)
        .expect("IMP-111c: Tiled causal attention should succeed");

    // Verify output shape
    assert_eq!(
        tiled_output.len(),
        seq_len * head_dim,
        "IMP-111c: Output should have seq_len * head_dim elements"
    );

    // Verify finite values
    assert!(
        tiled_output.iter().all(|x| x.is_finite()),
        "IMP-111c: All outputs should be finite"
    );

    // Verify causality: output at position i should only depend on positions 0..=i
    // We test this by checking that changing K/V at position j > i doesn't affect output[i]
    let mut k_modified = k.clone();
    // Modify K at last position
    for d in 0..head_dim {
        k_modified[(seq_len - 1) * head_dim + d] = 999.0;
    }

    let modified_output = model
        .tiled_causal_attention(&q, &k_modified, &v, seq_len, head_dim, scale, tile_size)
        .expect("Modified attention should succeed");

    // Positions 0 to seq_len-2 should be unchanged (they don't attend to position seq_len-1)
    for pos in 0..seq_len - 1 {
        for d in 0..head_dim {
            let idx = pos * head_dim + d;
            let diff = (tiled_output[idx] - modified_output[idx]).abs();
            assert!(
                diff < 1e-6,
                "IMP-111c: Position {} should not be affected by future positions, diff={}",
                pos,
                diff
            );
        }
    }
}

#[test]
fn test_imp_111d_tiled_attention_various_tile_sizes() {
    // IMP-111d: Verify tiled attention works with various tile sizes
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let seq_len = 16;
    let head_dim = config.hidden_dim / config.num_heads;

    let q: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * head_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    let scale = 1.0 / (head_dim as f32).sqrt();

    // Get reference with tile_size = 1 (equivalent to standard)
    let reference = model
        .tiled_causal_attention(&q, &k, &v, seq_len, head_dim, scale, 1)
        .expect("Reference should succeed");

    // Test various tile sizes
    for tile_size in [2, 4, 8, 16] {
        let output = model
            .tiled_causal_attention(&q, &k, &v, seq_len, head_dim, scale, tile_size)
            .unwrap_or_else(|_| panic!("Tile size {} should succeed", tile_size));

        assert_eq!(output.len(), reference.len());
        for i in 0..output.len() {
            let diff = (output[i] - reference[i]).abs();
            assert!(
                diff < 1e-4,
                "IMP-111d: Tile size {} differs at {}: ref={}, tiled={}, diff={}",
                tile_size,
                i,
                reference[i],
                output[i],
                diff
            );
        }
    }
}

// ========================================================================
// IMP-113: True Batched GPU Kernel Tests (Single Dispatch)
// ========================================================================

#[test]
#[cfg(feature = "gpu")]
fn test_imp_113a_batched_gemm_single_dispatch() {
    // IMP-113a: Verify batched GEMM processes all heads in single dispatch
    // This is the foundation for efficient multi-head attention
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    let num_heads = 4;
    let seq_len = 8;
    let head_dim = 16;

    // Create batched A: [num_heads, seq_len, head_dim]
    let batched_a: Vec<f32> = (0..num_heads * seq_len * head_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();

    // Create batched B: [num_heads, head_dim, seq_len]
    let batched_b: Vec<f32> = (0..num_heads * head_dim * seq_len)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();

    // Single dispatch batched GEMM
    let result = cached_model
        .batched_gemm_single_dispatch(
            &batched_a, &batched_b, num_heads, seq_len, head_dim, seq_len,
        )
        .expect("Batched GEMM should succeed");

    // Output: [num_heads, seq_len, seq_len]
    assert_eq!(
        result.len(),
        num_heads * seq_len * seq_len,
        "IMP-113a: Output should have shape [num_heads, seq_len, seq_len]"
    );

    // Verify by computing reference per-head
    for h in 0..num_heads {
        let a_start = h * seq_len * head_dim;
        let b_start = h * head_dim * seq_len;
        let out_start = h * seq_len * seq_len;

        for i in 0..seq_len {
            for j in 0..seq_len {
                let mut expected = 0.0f32;
                for k in 0..head_dim {
                    expected += batched_a[a_start + i * head_dim + k]
                        * batched_b[b_start + k * seq_len + j];
                }
                let actual = result[out_start + i * seq_len + j];
                let diff = (expected - actual).abs();
                assert!(
                    diff < 1e-3,
                    "IMP-113a: Head {} mismatch at ({},{}): expected={}, actual={}, diff={}",
                    h,
                    i,
                    j,
                    expected,
                    actual,
                    diff
                );
            }
        }
    }
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_113b_single_dispatch_attention_correctness() {
    // IMP-113b: Verify single-dispatch attention matches multi-dispatch
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model.clone());

    let seq_len = 8;
    let hidden_dim = config.hidden_dim;

    // Create Q, K, V: [seq_len, hidden_dim]
    let q: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    // Multi-dispatch reference (existing implementation)
    let reference = cached_model
        .parallel_multihead_attention_gpu_cached(&q, &k, &v, seq_len)
        .expect("Multi-dispatch attention should succeed");

    // Single-dispatch new implementation
    let result = cached_model
        .single_dispatch_multihead_attention(&q, &k, &v, seq_len)
        .expect("Single-dispatch attention should succeed");

    // Compare outputs
    assert_eq!(result.len(), reference.len());
    for i in 0..result.len() {
        let diff = (result[i] - reference[i]).abs();
        assert!(
            diff < 1e-3,
            "IMP-113b: Single-dispatch differs at {}: ref={}, single={}, diff={}",
            i,
            reference[i],
            result[i],
            diff
        );
    }
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_113c_single_dispatch_dispatch_count() {
    // IMP-113c: Verify single-dispatch uses fewer GPU dispatches
    // This test validates the architectural improvement
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 8, // More heads = bigger benefit
        num_kv_heads: 8,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    let seq_len = 16;
    let hidden_dim = config.hidden_dim;

    let q: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    // Both should succeed and produce valid output
    let single_result = cached_model
        .single_dispatch_multihead_attention(&q, &k, &v, seq_len)
        .expect("Single-dispatch should succeed");

    // Validate output dimensions
    assert_eq!(
        single_result.len(),
        seq_len * hidden_dim,
        "IMP-113c: Output should have shape [seq_len, hidden_dim]"
    );

    // Validate output is not all zeros (sanity check)
    let sum: f32 = single_result.iter().map(|x| x.abs()).sum();
    assert!(
        sum > 0.01,
        "IMP-113c: Output should have non-trivial values, got sum={}",
        sum
    );
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_113d_batched_softmax_correctness() {
    // IMP-113d: Verify batched softmax with causal mask
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 32,
        intermediate_dim: 64,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    let num_heads = 4;
    let seq_len = 8;

    // Create batched scores: [num_heads, seq_len, seq_len]
    let batched_scores: Vec<f32> = (0..num_heads * seq_len * seq_len)
        .map(|i| ((i % 17) as f32 - 8.0) * 0.2)
        .collect();

    // Apply batched causal softmax
    let result = cached_model
        .batched_causal_softmax(&batched_scores, num_heads, seq_len)
        .expect("Batched causal softmax should succeed");

    // Verify dimensions
    assert_eq!(result.len(), num_heads * seq_len * seq_len);

    // Verify each row sums to 1.0 (within causal mask)
    for h in 0..num_heads {
        for i in 0..seq_len {
            let row_start = h * seq_len * seq_len + i * seq_len;
            let row_sum: f32 = (0..=i).map(|j| result[row_start + j]).sum();
            assert!(
                (row_sum - 1.0).abs() < 1e-5,
                "IMP-113d: Head {} row {} should sum to 1.0, got {}",
                h,
                i,
                row_sum
            );

            // Verify causal mask: positions > i should be 0
            for j in (i + 1)..seq_len {
                assert!(
                    result[row_start + j].abs() < 1e-6,
                    "IMP-113d: Head {} pos ({},{}) should be masked, got {}",
                    h,
                    i,
                    j,
                    result[row_start + j]
                );
            }
        }
    }
}

// ========================================================================
// IMP-114: True GPU Batched GEMM Kernel Tests
// ========================================================================

#[test]
#[cfg(feature = "gpu")]
fn test_imp_114a_flattened_batched_gemm_correctness() {
    // IMP-114a: Verify flattened batched GEMM computes correct results
    // Strategy: Flatten [batch, m, k] @ [batch, k, n] into single large matmul
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    let batch_size = 4;
    let m = 8;
    let k = 16;
    let n = 8;

    // Create batched matrices
    let batched_a: Vec<f32> = (0..batch_size * m * k)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let batched_b: Vec<f32> = (0..batch_size * k * n)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();

    // Use flattened batched GEMM (true single dispatch)
    let result = cached_model
        .flattened_batched_gemm(&batched_a, &batched_b, batch_size, m, k, n)
        .expect("Flattened batched GEMM should succeed");

    // Output should be [batch_size, m, n]
    assert_eq!(
        result.len(),
        batch_size * m * n,
        "IMP-114a: Output should have shape [batch, m, n]"
    );

    // Verify by computing reference per-batch
    for b in 0..batch_size {
        let a_start = b * m * k;
        let b_start = b * k * n;
        let out_start = b * m * n;

        for i in 0..m {
            for j in 0..n {
                let mut expected = 0.0f32;
                for kk in 0..k {
                    expected += batched_a[a_start + i * k + kk] * batched_b[b_start + kk * n + j];
                }
                let actual = result[out_start + i * n + j];
                let diff = (expected - actual).abs();
                assert!(
                    diff < 1e-3,
                    "IMP-114a: Batch {} mismatch at ({},{}): expected={}, actual={}, diff={}",
                    b,
                    i,
                    j,
                    expected,
                    actual,
                    diff
                );
            }
        }
    }
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_114b_flattened_matches_loop() {
    // IMP-114b: Verify flattened approach matches loop-based approach
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 8,
        num_kv_heads: 8,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    let batch_size = 8;
    let m = 16;
    let k = 8;
    let n = 16;

    let batched_a: Vec<f32> = (0..batch_size * m * k)
        .map(|i| ((i % 17) as f32 - 8.0) * 0.05)
        .collect();
    let batched_b: Vec<f32> = (0..batch_size * k * n)
        .map(|i| ((i % 19) as f32 - 9.0) * 0.05)
        .collect();

    // Loop-based (IMP-113)
    let loop_result = cached_model
        .batched_gemm_single_dispatch(&batched_a, &batched_b, batch_size, m, k, n)
        .expect("Loop GEMM should succeed");

    // Flattened (IMP-114)
    let flat_result = cached_model
        .flattened_batched_gemm(&batched_a, &batched_b, batch_size, m, k, n)
        .expect("Flattened GEMM should succeed");

    assert_eq!(loop_result.len(), flat_result.len());
    for i in 0..loop_result.len() {
        let diff = (loop_result[i] - flat_result[i]).abs();
        assert!(
            diff < 1e-3,
            "IMP-114b: Results differ at {}: loop={}, flat={}, diff={}",
            i,
            loop_result[i],
            flat_result[i],
            diff
        );
    }
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_114c_flattened_attention_correctness() {
    // IMP-114c: Verify flattened attention matches reference
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 64,
        intermediate_dim: 128,
        num_layers: 1,
        num_heads: 4,
        num_kv_heads: 4,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
        bos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let cached_model = OwnedQuantizedModelCached::new(model);

    let seq_len = 8;
    let hidden_dim = config.hidden_dim;

    let q: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 13) as f32 - 6.0) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 11) as f32 - 5.0) * 0.1)
        .collect();
    let v: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 7) as f32 - 3.0) * 0.1)
        .collect();

    // Reference (IMP-113 single dispatch)
    let reference = cached_model
        .single_dispatch_multihead_attention(&q, &k, &v, seq_len)
        .expect("Reference attention should succeed");

    // Flattened (IMP-114)
    let result = cached_model
        .flattened_multihead_attention(&q, &k, &v, seq_len)
        .expect("Flattened attention should succeed");

    assert_eq!(result.len(), reference.len());
    for i in 0..result.len() {
        let diff = (result[i] - reference[i]).abs();
        assert!(
            diff < 1e-3,
            "IMP-114c: Attention differs at {}: ref={}, flat={}, diff={}",
            i,
            reference[i],
            result[i],
            diff
        );
    }
}

#[test]
#[cfg(feature = "gpu")]
fn test_imp_114d_large_batch_flattened() {
    // IMP-114d: Test with larger batch sizes where flattening benefits
    let config = GGUFConfig {
        architecture: "test".to_string(),
        hidden_dim: 128,
        intermediate_dim: 256,
        num_layers: 1,
        num_heads: 16, // Larger number of heads
        num_kv_heads: 16,
        vocab_size: 50,
        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 cached_model = OwnedQuantizedModelCached::new(model);

    let batch_size = 16;
    let m = 8;
    let k = 8;
    let n = 8;

    let batched_a: Vec<f32> = (0..batch_size * m * k)
        .map(|i| ((i % 23) as f32 - 11.0) * 0.04)
        .collect();
    let batched_b: Vec<f32> = (0..batch_size * k * n)
        .map(|i| ((i % 29) as f32 - 14.0) * 0.04)
        .collect();

    let result = cached_model
        .flattened_batched_gemm(&batched_a, &batched_b, batch_size, m, k, n)
        .expect("Large batch flattened GEMM should succeed");

    assert_eq!(
        result.len(),
        batch_size * m * n,
        "IMP-114d: Output should have correct dimensions"
    );

    // Verify non-trivial output
    let sum: f32 = result.iter().map(|x| x.abs()).sum();
    assert!(
        sum > 0.01,
        "IMP-114d: Output should have non-trivial values, got sum={}",
        sum
    );
}
