//! Medical imaging benchmarks
//!
//! Performance benchmarks for:
//! - DICOM processing and parsing
//! - Medical volume segmentation
//! - Multi-modal fusion (CT + MRI)
//! - 3D reconstruction and visualization
//! - Real-time medical imaging workflows

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

/// Medical imaging benchmark helper
struct MedicalBenchHelper;

impl MedicalBenchHelper {
    fn create_dicom_volume(slices: usize, height: usize, width: usize) -> Tensor {
        // Simulate realistic medical imaging data with appropriate intensity ranges
        let mut volume = Tensor::randn(&[slices, height, width], &Device::default())
            .expect("Failed to create DICOM volume");

        // Scale to typical CT Hounsfield units (-1000 to +3000)
        volume = volume
            .mul_scalar(500.0)
            .expect("Failed to scale volume")
            .add_scalar(-200.0)
            .expect("Failed to offset volume");

        volume
    }

    fn create_medical_metadata(modality: medical::ImagingModality) -> medical::MedicalMetadata {
        medical::MedicalMetadata {
            patient_id: "BENCH_001".to_string(),
            study_date: "20241201".to_string(),
            modality,
            series_description: "Benchmark Series".to_string(),
            slice_thickness: 1.25,
            pixel_spacing: [0.625, 0.625],
            window_center: if matches!(modality, medical::ImagingModality::CT) {
                40.0
            } else {
                300.0
            },
            window_width: if matches!(modality, medical::ImagingModality::CT) {
                400.0
            } else {
                600.0
            },
        }
    }
}

/// Benchmark DICOM file processing
fn bench_dicom_processing(c: &mut Criterion) {
    let mut group = c.benchmark_group("dicom_processing");

    let dicom_processor = medical::DicomProcessor::new().expect("Failed to create DICOM processor");

    // Different DICOM volume sizes (typical clinical scenarios)
    let volume_configs = vec![
        ("ct_chest", 100, 512, 512),   // Typical chest CT
        ("ct_abdomen", 150, 512, 512), // Abdominal CT
        ("mri_brain", 180, 256, 256),  // Brain MRI
        ("ct_spine", 200, 512, 512),   // Spine CT
    ];

    for (config_name, slices, height, width) in volume_configs {
        let volume_data = MedicalBenchHelper::create_dicom_volume(slices, height, width);
        let metadata = MedicalBenchHelper::create_medical_metadata(medical::ImagingModality::CT);

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

        group.bench_with_input(
            BenchmarkId::new("dicom_parse", config_name),
            &(volume_data, metadata),
            |b, (volume, meta)| {
                b.iter(|| {
                    // Simulate DICOM parsing and processing
                    let medical_volume = medical::MedicalVolume {
                        data: volume.clone(),
                        metadata: meta.clone(),
                        quality_metrics: medical::QualityMetrics {
                            snr: 25.0,
                            cnr: 15.0,
                            uniformity: 0.95,
                            noise_level: 0.05,
                        },
                    };

                    // Simulate quality assessment
                    let _ = medical_volume.quality_metrics.snr > 20.0;
                });
            },
        );
    }

    group.finish();
}

/// Benchmark medical image preprocessing
fn bench_medical_preprocessing(c: &mut Criterion) {
    let mut group = c.benchmark_group("medical_preprocessing");

    // Different preprocessing operations
    let preprocessing_ops = vec![
        ("windowing", "window_level_adjustment"),
        ("normalization", "intensity_normalization"),
        ("denoising", "gaussian_denoising"),
        ("registration", "volume_registration"),
        ("resampling", "isotropic_resampling"),
    ];

    let volume_data = MedicalBenchHelper::create_dicom_volume(64, 256, 256);

    for (op_name, operation) in preprocessing_ops {
        group.bench_with_input(
            BenchmarkId::new("preprocessing", op_name),
            &volume_data,
            |b, volume| {
                b.iter(|| {
                    match operation {
                        "window_level_adjustment" => {
                            // Window/Level adjustment (common in CT/MRI)
                            let window_center = 40.0;
                            let window_width = 400.0;
                            let windowed = volume
                                .clamp(
                                    window_center - window_width / 2.0,
                                    window_center + window_width / 2.0,
                                )
                                .expect("Windowing failed");
                            windowed
                        }
                        "intensity_normalization" => {
                            // Z-score normalization
                            let mean = volume
                                .mean(None)
                                .expect("Mean calculation failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            let std = volume
                                .std(None, false)
                                .expect("Std calculation failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            volume
                                .sub_scalar(mean)
                                .expect("Subtraction failed")
                                .div_scalar(std + 1e-8)
                                .expect("Division failed")
                        }
                        "gaussian_denoising" => {
                            // Simplified Gaussian smoothing
                            // In practice, this would use proper 3D Gaussian kernels
                            let smoothed = volume.clone(); // Placeholder
                            smoothed
                        }
                        "volume_registration" => {
                            // Volume registration preprocessing
                            // This would involve affine transformations, mutual information, etc.
                            let registered = volume.clone(); // Placeholder
                            registered
                        }
                        "isotropic_resampling" => {
                            // Resample to isotropic voxel spacing
                            // This would involve interpolation to make all voxel dimensions equal
                            let resampled = volume.clone(); // Placeholder
                            resampled
                        }
                        _ => volume.clone(),
                    }
                });
            },
        );
    }

    group.finish();
}

/// Benchmark 3D medical image segmentation
fn bench_3d_medical_segmentation(c: &mut Criterion) {
    let mut group = c.benchmark_group("3d_medical_segmentation");

    // Different anatomical structures with typical volume sizes
    let segmentation_tasks = vec![
        ("liver_ct", 80, 512, 512),   // Liver segmentation
        ("brain_mri", 160, 256, 256), // Brain tissue segmentation
        ("lung_ct", 100, 512, 512),   // Lung segmentation
        ("cardiac_ct", 60, 512, 512), // Cardiac segmentation
    ];

    for (task_name, depth, height, width) in segmentation_tasks {
        let volume = MedicalBenchHelper::create_dicom_volume(depth, height, width);
        let total_voxels = (depth * height * width) as u64;

        group.throughput(Throughput::Elements(total_voxels));
        group.bench_with_input(
            BenchmarkId::new("3d_segmentation", task_name),
            &volume,
            |b, vol| {
                b.iter(|| {
                    // Simulate 3D medical segmentation
                    // This would typically involve:
                    // 1. 3D U-Net or similar architecture
                    // 2. Patch-based processing for large volumes
                    // 3. Post-processing (connected components, morphology)

                    // For benchmark, simulate processing each slice
                    let mut segmented_slices = Vec::new();
                    for slice_idx in 0..vol.shape()[0] {
                        let slice = vol
                            .narrow(0, slice_idx, 1)
                            .expect("Failed to extract slice");

                        // Simulate 2D segmentation on slice
                        let segmented = slice
                            .gt_scalar(0.0)
                            .expect("Thresholding failed")
                            .to_dtype(DType::U8)
                            .expect("Type conversion failed");

                        segmented_slices.push(segmented);
                    }

                    // Simulate 3D reconstruction
                    let _ = segmented_slices.len();
                });
            },
        );
    }

    group.finish();
}

/// Benchmark multi-modal medical image fusion
fn bench_multimodal_fusion(c: &mut Criterion) {
    let mut group = c.benchmark_group("multimodal_fusion");

    // Common multi-modal combinations in medical imaging
    let fusion_scenarios = vec![
        (
            "ct_pet",
            medical::ImagingModality::CT,
            medical::ImagingModality::PET,
        ),
        (
            "mri_t1_t2",
            medical::ImagingModality::MRI,
            medical::ImagingModality::MRI,
        ),
        (
            "ct_mri",
            medical::ImagingModality::CT,
            medical::ImagingModality::MRI,
        ),
    ];

    for (scenario_name, modality1, modality2) in fusion_scenarios {
        let volume1 = MedicalBenchHelper::create_dicom_volume(64, 256, 256);
        let volume2 = MedicalBenchHelper::create_dicom_volume(64, 256, 256);

        let metadata1 = MedicalBenchHelper::create_medical_metadata(modality1);
        let metadata2 = MedicalBenchHelper::create_medical_metadata(modality2);

        group.bench_with_input(
            BenchmarkId::new("multimodal_fusion", scenario_name),
            &(volume1, volume2, metadata1, metadata2),
            |b, (vol1, vol2, meta1, meta2)| {
                b.iter(|| {
                    // Simulate multi-modal fusion
                    // 1. Registration (align modalities)
                    // 2. Intensity normalization
                    // 3. Feature extraction
                    // 4. Fusion strategy (pixel-level, feature-level, or decision-level)

                    // Registration step (simplified)
                    let registered_vol2 = vol2.clone(); // Would involve actual registration

                    // Intensity normalization
                    let norm_vol1 = vol1
                        .sub(vol1.mean(None).expect("Mean failed"))
                        .expect("Normalization failed");
                    let norm_vol2 = registered_vol2
                        .sub(registered_vol2.mean(None).expect("Mean failed"))
                        .expect("Normalization failed");

                    // Fusion (simple averaging for benchmark)
                    let fused = norm_vol1
                        .add(&norm_vol2)
                        .expect("Addition failed")
                        .div_scalar(2.0)
                        .expect("Division failed");

                    let _ = fused;
                });
            },
        );
    }

    group.finish();
}

/// Benchmark medical image quality assessment
fn bench_quality_assessment(c: &mut Criterion) {
    let mut group = c.benchmark_group("quality_assessment");

    // Different quality metrics commonly used in medical imaging
    let quality_metrics = vec![
        ("snr_calculation", "signal_to_noise_ratio"),
        ("cnr_calculation", "contrast_to_noise_ratio"),
        ("uniformity_assessment", "intensity_uniformity"),
        ("artifact_detection", "motion_artifacts"),
        ("sharpness_measure", "edge_sharpness"),
    ];

    let volume = MedicalBenchHelper::create_dicom_volume(32, 256, 256);

    for (metric_name, metric_type) in quality_metrics {
        group.bench_with_input(
            BenchmarkId::new("quality_metrics", metric_name),
            &volume,
            |b, vol| {
                b.iter(|| {
                    match metric_type {
                        "signal_to_noise_ratio" => {
                            // SNR calculation: mean / std
                            let mean = vol
                                .mean(None)
                                .expect("Mean calculation failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            let std = vol
                                .std(None, false)
                                .expect("Std calculation failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            let snr = mean / (std + 1e-8);
                            snr
                        }
                        "contrast_to_noise_ratio" => {
                            // CNR: (signal - background) / noise_std
                            let signal_roi = vol.narrow(0, 0, 16).expect("ROI extraction failed");
                            let background_roi =
                                vol.narrow(0, 16, 16).expect("Background extraction failed");

                            let signal_mean = signal_roi
                                .mean(None)
                                .expect("Signal mean failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            let bg_mean = background_roi
                                .mean(None)
                                .expect("Background mean failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            let noise_std = background_roi
                                .std(None, false)
                                .expect("Noise std failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");

                            let cnr = (signal_mean - bg_mean) / (noise_std + 1e-8);
                            cnr
                        }
                        "intensity_uniformity" => {
                            // Uniformity: 1 - (std / mean)
                            let mean = vol
                                .mean(None)
                                .expect("Mean calculation failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            let std = vol
                                .std(None, false)
                                .expect("Std calculation failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            let uniformity = 1.0 - (std / (mean.abs() + 1e-8));
                            uniformity.max(0.0).min(1.0)
                        }
                        "motion_artifacts" => {
                            // Motion artifact detection (simplified)
                            // Would typically analyze frequency domain or gradient variations
                            let grad = vol.diff(-1, None).expect("Gradient calculation failed");
                            let artifact_score = grad
                                .abs()
                                .mean(None)
                                .expect("Artifact score failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            artifact_score
                        }
                        "edge_sharpness" => {
                            // Edge sharpness measurement
                            let grad = vol.diff(-1, None).expect("Gradient calculation failed");
                            let sharpness = grad
                                .abs()
                                .max(None)
                                .expect("Max gradient failed")
                                .to_scalar::<f32>()
                                .expect("Scalar conversion failed");
                            sharpness
                        }
                        _ => 0.0,
                    }
                });
            },
        );
    }

    group.finish();
}

/// Benchmark real-time medical imaging workflows
fn bench_realtime_medical_workflow(c: &mut Criterion) {
    let mut group = c.benchmark_group("realtime_medical_workflow");

    // Simulate real-time medical imaging scenarios
    let workflow_scenarios = vec![
        ("ultrasound_realtime", 30, 480, 640),    // 30 FPS ultrasound
        ("fluoroscopy_realtime", 15, 1024, 1024), // 15 FPS fluoroscopy
        ("mri_slice_realtime", 5, 256, 256),      // 5 slices/sec MRI acquisition
    ];

    for (scenario_name, fps, height, width) in workflow_scenarios {
        let frame = MedicalBenchHelper::create_dicom_volume(1, height, width);
        let frames_per_bench = fps; // Simulate 1 second worth of frames

        group.throughput(Throughput::Elements(frames_per_bench as u64));
        group.bench_with_input(
            BenchmarkId::new("realtime_workflow", scenario_name),
            &frame,
            |b, f| {
                b.iter(|| {
                    // Simulate real-time processing pipeline
                    for _ in 0..frames_per_bench {
                        // 1. Frame acquisition (simulated)
                        let current_frame = f.clone();

                        // 2. Real-time preprocessing
                        let processed = current_frame
                            .clamp(-1000.0, 3000.0)
                            .expect("Clamping failed") // Window
                            .add_scalar(1000.0)
                            .expect("Offset failed") // Normalize
                            .div_scalar(4000.0)
                            .expect("Scale failed"); // Scale to 0-1

                        // 3. Real-time analysis (if needed)
                        let analysis = processed
                            .mean(None)
                            .expect("Analysis failed")
                            .to_scalar::<f32>()
                            .expect("Scalar conversion failed");

                        // 4. Display preparation
                        let _ = analysis > 0.5; // Simple threshold for display
                    }
                });
            },
        );
    }

    group.finish();
}

/// Benchmark memory usage in medical imaging
fn bench_medical_memory_usage(c: &mut Criterion) {
    let mut group = c.benchmark_group("medical_memory_usage");

    // Test memory usage with different volume sizes
    let volume_sizes = vec![
        ("small_volume", 32, 128, 128),
        ("medium_volume", 64, 256, 256),
        ("large_volume", 128, 512, 512),
        ("huge_volume", 256, 512, 512),
    ];

    for (size_name, depth, height, width) in volume_sizes {
        let volume = MedicalBenchHelper::create_dicom_volume(depth, height, width);
        let total_voxels = (depth * height * width) as u64;

        group.throughput(Throughput::Bytes(total_voxels * 4)); // 4 bytes per float32
        group.bench_with_input(
            BenchmarkId::new("memory_usage", size_name),
            &volume,
            |b, vol| {
                b.iter(|| {
                    // Simulate memory-intensive operations

                    // 1. Volume duplication (common in processing)
                    let vol_copy = vol.clone();

                    // 2. Multi-scale processing
                    let downsampled = vol
                        .narrow(0, 0, vol.shape()[0] / 2)
                        .expect("Downsampling failed");

                    // 3. Gradient calculation (memory intensive)
                    let grad_x = vol.diff(-1, None).expect("Gradient X failed");
                    let grad_y = vol.diff(-2, None).expect("Gradient Y failed");

                    // 4. Temporary buffers for processing
                    let temp_buffer = Tensor::zeros_like(&vol).expect("Buffer creation failed");

                    // Clean up references to measure memory usage
                    drop(vol_copy);
                    drop(downsampled);
                    drop(grad_x);
                    drop(grad_y);
                    drop(temp_buffer);
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_dicom_processing,
    bench_medical_preprocessing,
    bench_3d_medical_segmentation,
    bench_multimodal_fusion,
    bench_quality_assessment,
    bench_realtime_medical_workflow,
    bench_medical_memory_usage
);

criterion_main!(benches);