//! Mamba (State Space Model) Demonstration
//!
//! This example demonstrates the Mamba architecture implementation,
//! showcasing its linear complexity and integration with transformer components.

use rtx_tensor::{Device, Tensor};
use rtx_transformers::layers::{MambaBlock, MambaConfig, MambaTransformer, MambaTransformerBlock};
use std::time::Instant;

/// Demonstrates basic Mamba block functionality
async fn demo_mamba_block() -> Result<(), Box<dyn std::error::Error>> {
    println!("🔬 Mamba Block Demo");
    println!("===================");

    let device = Device::cpu();

    // Create Mamba configuration
    let config = MambaConfig::new(512, 16, 4); // d_model=512, d_state=16, d_conv=4
    println!(
        "Configuration: d_model={}, d_state={}, d_conv={}",
        config.d_model, config.d_state, config.d_conv
    );

    // Initialize Mamba block
    let mamba = MambaBlock::new(config.clone(), &device)?;
    println!(
        "✅ Mamba block created with {} parameters",
        mamba.parameters().len()
    );

    // Test with different sequence lengths
    let sequence_lengths = vec![10, 50, 100, 200];
    println!("\n📊 Linear Complexity Demonstration:");

    for seq_len in sequence_lengths {
        let batch_size = 2;
        let input = Tensor::randn([batch_size, seq_len, config.d_model], &device)?;

        let start = Instant::now();
        let output = mamba.forward(&input)?;
        let elapsed = start.elapsed();

        println!(
            "Seq length {}: {:.2}ms (output shape: {:?})",
            seq_len,
            elapsed.as_millis(),
            output.output.shape().dims()
        );
    }

    Ok(())
}

/// Demonstrates hybrid Mamba-Transformer architecture
async fn demo_mamba_transformer() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n🏗️  Mamba Transformer Demo");
    println!("==========================");

    let device = Device::cpu();

    // Create hybrid transformer block
    let d_model = 256;
    let ff_dim = 1024;
    let mamba_config = MambaConfig::new(d_model, 16, 4);

    let hybrid_block = MambaTransformerBlock::new(d_model, ff_dim, mamba_config, &device)?;
    println!("✅ Hybrid Mamba-Transformer block created");

    // Test forward pass
    let batch_size = 1;
    let seq_len = 20;
    let input = Tensor::randn([batch_size, seq_len, d_model], &device)?;

    let start = Instant::now();
    let output = hybrid_block.forward(&input)?;
    let elapsed = start.elapsed();

    println!("Forward pass: {:.2}ms", elapsed.as_millis());
    println!("Input shape: {:?}", input.shape().dims());
    println!("Output shape: {:?}", output.shape().dims());

    // Verify residual connections work
    let output_data = output.to_vec()?;
    let input_data = input.to_vec()?;
    let are_different = output_data
        .iter()
        .zip(input_data.iter())
        .any(|(&a, &b)| (a - b).abs() > 1e-6);

    println!("Residual connection working: {}", are_different);

    Ok(())
}

/// Demonstrates full Mamba transformer model
async fn demo_full_transformer() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n🧠 Full Mamba Transformer Demo");
    println!("===============================");

    let device = Device::cpu();

    // Model configuration
    let vocab_size = 1000;
    let d_model = 128;
    let num_layers = 3;
    let max_seq_len = 100;

    // Create full transformer
    let transformer = MambaTransformer::new(vocab_size, d_model, num_layers, max_seq_len, &device)?;

    println!("✅ Full Mamba Transformer created:");
    println!("   - Vocabulary size: {}", vocab_size);
    println!("   - Model dimension: {}", d_model);
    println!("   - Number of layers: {}", num_layers);
    println!("   - Max sequence length: {}", max_seq_len);

    // Test inference
    let batch_size = 1;
    let seq_len = 15;
    let input_ids = Tensor::randint(0, vocab_size as i32, &[batch_size, seq_len], &device)?;

    let start = Instant::now();
    let logits = transformer.forward(&input_ids)?;
    let elapsed = start.elapsed();

    println!("Inference time: {:.2}ms", elapsed.as_millis());
    println!("Input shape: {:?}", input_ids.shape().dims());
    println!("Output logits shape: {:?}", logits.shape().dims());

    // Verify output is reasonable
    let logits_data = logits.to_vec()?;
    let max_logit = logits_data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
    let min_logit = logits_data.iter().fold(f32::INFINITY, |a, &b| a.min(b));

    println!("Logit range: [{:.3}, {:.3}]", min_logit, max_logit);
    println!(
        "All logits finite: {}",
        logits_data.iter().all(|&x| x.is_finite())
    );

    Ok(())
}

/// Compares Mamba efficiency with theoretical attention complexity
async fn demo_efficiency_comparison() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n⚡ Efficiency Comparison");
    println!("========================");

    let device = Device::cpu();
    let d_model = 256;
    let config = MambaConfig::new(d_model, 16, 4);
    let mamba = MambaBlock::new(config, &device)?;

    println!("Comparing Mamba (linear) vs theoretical Attention (quadratic):");
    println!("Seq Length | Mamba Time (ms) | Theoretical Attention FLOPs");
    println!("-----------|-----------------|-----------------------------");

    for seq_len in [50, 100, 200, 400, 800] {
        let input = Tensor::randn([1, seq_len, d_model], &device)?;

        let start = Instant::now();
        let _output = mamba.forward(&input)?;
        let mamba_time = start.elapsed().as_millis();

        // Theoretical attention FLOPs (simplified): O(seq_len^2 * d_model)
        let attention_flops = (seq_len * seq_len * d_model) as f64 / 1e6; // In millions

        println!(
            "{:10} | {:15} | {:.2}M FLOPs",
            seq_len, mamba_time, attention_flops
        );
    }

    println!("\n📈 Mamba scales linearly O(n) while attention scales quadratically O(n²)");

    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🐍 Mamba: Linear-Time Sequence Modeling Demo");
    println!("===========================================");
    println!("State Space Models for Transformers");
    println!("Paper: https://arxiv.org/abs/2312.00752\n");

    demo_mamba_block().await?;
    demo_mamba_transformer().await?;
    demo_full_transformer().await?;
    demo_efficiency_comparison().await?;

    println!("\n🎉 All Mamba demonstrations completed successfully!");
    println!("Key advantages demonstrated:");
    println!("- ✅ Linear complexity O(n) vs attention's O(n²)");
    println!("- ✅ Selective state space modeling");
    println!("- ✅ Integration with transformer architecture");
    println!("- ✅ Efficient caching for inference");

    Ok(())
}