//! Sliding Window Attention Feature Demonstration
//!
//! This demonstrates all the key features requested:
//! 1. Configurable window sizes
//! 2. Efficient computation for long sequences  
//! 3. Proper boundary handling at sequence start/end
//! 4. Option for overlapping windows (bidirectional)
//! 5. Integration with existing attention types (MQA/GQA)

use rtx_transformers::layers::{SlidingWindowAttention, SlidingWindowConfig, LayerNorm};
use rtx_transformers::architectures::TransformerConfig;
use rtx_transformers::{Result, TransformerError};
use rtx_tensor::{Tensor, Device, DType};
use rtx_autograd::{TensorAutograd, backward};

fn main() -> Result<()> {
    println!("🚀 Sliding Window Attention Demonstration");
    println!("==========================================");
    
    let device = Device::cuda(0).unwrap_or(Device::default());
    
    // Feature 1: Configurable window sizes
    demonstrate_configurable_window_sizes(&device)?;
    
    // Feature 2: Efficient computation for long sequences
    demonstrate_long_sequence_efficiency(&device)?;
    
    // Feature 3: Proper boundary handling
    demonstrate_boundary_handling(&device)?;
    
    // Feature 4: Overlapping windows (bidirectional vs causal)
    demonstrate_overlapping_windows(&device)?;
    
    // Feature 5: Integration with MQA/GQA
    demonstrate_mqa_gqa_integration(&device)?;
    
    // Bonus: Full transformer block simulation
    demonstrate_transformer_block_integration(&device)?;
    
    println!("\n✅ All sliding window attention features demonstrated successfully!");
    println!("\n📊 Summary:");
    println!("  • Configurable window sizes (16, 64, 256, etc.)");
    println!("  • Efficient O(n*w) complexity instead of O(n²)");
    println!("  • Proper boundary handling for short sequences");
    println!("  • Causal and bidirectional attention modes");
    println!("  • Full compatibility with MQA and GQA");
    println!("  • Integration with existing transformer infrastructure");
    
    Ok(())
}

fn demonstrate_configurable_window_sizes(device: &Device) -> Result<()> {
    println!("\n1. 🎛️  Configurable Window Sizes");
    println!("   Testing different window sizes for efficiency trade-offs");
    
    let window_sizes = vec![16, 64, 256];
    let seq_len = 512;
    
    for window_size in window_sizes {
        let config = SlidingWindowConfig::new(8, 512, window_size)?;
        let mut attention = SlidingWindowAttention::new(config, device)?;
        attention.initialize_parameters()?;
        
        let input = Tensor::ones_typed(&[1, seq_len, 512], DType::F32, device)?;
        let output = attention.forward(&input, None, None)?;
        
        let memory_savings = attention.memory_savings_ratio(seq_len);
        let effective_span = attention.effective_attention_span(seq_len);
        
        println!("   Window size {:3}: Memory savings {:5.1}%, Effective span {:3}", 
                window_size, memory_savings * 100.0, effective_span);
    }
    
    Ok(())
}

fn demonstrate_long_sequence_efficiency(device: &Device) -> Result<()> {
    println!("\n2. ⚡ Long Sequence Efficiency");
    println!("   Comparing efficiency for increasingly long sequences");
    
    let config = SlidingWindowConfig::new(12, 768, 128)?; // 128 token window
    let mut attention = SlidingWindowAttention::new(config, device)?;
    attention.initialize_parameters()?;
    
    let sequence_lengths = vec![256, 512, 1024, 2048, 4096];
    
    for seq_len in sequence_lengths {
        let input = Tensor::ones_typed(&[1, seq_len, 768], DType::F32, device)?;
        
        let start = std::time::Instant::now();
        let output = attention.forward(&input, None, None)?;
        let duration = start.elapsed();
        
        let memory_savings = attention.memory_savings_ratio(seq_len);
        
        println!("   Seq len {:4}: {:6.2}ms, {:5.1}% memory saved, O({}*128) vs O({}²)", 
                seq_len, 
                duration.as_secs_f64() * 1000.0,
                memory_savings * 100.0,
                seq_len,
                seq_len);
    }
    
    Ok(())
}

fn demonstrate_boundary_handling(device: &Device) -> Result<()> {
    println!("\n3. 🎯 Boundary Handling");
    println!("   Testing behavior at sequence boundaries");
    
    let window_size = 64;
    let config = SlidingWindowConfig::new(8, 512, window_size)?;
    let mut attention = SlidingWindowAttention::new(config, device)?;
    attention.initialize_parameters()?;
    
    // Test cases: sequence lengths relative to window size
    let test_cases = vec![
        (32, "Short (< window)"),
        (64, "Exact window size"),
        (128, "Long (> window)"),
    ];
    
    for (seq_len, description) in test_cases {
        let input = Tensor::ones_typed(&[1, seq_len, 512], DType::F32, device)?;
        let output = attention.forward(&input, None, None)?;
        
        let effective_span = attention.effective_attention_span(seq_len);
        
        println!("   {:18}: seq_len={:3}, effective_span={:3}, window_size={:2}", 
                description, seq_len, effective_span, window_size);
        
        // Verify output shape is preserved
        assert_eq!(output.shape().dims(), &[1, seq_len, 512]);
    }
    
    Ok(())
}

fn demonstrate_overlapping_windows(device: &Device) -> Result<()> {
    println!("\n4. 🔄 Overlapping Windows (Causal vs Bidirectional)");
    println!("   Comparing causal and bidirectional attention modes");
    
    let window_size = 32;
    let seq_len = 128;
    
    // Causal (traditional autoregressive)
    let causal_config = SlidingWindowConfig::new(8, 512, window_size)?
        .with_causal(true);
    let mut causal_attention = SlidingWindowAttention::new(causal_config, device)?;
    causal_attention.initialize_parameters()?;
    
    // Bidirectional (encoder-style)
    let bidirectional_config = SlidingWindowConfig::new(8, 512, window_size)?
        .with_causal(false);
    let mut bidirectional_attention = SlidingWindowAttention::new(bidirectional_config, device)?;
    bidirectional_attention.initialize_parameters()?;
    
    let input = Tensor::ones_typed(&[1, seq_len, 512], DType::F32, device)?;
    
    let causal_output = causal_attention.forward(&input, None, None)?;
    let bidirectional_output = bidirectional_attention.forward(&input, None, None)?;
    
    let causal_span = causal_attention.effective_attention_span(seq_len);
    let bidirectional_span = bidirectional_attention.effective_attention_span(seq_len);
    
    println!("   Causal mode:       effective_span={:2} (window_size + 1)", causal_span);
    println!("   Bidirectional mode: effective_span={:2} (2 * window_size + 1)", bidirectional_span);
    
    // Verify outputs are different
    let causal_sum = causal_output.sum_all()?.to_scalar::<f32>()?;
    let bidirectional_sum = bidirectional_output.sum_all()?.to_scalar::<f32>()?;
    
    println!("   Output sums differ: causal={:.2}, bidirectional={:.2}", 
            causal_sum, bidirectional_sum);
    
    Ok(())
}

fn demonstrate_mqa_gqa_integration(device: &Device) -> Result<()> {
    println!("\n5. 🤝 MQA/GQA Integration");
    println!("   Testing integration with Multi-Query and Grouped-Query Attention");
    
    let window_size = 64;
    
    // Standard Multi-Head Attention (baseline)
    let mha_config = SlidingWindowConfig::new(12, 768, window_size)?;
    let mha_attention = SlidingWindowAttention::new(mha_config, device)?;
    
    // Multi-Query Attention (1 KV head)
    let mqa_config = SlidingWindowConfig::new(12, 768, window_size)?
        .with_kv_heads(1)?;
    let mqa_attention = SlidingWindowAttention::new(mqa_config, device)?;
    
    // Grouped-Query Attention (4 KV heads)
    let gqa_config = SlidingWindowConfig::new(12, 768, window_size)?
        .with_kv_heads(4)?;
    let gqa_attention = SlidingWindowAttention::new(gqa_config, device)?;
    
    println!("   MHA: {} query heads, {} KV heads", 
            mha_attention.config().num_heads, 
            mha_attention.config().get_num_kv_heads());
    
    println!("   MQA: {} query heads, {} KV heads ({}% KV memory vs MHA)", 
            mqa_attention.config().num_heads,
            mqa_attention.config().get_num_kv_heads(),
            (100 / 12)); // 1/12 of the KV memory
    
    println!("   GQA: {} query heads, {} KV heads ({}% KV memory vs MHA)", 
            gqa_attention.config().num_heads,
            gqa_attention.config().get_num_kv_heads(),
            (4 * 100 / 12)); // 4/12 of the KV memory
    
    // Test forward passes
    let input = Tensor::ones_typed(&[1, 128, 768], DType::F32, device)?;
    
    let mut mha_mut = SlidingWindowAttention::new(mha_config, device)?;
    let mut mqa_mut = SlidingWindowAttention::new(mqa_config, device)?;
    let mut gqa_mut = SlidingWindowAttention::new(gqa_config, device)?;
    
    mha_mut.initialize_parameters()?;
    mqa_mut.initialize_parameters()?;
    gqa_mut.initialize_parameters()?;
    
    let _mha_output = mha_mut.forward(&input, None, None)?;
    let _mqa_output = mqa_mut.forward(&input, None, None)?;
    let _gqa_output = gqa_mut.forward(&input, None, None)?;
    
    println!("   ✓ All attention variants process the same input successfully");
    
    Ok(())
}

fn demonstrate_transformer_block_integration(device: &Device) -> Result<()> {
    println!("\n6. 🏗️  Transformer Block Integration");
    println!("   Simulating full transformer block with sliding window attention");
    
    let config = SlidingWindowConfig::new(8, 512, 64)?;
    let mut attention = SlidingWindowAttention::new(config, device)?;
    attention.initialize_parameters()?;
    
    // Layer norms for transformer block
    let ln1 = LayerNorm::new(512, 1e-5, true, device)?;
    let ln2 = LayerNorm::new(512, 1e-5, true, device)?;
    
    let batch_size = 2;
    let seq_len = 128;
    let d_model = 512;
    
    let input = Tensor::ones_typed(&[batch_size, seq_len, d_model], DType::F32, device)?;
    let input_grad = input.set_requires_grad(true);
    
    // Transformer block: x = x + attention(ln(x)); x = x + mlp(ln(x))
    
    // Attention sublayer
    let norm1 = ln1.forward(&input_grad)?;
    let attn_out = attention.forward(&norm1, None, None)?;
    let residual1 = input_grad.add(&attn_out)?;
    
    // MLP sublayer (simplified)
    let norm2 = ln2.forward(&residual1)?;
    let mlp_out = norm2.mul_scalar(1.5)?.add(&norm2.mul_scalar(-0.5)?)?; // Simulate MLP
    let final_output = residual1.add(&mlp_out)?;
    
    println!("   Input shape:  {:?}", input_grad.shape().dims());
    println!("   Output shape: {:?}", final_output.shape().dims());
    
    // Test gradient flow
    let loss = final_output.mean_all()?;
    let gradients = backward(&loss, &[&input_grad])?;
    
    let grad_norm = gradients[0].sum_all()?.to_scalar::<f32>()?.abs();
    println!("   Gradient norm: {:.6} (indicating proper gradient flow)", grad_norm);
    
    println!("   ✓ Full transformer block with sliding window attention works!");
    
    Ok(())
}

#[cfg(test)]
mod demo_tests {
    use super::*;
    
    #[test]
    fn test_demo_runs_successfully() {
        // Ensure the demo runs without panics
        let result = main();
        assert!(result.is_ok(), "Demo should run successfully: {:?}", result);
    }
}