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

#[cfg(test)]
mod tests {
    use rtx_vision::architectures::{VGG, VGGConfig, VGGVariant};
    use rtx_vision::{Device, Tensor};

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

    #[test]
    fn test_vgg_config_default() {
        let config = VGGConfig::default();
        assert!(matches!(config.variant, VGGVariant::VGG16));
        assert_eq!(config.num_classes, 1000);
        assert!(!config.batch_norm);
    }

    #[test]
    fn test_vgg_config_variants() {
        let config11 = VGGConfig::vgg11();
        assert!(matches!(config11.variant, VGGVariant::VGG11));

        let config13 = VGGConfig::vgg13();
        assert!(matches!(config13.variant, VGGVariant::VGG13));

        let config16 = VGGConfig::vgg16();
        assert!(matches!(config16.variant, VGGVariant::VGG16));

        let config19 = VGGConfig::vgg19();
        assert!(matches!(config19.variant, VGGVariant::VGG19));
    }

    #[test]
    fn test_vgg_config_builders() {
        let config = VGGConfig::vgg16()
            .with_num_classes(10)
            .with_dropout(0.3)
            .with_init_weights(false);

        assert_eq!(config.num_classes, 10);
        assert_eq!(config.dropout, 0.3);
        assert!(!config.init_weights);
    }

    #[test]
    fn test_vgg11_creation() {
        let device = get_test_device();
        let config = VGGConfig::vgg11();
        let model = VGG::new(config);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 132_863_336);
    }

    #[test]
    fn test_vgg13_creation() {
        let device = get_test_device();
        let config = VGGConfig::vgg13();
        let model = VGG::new(config);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 133_047_848);
    }

    #[test]
    fn test_vgg16_creation() {
        let device = get_test_device();
        let config = VGGConfig::vgg16();
        let model = VGG::new(config);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 138_357_544);
    }

    #[test]
    fn test_vgg19_creation() {
        let device = get_test_device();
        let config = VGGConfig::vgg19();
        let model = VGG::new(config);

        assert!(model.is_ok());
        let model = model.unwrap();
        assert_eq!(model.num_parameters(), 143_667_240);
    }

    #[test]
    fn test_vgg_with_batch_norm_creation() {
        let device = get_test_device();
        let variants = vec![
            VGGConfig::vgg11_bn(),
            VGGConfig::vgg13_bn(),
            VGGConfig::vgg16_bn(),
            VGGConfig::vgg19_bn(),
        ];

        for config in variants {
            let model = VGG::new(config);
            assert!(model.is_ok());
        }
    }

    #[test]
    fn test_vgg16_forward_pass() {
        let device = get_test_device();
        let config = VGGConfig::vgg16();
        let model = VGG::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_vgg_batch_forward() {
        let device = get_test_device();
        let config = VGGConfig::vgg11();
        let model = VGG::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_vgg_custom_classes() {
        let device = get_test_device();
        let config = VGGConfig::vgg11().with_num_classes(10);
        let model = VGG::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_vgg_different_input_sizes() {
        let device = get_test_device();
        let config = VGGConfig::vgg11();
        let model = VGG::new(config, &device).unwrap();

        let sizes = vec![
            (1, 3, 32, 32),   // CIFAR-10 size
            (1, 3, 64, 64),   // Small images
            (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_vgg_feature_extraction() {
        let device = get_test_device();
        let config = VGGConfig::vgg16();
        let model = VGG::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();
        // Features should be [batch_size, 512, 7, 7] for VGG
        assert_eq!(features.shape().dims(), &[1, 512, 7, 7]);
    }

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

        // Test minimum viable input
        let config = VGGConfig::vgg11().with_num_classes(1);
        let model = VGG::new(config, &device).unwrap();

        let input = Tensor::randn(&[1, 3, 32, 32], &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_vgg_deterministic_output() {
        let device = get_test_device();
        let config = VGGConfig::vgg11();
        let model = VGG::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_vgg_large_batch() {
        let device = get_test_device();
        let config = VGGConfig::vgg11();
        let model = VGG::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_vgg_memory_consistency() {
        let device = get_test_device();
        let config = VGGConfig::vgg16();
        let model = VGG::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_vgg_dropout_variations() {
        let device = get_test_device();
        let dropouts = vec![0.0, 0.3, 0.5, 0.7];

        for dropout in dropouts {
            let config = VGGConfig::vgg11().with_dropout(dropout);
            let model = VGG::new(config);
            assert!(model.is_ok());
        }
    }

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

        let models = vec![
            (VGGConfig::vgg11(), 132_863_336),
            (VGGConfig::vgg13(), 133_047_848),
            (VGGConfig::vgg16(), 138_357_544),
            (VGGConfig::vgg19(), 143_667_240),
        ];

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

    #[test]
    fn test_vgg_serialization() {
        let config = VGGConfig::vgg16().with_num_classes(10);
        let serialized = serde_json::to_string(&config);
        assert!(serialized.is_ok());

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

        let deserialized = deserialized.unwrap();
        assert!(matches!(deserialized.variant, VGGVariant::VGG16));
        assert_eq!(deserialized.num_classes, 10);
    }

    #[test]
    fn test_vgg_forward_consistency() {
        let device = get_test_device();
        let config = VGGConfig::vgg11();
        let model = VGG::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::{VGG, VGGConfig, VGGVariant};
    use rtx_vision::{Device, Tensor};

    #[test]
    fn test_vgg_transfer_learning_setup() {
        let device = Device::cpu();
        let config = VGGConfig::vgg16()
            .with_num_classes(10)  // Fine-tune for CIFAR-10
            .with_dropout(0.3);

        let model = VGG::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_vgg_feature_extraction_pipeline() {
        let device = Device::cpu();
        let config = VGGConfig::vgg19();
        let model = VGG::new(config, &device).unwrap();

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

        // Extract features for downstream tasks
        let features = model.extract_features(&input).unwrap();
        assert_eq!(features.shape().dims(), &[1, 512, 7, 7]);
    }

    #[test]
    fn test_all_vgg_variants_consistency() {
        let device = Device::cpu();
        let variants = vec![
            VGGVariant::VGG11,
            VGGVariant::VGG11WithBN,
            VGGVariant::VGG13,
            VGGVariant::VGG13WithBN,
            VGGVariant::VGG16,
            VGGVariant::VGG16WithBN,
            VGGVariant::VGG19,
            VGGVariant::VGG19WithBN,
        ];

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

            let model = VGG::new(config, &device).unwrap();
            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]);
        }
    }

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

        // Test VGG-16 without batch norm
        let config_no_bn = VGGConfig::vgg16();
        let model_no_bn = VGG::new(config_no_bn, &device).unwrap();

        // Test VGG-16 with batch norm
        let config_bn = VGGConfig::vgg16_bn();
        let model_bn = VGG::new(config_bn, &device).unwrap();

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

        let output_no_bn = model_no_bn.forward(&input).unwrap();
        let output_bn = model_bn.forward(&input).unwrap();

        // Both should produce outputs with same shape
        assert_eq!(output_no_bn.shape().dims(), output_bn.shape().dims());
    }

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

        let device = Device::cpu();

        let configs = vec![
            VGGConfig::vgg11(),
            VGGConfig::vgg16(),
            VGGConfig::vgg19(),
        ];

        for config in configs {
            let model = VGG::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
        }
    }
}