//! Edge case tests for Flash Attention - numerical stability and boundary conditions
//!
//! These tests validate Flash Attention behavior under extreme conditions that commonly
//! cause numerical instability or failure in attention implementations.

use rtx_flash_attention::*;
use rtx_tensor::{DType, Device, Tensor};

/// Test numerical stability with very small softmax scales
#[test]
fn test_numerical_stability_small_softmax_scale() {
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let mut config = FlashAttentionConfig::new(8, 64);
    config.softmax_scale = Some(1e-8); // Extremely small scale

    let q = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
    let k = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
    let v = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();

    // This should not crash or produce NaN/Inf values
    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(
        result.is_ok(),
        "Forward pass should handle small softmax scale"
    );

    let output = result.unwrap();
    assert!(
        is_finite(&output),
        "Output should be finite with small softmax scale"
    );
    assert!(!has_nan(&output), "Output should not contain NaN values");

    println!("✓ Small softmax scale test passed");
}

/// Test numerical stability with very large softmax scales
#[test]
fn test_numerical_stability_large_softmax_scale() {
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let mut config = FlashAttentionConfig::new(8, 64);
    config.softmax_scale = Some(1e8); // Extremely large scale

    let q = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
    let k = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
    let v = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();

    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(
        result.is_ok(),
        "Forward pass should handle large softmax scale"
    );

    let output = result.unwrap();
    assert!(
        is_finite(&output),
        "Output should be finite with large softmax scale"
    );
    assert!(!has_nan(&output), "Output should not contain NaN values");

    println!("✓ Large softmax scale test passed");
}

/// Test boundary condition: single token sequence
#[test]
fn test_single_token_sequence() {
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let config = FlashAttentionConfig::new(8, 64);

    // Single token in sequence
    let q = Tensor::randn(&[1, 8, 1, 64], &device).unwrap();
    let k = Tensor::randn(&[1, 8, 1, 64], &device).unwrap();
    let v = Tensor::randn(&[1, 8, 1, 64], &device).unwrap();

    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(result.is_ok(), "Should handle single token sequences");

    let output = result.unwrap();
    assert_eq!(output.shape().dims(), &[1, 8, 1, 64]);
    assert!(is_finite(&output), "Single token output should be finite");

    println!("✓ Single token sequence test passed");
}

/// Test boundary condition: very long sequences testing memory efficiency
#[test]
fn test_long_sequence_memory_efficiency() {
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let config = FlashAttentionConfig::new(8, 64);

    // Very long sequence that would be prohibitive for standard attention O(n²) memory
    let seq_len = 8192; // 8K sequence length
    let q = Tensor::randn(&[1, 8, seq_len, 64], &device).unwrap();
    let k = Tensor::randn(&[1, 8, seq_len, 64], &device).unwrap();
    let v = Tensor::randn(&[1, 8, seq_len, 64], &device).unwrap();

    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(result.is_ok(), "Should handle long sequences efficiently");

    let output = result.unwrap();
    assert_eq!(output.shape().dims(), &[1, 8, seq_len, 64]);
    assert!(is_finite(&output), "Long sequence output should be finite");

    println!("✓ Long sequence memory efficiency test passed");
}

/// Test numerical stability with extreme input values
#[test]
fn test_extreme_input_values() {
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let config = FlashAttentionConfig::new(8, 64);

    // Create tensors with extreme values
    let q = create_extreme_tensor(&[2, 8, 128, 64], &device, 100.0);
    let k = create_extreme_tensor(&[2, 8, 128, 64], &device, -100.0);
    let v = create_extreme_tensor(&[2, 8, 128, 64], &device, 50.0);

    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(result.is_ok(), "Should handle extreme input values");

    let output = result.unwrap();
    assert!(
        is_finite(&output),
        "Output should be finite with extreme inputs"
    );
    assert!(
        !has_nan(&output),
        "Output should not contain NaN with extreme inputs"
    );

    println!("✓ Extreme input values test passed");
}

/// Test causal masking correctness for different sequence lengths
#[test]
fn test_causal_masking_correctness() {
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let mut config = FlashAttentionConfig::new(4, 32);
    config.causal = true;

    let seq_len = 16;
    let q = Tensor::randn(&[1, 4, seq_len, 32], &device).unwrap();
    let k = Tensor::randn(&[1, 4, seq_len, 32], &device).unwrap();
    let v = Tensor::randn(&[1, 4, seq_len, 32], &device).unwrap();

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

    config.causal = false;
    let non_causal_result = flash_attention_forward(&q, &k, &v, &config).unwrap();

    // Causal and non-causal should produce different results
    assert!(
        !tensors_equal(&causal_result, &non_causal_result),
        "Causal and non-causal attention should produce different results"
    );

    assert!(is_finite(&causal_result), "Causal result should be finite");
    assert!(
        is_finite(&non_causal_result),
        "Non-causal result should be finite"
    );

    println!("✓ Causal masking correctness test passed");
}

/// Test gradient flow in backward pass
#[test]
fn test_gradient_flow_correctness() {
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let config = FlashAttentionConfig::new(4, 32);

    let q = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();
    let k = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();
    let v = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();
    let grad_output = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();

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

    // Gradients should have the same shape as inputs
    assert_eq!(grad_q.shape().dims(), q.shape().dims());
    assert_eq!(grad_k.shape().dims(), k.shape().dims());
    assert_eq!(grad_v.shape().dims(), v.shape().dims());

    // Gradients should be finite and non-zero (indicating proper flow)
    assert!(
        is_finite(&grad_q) && !is_zero(&grad_q),
        "grad_q should be finite and non-zero"
    );
    assert!(
        is_finite(&grad_k) && !is_zero(&grad_k),
        "grad_k should be finite and non-zero"
    );
    assert!(
        is_finite(&grad_v) && !is_zero(&grad_v),
        "grad_v should be finite and non-zero"
    );

    println!("✓ Gradient flow correctness test passed");
}

/// Test memory efficiency by comparing Flash vs naive attention memory usage
#[test]
fn test_memory_efficiency_validation() {
    let batch_size = 2;
    let num_heads = 8;
    let seq_len = 2048;
    let head_dim = 64;

    let (standard_memory, flash_memory) =
        utils::compare_memory_usage(batch_size, num_heads, seq_len, head_dim);

    // Flash Attention should use significantly less memory for long sequences
    let memory_reduction = standard_memory as f64 / flash_memory as f64;
    assert!(
        memory_reduction > 2.0,
        "Flash Attention should reduce memory usage by at least 2x for seq_len=2048, got {:.2}x",
        memory_reduction
    );

    println!(
        "✓ Memory efficiency validation passed: {:.2}x reduction",
        memory_reduction
    );
}

// Helper functions for edge case testing

/// Check if all values in tensor are finite (not NaN or Inf)
fn is_finite(tensor: &Tensor) -> bool {
    match tensor.to_cpu() {
        Ok(data) => data.iter().all(|&x| x.is_finite()),
        Err(_) => false, // If we can't get CPU data, assume not finite
    }
}

/// Check if tensor contains any NaN values
fn has_nan(tensor: &Tensor) -> bool {
    match tensor.to_cpu() {
        Ok(data) => data.iter().any(|&x| x.is_nan()),
        Err(_) => true, // If we can't get CPU data, assume NaN for safety
    }
}

/// Check if tensor is all zeros
fn is_zero(tensor: &Tensor) -> bool {
    match tensor.to_cpu() {
        Ok(data) => data.iter().all(|&x| x == 0.0),
        Err(_) => false,
    }
}

/// Check if two tensors are equal within tolerance
fn tensors_equal(a: &Tensor, b: &Tensor) -> bool {
    if a.shape() != b.shape() {
        return false;
    }

    match (a.to_cpu(), b.to_cpu()) {
        (Ok(data_a), Ok(data_b)) => {
            const TOLERANCE: f32 = 1e-6;
            data_a
                .iter()
                .zip(data_b.iter())
                .all(|(&x, &y)| (x - y).abs() < TOLERANCE)
        }
        _ => false,
    }
}

/// Create tensor with extreme values for stress testing
fn create_extreme_tensor(shape: &[usize], device: &Device, scale: f32) -> Tensor {
    // Create a tensor with random values and scale them to extreme values
    let tensor = Tensor::randn(shape, device).unwrap();

    // For now, return the tensor as-is since we don't have a scalar multiply operation
    // This will be a test that should fail in RED phase until we implement proper scaling
    tensor
}