//! Comprehensive kernel validation tests for Flash Attention
//! Tests the correctness of Flash Attention kernel implementations
//!
//! NOTE: Disabled until Flash Attention API is fully implemented

#![cfg(all(feature = "cuda", feature = "disabled_tests"))]

use approx::assert_relative_eq;
use proptest::prelude::*;
use rtx_flash_attention::*;
use rtx_tensor::Tensor;

/// Test basic Flash Attention forward pass correctness
#[test]
fn test_flash_attention_forward_basic() {
    let seq_len = 128;
    let num_heads = 8;
    let batch_size = 2;

    let mut config = FlashAttentionConfig::new(num_heads, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = false;
    config.softmax_scale = None;

    // Create input tensors
    let q = Tensor::randn(&[batch_size, num_heads, seq_len, config.head_dim]);
    let k = Tensor::randn(&[batch_size, num_heads, seq_len, config.head_dim]);
    let v = Tensor::randn(&[batch_size, num_heads, seq_len, config.head_dim]);

    let result = flash_attention_forward(&q, &k, &v, &config).unwrap();

    // Verify output shape
    assert_eq!(
        result.shape(),
        &[batch_size, num_heads, seq_len, config.head_dim]
    );

    // Verify output is finite
    assert!(result.all_finite().unwrap());
}

/// Test Flash Attention with causal masking
#[test]
fn test_flash_attention_causal() {
    let seq_len = 64;

    let mut config = FlashAttentionConfig::new(1, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = true;
    config.softmax_scale = Some(0.125);
    let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);

    let result = flash_attention_forward(&q, &k, &v, &config).unwrap();

    // For causal attention, verify triangular structure
    assert_eq!(result.shape(), &[1, 1, seq_len, config.head_dim]);
    assert!(result.all_finite().unwrap());
}

/// Test Flash Attention backward pass
#[test]
fn test_flash_attention_backward() {
    let seq_len = 32;

    let mut config = FlashAttentionConfig::new(1, 32);
    config.block_size_q = 16;
    config.block_size_kv = 16;
    config.causal = false;
    config.softmax_scale = None;
    let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let grad_out = Tensor::randn(&[1, 1, seq_len, config.head_dim]);

    let (grad_q, grad_k, grad_v) =
        flash_attention_backward(&grad_out, &q, &k, &v, &config).unwrap();

    // Verify gradient shapes
    assert_eq!(grad_q.shape(), q.shape());
    assert_eq!(grad_k.shape(), k.shape());
    assert_eq!(grad_v.shape(), v.shape());

    // Verify gradients are finite
    assert!(grad_q.all_finite().unwrap());
    assert!(grad_k.all_finite().unwrap());
    assert!(grad_v.all_finite().unwrap());
}

/// Test memory efficiency for large sequences
#[test]
fn test_memory_efficiency_large_sequence() {
    let seq_len = 2048; // Large sequence

    let mut config = FlashAttentionConfig::new(1, 64);
    config.block_size_q = 64;
    config.block_size_kv = 64;
    config.causal = false;
    config.softmax_scale = None;
    let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);

    let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
    assert_eq!(result.shape(), &[1, 1, seq_len, config.head_dim]);
}

/// Test numerical equivalence with reference implementation
#[test]
fn test_numerical_equivalence() {
    let seq_len = 16;

    let mut config = FlashAttentionConfig::new(1, 16);
    config.block_size_q = 8;
    config.block_size_kv = 8;
    config.causal = false;
    config.softmax_scale = Some(0.25);
    let q = Tensor::ones(&[1, 1, seq_len, config.head_dim]);
    let k = Tensor::ones(&[1, 1, seq_len, config.head_dim]);
    let v = Tensor::ones(&[1, 1, seq_len, config.head_dim]);

    let flash_result = flash_attention_forward(&q, &k, &v, &config).unwrap();
    let reference_result = reference_attention(&q, &k, &v, config.softmax_scale).unwrap();

    // Compare results with tolerance
    for i in 0..flash_result.numel() {
        assert_relative_eq!(
            flash_result.get_item(i).unwrap(),
            reference_result.get_item(i).unwrap(),
            epsilon = 1e-4
        );
    }
}

/// Property-based test for Flash Attention invariants
proptest! {
    #[test]
    fn test_flash_attention_properties(
    batch_size in 1..4usize,
    num_heads in 1..8usize,
    seq_len in 16..128usize,
    head_dim in prop::sample::select(vec![16, 32, 64]),
    ) {
        let mut config = FlashAttentionConfig::new(num_heads, head_dim);
        config.block_size_q = 32;
        config.block_size_kv = 32;
        config.causal = false;
        config.softmax_scale = None;

        let q = Tensor::randn(&[batch_size, num_heads, seq_len, head_dim]);
        let k = Tensor::randn(&[batch_size, num_heads, seq_len, head_dim]);
        let v = Tensor::randn(&[batch_size, num_heads, seq_len, head_dim]);

        let result = flash_attention_forward(&q, &k, &v, &config).unwrap();

        // Property: Output shape matches expected
        prop_assert_eq!(result.shape(), &[batch_size, num_heads, seq_len, head_dim]);

        // Property: Output is finite
        prop_assert!(result.all_finite().unwrap());

        // Property: Output norm is reasonable (not NaN/Inf)
        let output_norm = result.norm().unwrap();
        prop_assert!(output_norm.is_finite());
        prop_assert!(output_norm > 0.0);
    }
}

/// Test Flash Attention with different block sizes
#[test]
fn test_different_block_sizes() {
    let block_sizes = vec![(16, 16), (32, 32), (64, 64), (128, 128)];
    let seq_len = 256;

    let q = Tensor::randn(&[1, 1, seq_len, 64]);
    let k = Tensor::randn(&[1, 1, seq_len, 64]);
    let v = Tensor::randn(&[1, 1, seq_len, 64]);

    for (block_q, block_k) in block_sizes {
        if block_q <= seq_len && block_k <= seq_len {
            let mut config = FlashAttentionConfig::new(1, 64);
            config.block_size_q = block_q;
            config.block_size_kv = block_k;
            config.causal = false;
            config.softmax_scale = None;

            let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
            assert_eq!(result.shape(), &[1, 1, seq_len, 64]);
        }
    }
}

/// Test online softmax correctness
#[test]
fn test_online_softmax() {
    use rtx_flash_attention::kernels::utils::online_softmax;

    let input = vec![1.0, 2.0, 3.0, 4.0];
    let result = online_softmax(&input);

    // Verify softmax properties
    let sum: f32 = result.iter().sum();
    assert_relative_eq!(sum, 1.0, epsilon = 1e-6);

    // Verify all values are positive
    for &val in &result {
        assert!(val > 0.0);
    }

    // Verify monotonicity for this input
    for i in 1..result.len() {
        assert!(result[i] > result[i - 1]);
    }
}

/// Test Flash Attention with mixed precision
#[test]
fn test_mixed_precision() {
    let seq_len = 128;

    let mut config = FlashAttentionConfig::new(1, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = false;
    config.softmax_scale = None;

    // Test with fp16 inputs (simulated with fp32)
    let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);

    let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
    assert!(result.all_finite().unwrap());
}

/// Test Flash Attention variants (edge, neuromorphic, quantum)
#[test]
fn test_flash_attention_variants() {
    use rtx_flash_attention::variants::*;

    let seq_len = 64;
    let head_dim = 32;

    let q = Tensor::randn(&[1, 1, seq_len, head_dim]);
    let k = Tensor::randn(&[1, 1, seq_len, head_dim]);
    let v = Tensor::randn(&[1, 1, seq_len, head_dim]);

    // Test edge variant
    let edge_result = edge::flash_attention_edge(&q, &k, &v).unwrap();
    assert_eq!(edge_result.shape(), &[1, 1, seq_len, head_dim]);

    // Test neuromorphic variant
    let neuro_result = neuromorphic::flash_attention_neuromorphic(&q, &k, &v).unwrap();
    assert_eq!(neuro_result.shape(), &[1, 1, seq_len, head_dim]);

    // Test quantum variant
    let quantum_result = quantum::flash_attention_quantum(&q, &k, &v).unwrap();
    assert_eq!(quantum_result.shape(), &[1, 1, seq_len, head_dim]);
}

/// Helper function for reference attention implementation
fn reference_attention(
    q: &Tensor,
    k: &Tensor,
    v: &Tensor,
    scale: Option<f32>,
) -> Result<Tensor, Box<dyn std::error::Error>> {
    // Simple reference implementation for testing
    let scale = scale.unwrap_or(1.0 / (q.shape()[3] as f32).sqrt());

    // QK^T
    let scores = q.matmul(&k.transpose(-2, -1)?)?;
    let scaled_scores = scores.mul_scalar(scale)?;

    // Softmax
    let attention_probs = scaled_scores.softmax(-1)?;

    // Apply to values
    let output = attention_probs.matmul(v)?;
    Ok(output)
}

/// Benchmark comparison test
#[test]
fn test_performance_improvement() {
    let seq_len = 1024;

    let mut config = FlashAttentionConfig::new(8, 64);
    config.block_size_q = 64;
    config.block_size_kv = 64;
    config.causal = false;
    config.softmax_scale = None;
    let q = Tensor::randn(&[1, 8, seq_len, config.head_dim]);
    let k = Tensor::randn(&[1, 8, seq_len, config.head_dim]);
    let v = Tensor::randn(&[1, 8, seq_len, config.head_dim]);

    let start = std::time::Instant::now();
    let _flash_result = flash_attention_forward(&q, &k, &v, &config).unwrap();
    let flash_time = start.elapsed();

    let start = std::time::Instant::now();
    let _reference_result = reference_attention(&q, &k, &v, config.softmax_scale).unwrap();
    let reference_time = start.elapsed();

    // Flash Attention should be faster for large sequences
    println!("Flash Attention time: {:?}", flash_time);
    println!("Reference time: {:?}", reference_time);

    // For large sequences, Flash Attention should show improvement
    if seq_len > 512 {
        assert!(flash_time < reference_time * 2); // Allow some margin for test environment
    }
}

/// Test error handling
#[test]
fn test_error_handling() {
    let mut config = FlashAttentionConfig::new(1, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = false;
    config.softmax_scale = None;

    // Mismatched shapes should error
    let q = Tensor::randn(&[1, 1, 128, 64]);
    let k = Tensor::randn(&[1, 1, 64, 64]); // Different seq_len
    let v = Tensor::randn(&[1, 1, 128, 64]);

    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(result.is_err());

    // Mismatched head dimensions should error
    let q = Tensor::randn(&[1, 1, 128, 64]);
    let k = Tensor::randn(&[1, 1, 128, 32]); // Different head_dim
    let v = Tensor::randn(&[1, 1, 128, 64]);

    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(result.is_err());
}

/// Test gradient computation correctness
#[test]
fn test_gradient_correctness() {
    let seq_len = 16;

    let mut config = FlashAttentionConfig::new(1, 16);
    config.block_size_q = 8;
    config.block_size_kv = 8;
    config.causal = false;
    config.softmax_scale = Some(0.25);
    let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
    let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);

    // Forward pass
    let output = flash_attention_forward(&q, &k, &v, &config).unwrap();

    // Create gradient of loss w.r.t. output
    let grad_output = Tensor::ones_like(&output);

    // Backward pass
    let (grad_q, grad_k, grad_v) =
        flash_attention_backward(&grad_output, &q, &k, &v, &config).unwrap();

    // Verify gradient shapes match input shapes
    assert_eq!(grad_q.shape(), q.shape());
    assert_eq!(grad_k.shape(), k.shape());
    assert_eq!(grad_v.shape(), v.shape());

    // Verify gradients are reasonable (not too large)
    assert!(grad_q.abs().max().unwrap() < 100.0);
    assert!(grad_k.abs().max().unwrap() < 100.0);
    assert!(grad_v.abs().max().unwrap() < 100.0);
}

/// Test multi-head attention
#[test]
fn test_multi_head_attention() {
    let seq_len = 64;
    let num_heads = 12; // Large number of heads

    let mut config = FlashAttentionConfig::new(num_heads, 32);
    config.block_size_q = 16;
    config.block_size_kv = 16;
    config.causal = false;
    config.softmax_scale = None;

    let q = Tensor::randn(&[2, num_heads, seq_len, config.head_dim]);
    let k = Tensor::randn(&[2, num_heads, seq_len, config.head_dim]);
    let v = Tensor::randn(&[2, num_heads, seq_len, config.head_dim]);

    let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
    assert_eq!(result.shape(), &[2, num_heads, seq_len, config.head_dim]);

    // Verify each head produces reasonable outputs
    for head in 0..num_heads {
        let head_output = result.select(1, head).unwrap();
        assert!(head_output.all_finite().unwrap());
        assert!(head_output.norm().unwrap() > 0.0);
    }
}