use rtx_multimodal::fusion::{FusionConfig, ModalityFusion};
use rtx_tensor::{DType, Device, Tensor};

#[cfg(test)]
mod fusion_tests {
    use super::*;

    #[test]
    fn test_modality_fusion_creation() {
        // RED: Write failing test first
        let device = Device::default();
        let config = FusionConfig::new(768, 768, 768, 768, 12);

        let fusion = ModalityFusion::new(config, &device);
        assert!(fusion.is_ok(), "Should create ModalityFusion successfully");

        let fusion = fusion.unwrap();
        assert_eq!(fusion.config().vision_dim, 768);
        assert_eq!(fusion.config().audio_dim, 768);
        assert_eq!(fusion.config().text_dim, 768);
        assert_eq!(fusion.config().output_dim, 768);
        assert_eq!(fusion.config().num_heads, 12);
    }

    #[test]
    fn test_bimodal_fusion_vision_audio() {
        // RED: Test bimodal fusion
        let device = Device::default();
        let config = FusionConfig::new(512, 512, 512, 512, 8);
        let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");

        // Create test tensors
        let batch_size = 2;
        let seq_len_vision = 197; // 196 patches + 1 cls token
        let seq_len_audio = 500;
        let hidden_dim = 512;

        let vision = Tensor::randn(&[batch_size, seq_len_vision, hidden_dim], &device)
            .expect("Failed to create vision tensor");
        let audio = Tensor::randn(&[batch_size, seq_len_audio, hidden_dim], &device)
            .expect("Failed to create audio tensor");

        // Create dummy text tensor for trimodal fusion
        let text = Tensor::randn(&[batch_size, 50, hidden_dim], &device)
            .expect("Failed to create text tensor");
        let result = fusion.forward_trimodal(&vision, &audio, &text);
        assert!(result.is_ok(), "Bimodal fusion should succeed");

        let output = result.unwrap();
        assert_eq!(output.shape()[0], batch_size);
        assert_eq!(output.shape()[1], hidden_dim);
    }

    #[test]
    fn test_bimodal_fusion_vision_text() {
        // RED: Test vision-text fusion
        let device = Device::default();
        let config = FusionConfig::new(384, 384, 384, 384, 6);
        let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");

        let batch_size = 4;
        let seq_len_vision = 100;
        let seq_len_text = 64;
        let hidden_dim = 384;

        let vision = Tensor::randn(&[batch_size, seq_len_vision, hidden_dim], &device)
            .expect("Failed to create vision tensor");
        let text = Tensor::randn(&[batch_size, seq_len_text, hidden_dim], &device)
            .expect("Failed to create text tensor");

        // Create dummy audio tensor for trimodal fusion
        let audio = Tensor::randn(&[batch_size, 80, hidden_dim], &device)
            .expect("Failed to create audio tensor");
        let result = fusion.forward_trimodal(&vision, &audio, &text);
        assert!(result.is_ok(), "Vision-text fusion should succeed");

        let output = result.unwrap();
        assert_eq!(output.shape()[0], batch_size);
        assert_eq!(output.shape()[1], hidden_dim);
    }

    #[test]
    fn test_trimodal_fusion() {
        // RED: Test trimodal fusion
        let device = Device::default();
        let config = FusionConfig::new(256, 256, 256, 256, 4);
        let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");

        let batch_size = 1;
        let vision =
            Tensor::randn(&[batch_size, 50, 256], &device).expect("Failed to create vision tensor");
        let audio =
            Tensor::randn(&[batch_size, 100, 256], &device).expect("Failed to create audio tensor");
        let text =
            Tensor::randn(&[batch_size, 30, 256], &device).expect("Failed to create text tensor");

        let result = fusion.forward_trimodal(&vision, &audio, &text);
        assert!(result.is_ok(), "Trimodal fusion should succeed");

        let output = result.unwrap();
        assert_eq!(output.shape()[0], batch_size);
        assert_eq!(output.shape()[1], 256); // output_dim
    }

    #[test]
    fn test_fusion_strategies() {
        // RED: Test different fusion strategies
        use rtx_multimodal::fusion::ModalityFusionStrategy;

        let device = Device::default();
        let strategies = vec![
            ModalityFusionStrategy::Concatenation,
            ModalityFusionStrategy::BilinearFusion,
            ModalityFusionStrategy::TensorFusion,
            ModalityFusionStrategy::AttentionFusion,
            ModalityFusionStrategy::HierarchicalAttention,
            ModalityFusionStrategy::GatedFusion,
            ModalityFusionStrategy::ContrastiveFusion,
        ];

        for strategy in strategies {
            let mut config = FusionConfig::new(128, 128, 128, 128, 2);
            config.fusion_strategy = strategy.clone();

            let mut fusion = ModalityFusion::new(config, &device).expect(&format!(
                "Failed to create fusion with strategy {:?}",
                strategy
            ));

            let vision = Tensor::randn(&[1, 10, 128], &device).unwrap();
            let audio = Tensor::randn(&[1, 20, 128], &device).unwrap();
            let text = Tensor::randn(&[1, 15, 128], &device).unwrap();

            let result = fusion.forward_trimodal(&vision, &audio, &text);
            assert!(result.is_ok(), "Strategy {:?} should work", strategy);

            let output = result.unwrap();
            assert_eq!(output.shape()[0], 1);
            assert_eq!(output.shape()[1], 128);
        }
    }

    #[test]
    fn test_fusion_output_consistency() {
        // RED: Test that fusion output is deterministic with same input
        let device = Device::default();
        let config = FusionConfig::new(64, 64, 64, 64, 2);
        let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");

        // Use zeros for deterministic testing
        let vision = Tensor::zeros(&[1, 5, 64], &device).expect("Failed to create vision tensor");
        let audio = Tensor::zeros(&[1, 10, 64], &device).expect("Failed to create audio tensor");
        let text = Tensor::zeros(&[1, 8, 64], &device).expect("Failed to create text tensor");

        let result1 = fusion
            .forward_trimodal(&vision, &audio, &text)
            .expect("First forward pass failed");
        let result2 = fusion
            .forward_trimodal(&vision, &audio, &text)
            .expect("Second forward pass failed");

        // With zero inputs and no randomness, outputs should be deterministic
        // This tests that the fusion mechanism itself is stable
        assert_eq!(result1.shape(), result2.shape());
    }

    #[test]
    fn test_empty_batch() {
        // RED: Test edge case with batch size 0
        let device = Device::default();
        let config = FusionConfig::new(64, 64, 64, 64, 2);
        let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");

        // Empty batch
        let vision =
            Tensor::zeros(&[0, 5, 64], &device).expect("Failed to create empty vision tensor");
        let audio =
            Tensor::zeros(&[0, 10, 64], &device).expect("Failed to create empty audio tensor");
        let text = Tensor::zeros(&[0, 8, 64], &device).expect("Failed to create empty text tensor");

        let result = fusion.forward_trimodal(&vision, &audio, &text);

        // Should handle empty batch gracefully
        assert!(result.is_ok(), "Should handle empty batch");
        let output = result.unwrap();
        assert_eq!(output.shape()[0], 0);
    }
}