//! Segmentation model benchmarks
//!
//! Performance benchmarks for:
//! - Semantic segmentation (DeepLabV3+)
//! - Instance segmentation (Mask R-CNN)
//! - Panoptic segmentation
//! - Medical image segmentation
//! - Real-time segmentation optimization

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rtx_tensor::{DType, Device, Tensor};
use rtx_vision_advanced::*;

/// Benchmark helper for segmentation tests
struct SegmentationBenchHelper;

impl SegmentationBenchHelper {
    fn create_image(batch_size: usize, height: usize, width: usize) -> Tensor {
        Tensor::randn(&[batch_size, 3, height, width], &Device::default())
            .expect("Failed to create benchmark image")
    }

    fn create_medical_volume(depth: usize, height: usize, width: usize) -> Tensor {
        Tensor::randn(&[1, depth, height, width], &Device::default())
            .expect("Failed to create benchmark medical volume")
    }
}

/// Benchmark semantic segmentation models
fn bench_semantic_segmentation(c: &mut Criterion) {
    let mut group = c.benchmark_group("semantic_segmentation");

    // DeepLabV3+ with different backbones
    let backbones = vec![("resnet50", "resnet50"), ("resnet101", "resnet101")];

    for (name, backbone) in backbones {
        let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, backbone)
            .expect("Failed to create DeepLabV3+");

        let test_image = SegmentationBenchHelper::create_image(1, 512, 512);
        let config = segmentation::SegmentationConfig {
            segmentation_type: segmentation::SegmentationType::Semantic,
            ..Default::default()
        };

        group.throughput(Throughput::Elements(512 * 512)); // Pixels
        group.bench_with_input(
            BenchmarkId::new("deeplabv3plus", name),
            &(test_image, config),
            |b, (image, cfg)| {
                b.iter(|| segmentor.segment(image, cfg).expect("Segmentation failed"));
            },
        );
    }

    group.finish();
}

/// Benchmark instance segmentation
fn bench_instance_segmentation(c: &mut Criterion) {
    let mut group = c.benchmark_group("instance_segmentation");

    let mut mask_rcnn = segmentation::SegmentationFactory::create_mask_rcnn(80)
        .expect("Failed to create Mask R-CNN");

    let input_sizes = vec![
        ("small", 512, 512),
        ("medium", 800, 1333),
        ("large", 1024, 1024),
    ];

    for (size_name, height, width) in input_sizes {
        let test_image = SegmentationBenchHelper::create_image(1, height, width);
        let config = segmentation::SegmentationConfig {
            segmentation_type: segmentation::SegmentationType::Instance,
            vision_config: VisionConfig {
                input_size: (height, width),
                ..Default::default()
            },
            ..Default::default()
        };

        group.throughput(Throughput::Elements((height * width) as u64));
        group.bench_with_input(
            BenchmarkId::new("mask_rcnn", size_name),
            &(test_image, config),
            |b, (image, cfg)| {
                b.iter(|| {
                    mask_rcnn
                        .segment(image, cfg)
                        .expect("Instance segmentation failed")
                });
            },
        );
    }

    group.finish();
}

/// Benchmark multi-scale segmentation
fn bench_multiscale_segmentation(c: &mut Criterion) {
    let mut group = c.benchmark_group("multiscale_segmentation");

    let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
        .expect("Failed to create segmentor");

    let scale_configs = vec![
        ("single_scale", vec![1.0]),
        ("multi_scale_3", vec![0.75, 1.0, 1.25]),
        ("multi_scale_5", vec![0.5, 0.75, 1.0, 1.25, 1.5]),
    ];

    let test_image = SegmentationBenchHelper::create_image(1, 512, 512);

    for (config_name, scales) in scale_configs {
        let config = segmentation::SegmentationConfig {
            multi_scale: scales.len() > 1,
            scales,
            ..Default::default()
        };

        group.bench_with_input(
            BenchmarkId::new("multiscale", config_name),
            &(test_image.clone(), config),
            |b, (image, cfg)| {
                b.iter(|| {
                    segmentor
                        .segment(image, cfg)
                        .expect("Multi-scale segmentation failed")
                });
            },
        );
    }

    group.finish();
}

/// Benchmark real-time segmentation
fn bench_realtime_segmentation(c: &mut Criterion) {
    let mut group = c.benchmark_group("realtime_segmentation");

    // Test different input resolutions for real-time performance
    let resolutions = vec![
        ("240p", 240, 320),
        ("360p", 360, 480),
        ("480p", 480, 640),
        ("720p", 720, 1280),
    ];

    let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
        .expect("Failed to create segmentor");

    for (res_name, height, width) in resolutions {
        let test_image = SegmentationBenchHelper::create_image(1, height, width);
        let config = segmentation::SegmentationConfig {
            vision_config: VisionConfig {
                input_size: (height, width),
                ..Default::default()
            },
            output_stride: 16, // Larger stride for speed
            ..Default::default()
        };

        group.throughput(Throughput::Elements((height * width) as u64));
        group.bench_with_input(
            BenchmarkId::new("realtime", res_name),
            &(test_image, config),
            |b, (image, cfg)| {
                b.iter(|| {
                    segmentor
                        .segment(image, cfg)
                        .expect("Real-time segmentation failed")
                });
            },
        );
    }

    group.finish();
}

/// Benchmark medical image segmentation
fn bench_medical_segmentation(c: &mut Criterion) {
    let mut group = c.benchmark_group("medical_segmentation");

    // Different medical volume sizes
    let volume_sizes = vec![
        ("ct_small", 32, 256, 256),
        ("ct_medium", 64, 512, 512),
        ("mri_large", 128, 256, 256),
    ];

    for (volume_name, depth, height, width) in volume_sizes {
        let medical_volume = SegmentationBenchHelper::create_medical_volume(depth, height, width);

        // Medical segmentation typically processes slice by slice
        let voxels = (depth * height * width) as u64;
        group.throughput(Throughput::Elements(voxels));

        group.bench_with_input(
            BenchmarkId::new("medical_3d", volume_name),
            &medical_volume,
            |b, volume| {
                b.iter(|| {
                    // Simulate medical volume segmentation
                    // In practice, this would involve:
                    // 1. Slice extraction
                    // 2. 2D segmentation on each slice
                    // 3. 3D reconstruction
                    // 4. Post-processing (morphological operations)

                    for slice_idx in 0..volume.shape()[1] {
                        let slice = volume
                            .narrow(1, slice_idx, 1)
                            .expect("Failed to extract slice");
                        let slice_2d = slice.squeeze_dim(1).expect("Failed to squeeze slice");
                        let slice_rgb = slice_2d
                            .unsqueeze_dim(1)
                            .expect("Failed to add channel")
                            .repeat(&[1, 3, 1, 1])
                            .expect("Failed to repeat channels");

                        // This is a placeholder for actual medical segmentation
                        let _ = slice_rgb;
                    }
                });
            },
        );
    }

    group.finish();
}

/// Benchmark batch segmentation processing
fn bench_batch_segmentation(c: &mut Criterion) {
    let mut group = c.benchmark_group("batch_segmentation");

    let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
        .expect("Failed to create segmentor");

    let batch_sizes = vec![1, 2, 4, 8];

    for batch_size in batch_sizes {
        let batch_images = SegmentationBenchHelper::create_image(batch_size, 512, 512);
        let config = segmentation::SegmentationConfig::default();

        group.throughput(Throughput::Elements(batch_size as u64));
        group.bench_with_input(
            BenchmarkId::new("batch_process", batch_size),
            &(batch_images, config),
            |b, (images, cfg)| {
                b.iter(|| {
                    segmentor
                        .segment_batch(images, cfg)
                        .expect("Batch segmentation failed")
                });
            },
        );
    }

    group.finish();
}

/// Benchmark post-processing operations
fn bench_postprocessing(c: &mut Criterion) {
    let mut group = c.benchmark_group("segmentation_postprocessing");

    let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
        .expect("Failed to create segmentor");

    let test_image = SegmentationBenchHelper::create_image(1, 512, 512);

    // Different post-processing configurations
    let postprocess_configs = vec![
        ("no_postprocess", false, false, 0),
        ("morphology_only", true, false, 0),
        ("crf_only", false, true, 10),
        ("full_postprocess", true, true, 10),
    ];

    for (config_name, use_morphology, use_crf, crf_iterations) in postprocess_configs {
        let config = segmentation::SegmentationConfig {
            post_process: segmentation::PostProcessConfig {
                use_morphology,
                use_crf,
                crf_iterations,
                ..Default::default()
            },
            ..Default::default()
        };

        group.bench_with_input(
            BenchmarkId::new("postprocess", config_name),
            &(test_image.clone(), config),
            |b, (image, cfg)| {
                b.iter(|| {
                    segmentor
                        .segment(image, cfg)
                        .expect("Post-processing benchmark failed")
                });
            },
        );
    }

    group.finish();
}

/// Benchmark segmentation accuracy vs speed trade-offs
fn bench_accuracy_speed_tradeoff(c: &mut Criterion) {
    let mut group = c.benchmark_group("segmentation_tradeoff");

    let test_image = SegmentationBenchHelper::create_image(1, 512, 512);

    // Different configurations representing accuracy vs speed trade-offs
    let tradeoff_configs = vec![
        ("speed_optimized", 16, false, vec![1.0]), // Large stride, no multi-scale
        ("balanced", 8, true, vec![0.75, 1.0, 1.25]), // Medium stride, some multi-scale
        (
            "accuracy_optimized",
            4,
            true,
            vec![0.5, 0.75, 1.0, 1.25, 1.5],
        ), // Small stride, full multi-scale
    ];

    for (config_name, output_stride, multi_scale, scales) in tradeoff_configs {
        let mut segmentor =
            segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet101")
                .expect("Failed to create segmentor");

        let config = segmentation::SegmentationConfig {
            output_stride,
            multi_scale,
            scales,
            ..Default::default()
        };

        group.bench_with_input(
            BenchmarkId::new("tradeoff", config_name),
            &(test_image.clone(), config),
            |b, (image, cfg)| {
                b.iter(|| {
                    segmentor
                        .segment(image, cfg)
                        .expect("Tradeoff benchmark failed")
                });
            },
        );
    }

    group.finish();
}

/// Benchmark memory usage during segmentation
fn bench_memory_usage(c: &mut Criterion) {
    let mut group = c.benchmark_group("segmentation_memory");

    // Test memory usage with different image sizes and batch sizes
    let test_configs = vec![
        ("small_single", 1, 256, 256),
        ("medium_single", 1, 512, 512),
        ("large_single", 1, 1024, 1024),
        ("small_batch", 4, 256, 256),
        ("medium_batch", 4, 512, 512),
    ];

    for (config_name, batch_size, height, width) in test_configs {
        let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
            .expect("Failed to create segmentor");

        let test_images = SegmentationBenchHelper::create_image(batch_size, height, width);
        let config = segmentation::SegmentationConfig {
            vision_config: VisionConfig {
                input_size: (height, width),
                ..Default::default()
            },
            ..Default::default()
        };

        let total_pixels = (batch_size * height * width) as u64;
        group.throughput(Throughput::Elements(total_pixels));

        group.bench_with_input(
            BenchmarkId::new("memory", config_name),
            &(test_images, config),
            |b, (images, cfg)| {
                b.iter(|| {
                    if cfg.vision_config.input_size.0 > 512 {
                        // For large images, test single image processing
                        let single_image =
                            images.narrow(0, 0, 1).expect("Failed to get single image");
                        segmentor
                            .segment(&single_image, cfg)
                            .expect("Memory benchmark failed")
                    } else {
                        // For smaller images, test batch processing
                        segmentor
                            .segment_batch(images, cfg)
                            .expect("Memory benchmark failed")
                    }
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_semantic_segmentation,
    bench_instance_segmentation,
    bench_multiscale_segmentation,
    bench_realtime_segmentation,
    bench_medical_segmentation,
    bench_batch_segmentation,
    bench_postprocessing,
    bench_accuracy_speed_tradeoff,
    bench_memory_usage
);

criterion_main!(benches);