//! Comprehensive tests for MobileNet implementation
//!
//! Following Test-Driven Development (TDD) methodology

#[cfg(test)]
mod tests {
    use rtx_vision::architectures::{MobileNet, MobileNetConfig, MobileNetVariant};
    use rtx_vision::{Device, Tensor};

    fn get_test_device() -> Device {
        Device::cpu()
    }

    #[test]
    fn test_mobilenet_config_default() {
        let config = MobileNetConfig::default();
        assert!(matches!(config.variant, MobileNetVariant::V1));
        assert_eq!(config.num_classes, 1000);
        assert_eq!(config.width_multiplier, 1.0);
        assert_eq!(config.resolution_multiplier, 1.0);
    }

    #[test]
    fn test_mobilenet_config_variants() {
        let config_v1 = MobileNetConfig::mobilenet_v1();
        assert!(matches!(config_v1.variant, MobileNetVariant::MobileNetV1));

        let config_v2 = MobileNetConfig::mobilenet_v2();
        assert!(matches!(config_v2.variant, MobileNetVariant::MobileNetV2));

        let config_v3_small = MobileNetConfig::mobilenet_v3_small();
        assert!(matches!(config_v3_small.variant, MobileNetVariant::MobileNetV3Small));

        let config_v3_large = MobileNetConfig::mobilenet_v3_large();
        assert!(matches!(config_v3_large.variant, MobileNetVariant::MobileNetV3Large));
    }

    #[test]
    fn test_mobilenet_config_builders() {
        let config = MobileNetConfig::mobilenet_v2()
            .with_num_classes(10)
            .with_width_multiplier(0.5)
            .with_dropout_rate(0.3)
            .with_reduced_tail(true);

        assert_eq!(config.num_classes, 10);
        assert_eq!(config.width_multiplier, 0.5);
        assert_eq!(config.dropout_rate, 0.3);
        assert!(config.reduced_tail);
    }

    #[test]
    fn test_mobilenet_v1_creation() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v1();
        let model = MobileNet::new(config, &device);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 4_253_864);
    }

    #[test]
    fn test_mobilenet_v2_creation() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v2();
        let model = MobileNet::new(config, &device);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 3_504_872);
    }

    #[test]
    fn test_mobilenet_v3_small_creation() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v3_small();
        let model = MobileNet::new(config, &device);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 2_945_000);
    }

    #[test]
    fn test_mobilenet_v3_large_creation() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v3_large();
        let model = MobileNet::new(config, &device);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 5_483_000);
    }

    #[test]
    fn test_mobilenet_v2_forward_pass() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v2();
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();
        let output = model.forward(&input);

        assert!(output.is_ok());
        let output = output.unwrap();
        assert_eq!(output.shape().dims(), &[1, 1000]);
    }

    #[test]
    fn test_mobilenet_batch_forward() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v1();
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::randn(&[4, 3, 224, 224], &device).unwrap();
        let output = model.forward(&input);

        assert!(output.is_ok());
        let output = output.unwrap();
        assert_eq!(output.shape().dims(), &[4, 1000]);
    }

    #[test]
    fn test_mobilenet_custom_classes() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v2().with_num_classes(10);
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::randn(&[2, 3, 224, 224], &device).unwrap();
        let output = model.forward(&input);

        assert!(output.is_ok());
        let output = output.unwrap();
        assert_eq!(output.shape().dims(), &[2, 10]);
    }

    #[test]
    fn test_mobilenet_width_multipliers() {
        let device = get_test_device();
        let width_multipliers = vec![0.25, 0.5, 0.75, 1.0, 1.25];

        for width_mult in width_multipliers {
            let config = MobileNetConfig::mobilenet_v2()
                .with_width_multiplier(width_mult);
            let model = MobileNet::new(config, &device);

            assert!(model.is_ok(), "Failed with width multiplier {}", width_mult);

            let model = model.unwrap();
            let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();
            let output = model.forward(&input);

            assert!(output.is_ok());
            assert_eq!(output.unwrap().shape().dims(), &[1, 1000]);
        }
    }

    #[test]
    fn test_mobilenet_different_input_sizes() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v2();
        let model = MobileNet::new(config, &device).unwrap();

        let sizes = vec![
            (1, 3, 32, 32),   // CIFAR-10 size
            (1, 3, 128, 128), // Medium images
            (1, 3, 224, 224), // Standard ImageNet
        ];

        for (b, c, h, w) in sizes {
            let input = Tensor::randn(&[b, c, h, w], &device).unwrap();
            let output = model.forward(&input);

            assert!(output.is_ok(), "Failed for input size {}x{}x{}x{}", b, c, h, w);
            let output = output.unwrap();
            assert_eq!(output.shape().dims(), &[b, 1000]);
        }
    }

    #[test]
    fn test_mobilenet_feature_extraction() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v2();
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();
        let features = model.extract_features(&input);

        assert!(features.is_ok());
        let features = features.unwrap();
        let feature_dims = features.shape().dims();
        // Features should be spatial size 1x1 after global average pooling
        assert_eq!(feature_dims[2], 1);
        assert_eq!(feature_dims[3], 1);
    }

    #[test]
    fn test_mobilenet_edge_cases() {
        let device = get_test_device();

        // Test minimum viable input
        let config = MobileNetConfig::mobilenet_v2()
            .with_num_classes(1)
            .with_width_multiplier(0.25);
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::randn(&[1, 3, 64, 64], &device).unwrap();
        let output = model.forward(&input);

        assert!(output.is_ok());
        let output = output.unwrap();
        assert_eq!(output.shape().dims(), &[1, 1]);
    }

    #[test]
    fn test_mobilenet_deterministic_output() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v1();
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::ones(&[1, 3, 224, 224], &device).unwrap();
        let output1 = model.forward(&input).unwrap();
        let output2 = model.forward(&input).unwrap();

        // In mock implementation, we check they have the same shape
        assert_eq!(output1.shape().dims(), output2.shape().dims());
    }

    #[test]
    fn test_mobilenet_large_batch() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v2();
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::randn(&[16, 3, 224, 224], &device).unwrap();
        let output = model.forward(&input);

        assert!(output.is_ok());
        let output = output.unwrap();
        assert_eq!(output.shape().dims(), &[16, 1000]);
    }

    #[test]
    fn test_mobilenet_memory_consistency() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v3_small();
        let model = MobileNet::new(config, &device).unwrap();

        for i in 1..=5 {
            let input = Tensor::randn(&[2, 3, 224, 224], &device).unwrap();
            let output = model.forward(&input);
            assert!(output.is_ok(), "Forward pass {} failed", i);
        }
    }

    #[test]
    fn test_mobilenet_dropout_variations() {
        let device = get_test_device();
        let dropout_rates = vec![0.0, 0.1, 0.2, 0.5];

        for dropout in dropout_rates {
            let config = MobileNetConfig::mobilenet_v2()
                .with_dropout_rate(dropout);
            let model = MobileNet::new(config, &device);
            assert!(model.is_ok(), "Failed with dropout rate {}", dropout);
        }
    }

    #[test]
    fn test_mobilenet_parameter_counting() {
        let device = get_test_device();

        let models = vec![
            (MobileNetConfig::mobilenet_v1(), 4_253_864),
            (MobileNetConfig::mobilenet_v2(), 3_504_872),
            (MobileNetConfig::mobilenet_v3_small(), 2_945_000),
            (MobileNetConfig::mobilenet_v3_large(), 5_483_000),
        ];

        for (config, expected_params) in models {
            let model = MobileNet::new(config, &device).unwrap();
            assert_eq!(model.num_parameters(), expected_params);
        }
    }

    #[test]
    fn test_mobilenet_serialization() {
        let config = MobileNetConfig::mobilenet_v2()
            .with_num_classes(10)
            .with_width_multiplier(0.75);

        let serialized = serde_json::to_string(&config);
        assert!(serialized.is_ok());

        let deserialized: Result<MobileNetConfig, _> = serde_json::from_str(&serialized.unwrap());
        assert!(deserialized.is_ok());

        let deserialized = deserialized.unwrap();
        assert!(matches!(deserialized.variant, MobileNetVariant::MobileNetV2));
        assert_eq!(deserialized.num_classes, 10);
        assert_eq!(deserialized.width_multiplier, 0.75);
    }

    #[test]
    fn test_mobilenet_forward_consistency() {
        let device = get_test_device();
        let config = MobileNetConfig::mobilenet_v2();
        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::zeros(&[1, 3, 224, 224], &device).unwrap();

        let results: Vec<_> = (0..3)
            .map(|_| model.forward(&input))
            .collect();

        assert!(results.iter().all(|r| r.is_ok()));
    }
}

#[cfg(test)]
mod integration_tests {
    use super::*;
    use rtx_vision::architectures::{MobileNet, MobileNetConfig, MobileNetVariant};
    use rtx_vision::{Device, Tensor};

    #[test]
    fn test_mobilenet_transfer_learning_setup() {
        let device = Device::cpu();
        let config = MobileNetConfig::mobilenet_v2()
            .with_num_classes(10)  // Fine-tune for CIFAR-10
            .with_width_multiplier(0.5)
            .with_dropout_rate(0.3);

        let model = MobileNet::new(config, &device).unwrap();

        let input = Tensor::randn(&[4, 3, 32, 32], &device).unwrap();
        let output = model.forward(&input);

        assert!(output.is_ok());
        assert_eq!(output.unwrap().shape().dims(), &[4, 10]);
    }

    #[test]
    fn test_mobilenet_efficiency_comparison() {
        let device = Device::cpu();

        // Compare parameter counts across variants
        let variants = vec![
            MobileNetVariant::MobileNetV1,
            MobileNetVariant::MobileNetV2,
            MobileNetVariant::MobileNetV3Small,
            MobileNetVariant::MobileNetV3Large,
        ];

        let mut param_counts = Vec::new();

        for variant in variants {
            let config = MobileNetConfig {
                variant: variant.clone(),
                ..Default::default()
            };

            let model = MobileNet::new(config, &device).unwrap();
            param_counts.push(model.num_parameters());

            let input = Tensor::randn(&[2, 3, 224, 224], &device).unwrap();
            let output = model.forward(&input);

            assert!(output.is_ok(), "Failed for variant {:?}", variant);
            assert_eq!(output.unwrap().shape().dims(), &[2, 1000]);
        }

        // Verify V3-Small has fewer parameters than others
        assert!(param_counts[2] < param_counts[0]); // V3-Small < V1
        assert!(param_counts[2] < param_counts[1]); // V3-Small < V2
    }

    #[test]
    fn test_mobilenet_width_scaling() {
        let device = Device::cpu();

        let width_multipliers = vec![0.5, 1.0, 1.25];
        let mut outputs = Vec::new();

        for width_mult in width_multipliers {
            let config = MobileNetConfig::mobilenet_v2()
                .with_width_multiplier(width_mult);
            let model = MobileNet::new(config, &device).unwrap();

            let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();
            let features = model.extract_features(&input).unwrap();
            outputs.push(features);
        }

        // All should produce valid outputs
        assert_eq!(outputs.len(), 3);
    }

    #[test]
    fn test_mobilenet_depthwise_separable_convolutions() {
        let device = Device::cpu();
        let config = MobileNetConfig::mobilenet_v1(); // Uses depthwise separable convs

        let model = MobileNet::new(config, &device).unwrap();
        let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();

        let output = model.forward(&input).unwrap();
        assert_eq!(output.shape().dims(), &[1, 1000]);

        // Test feature extraction
        let features = model.extract_features(&input).unwrap();
        assert!(features.shape().dims()[1] > 0); // Should have channels
    }

    #[test]
    fn test_mobilenet_inverted_residuals() {
        let device = Device::cpu();
        let config = MobileNetConfig::mobilenet_v2(); // Uses inverted residuals

        let model = MobileNet::new(config, &device).unwrap();
        let input = Tensor::randn(&[2, 3, 224, 224], &device).unwrap();

        let output = model.forward(&input).unwrap();
        assert_eq!(output.shape().dims(), &[2, 1000]);
    }

    #[test]
    fn test_mobilenet_computational_efficiency() {
        use std::time::Instant;

        let device = Device::cpu();

        let configs = vec![
            MobileNetConfig::mobilenet_v1(),
            MobileNetConfig::mobilenet_v2(),
            MobileNetConfig::mobilenet_v3_small(),
        ];

        for config in configs {
            let model = MobileNet::new(config, &device).unwrap();
            let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();

            let start = Instant::now();
            let _output = model.forward(&input).unwrap();
            let _duration = start.elapsed();

            // In a real implementation, we would assert on performance characteristics
            // MobileNets should be faster than standard CNNs
        }
    }
}