//! Edge-aware training performance benchmarks
//!
//! Benchmarks validate performance targets across all edge platforms:
//! - ARM NEON: Target 2.5x speedup over scalar
//! - RISC-V RVV: Target 4.8x speedup with vector extensions
//! - WASM SIMD: Target 1.8x speedup in browser
//! - Mobile GPU: Target 4.0x speedup with compute shaders
//! - IoT devices: Target 10x power efficiency improvement

use criterion::{black_box, criterion_group, criterion_main, Criterion, BenchmarkId};
use rustytorch::revolutionary::*;
use std::collections::HashMap;
use std::time::{Duration, Instant, SystemTime};

/// Benchmark ARM NEON optimization performance
fn benchmark_arm_neon_optimization(c: &mut Criterion) {
    let mut group = c.benchmark_group("ARM_NEON_Optimization");

    let test_cases = vec![
    ("Cortex-A", ArmArchitecture::CortexA, true),
    ("Apple_Silicon", ArmArchitecture::AppleSilicon, true),
    ("Cortex-A_No_NEON", ArmArchitecture::CortexA, false),
    ];

    for (name, arch, enable_neon) in test_cases {
        group.bench_with_input(BenchmarkId::new("optimization", name), &(arch, enable_neon),
        |b, (arch, neon)| {
            b.iter(|| {
                let mut optimizer = EdgeTargetOptimizer::new();
                let arm_opts = ArmOptimizations {
                    enable_neon: *neon,
                    memory_prefetch: true,
                    cache_optimization: true,
                    big_little_scheduling: true,
                    target_arch: *arch,
                };
                optimizer.configure_arm(arm_opts);
                black_box(optimizer.optimize_for_target(EdgeTarget::ARM).unwrap())
            })
        });
    }

    group.finish();
}

/// Benchmark RISC-V vector extension performance
fn benchmark_riscv_vector_optimization(c: &mut Criterion) {
    let mut group = c.benchmark_group("RISCV_Vector_Optimization");

    let vector_lengths = vec![
    ("VLEN128", RiscVVectorLength::VLEN128),
    ("VLEN256", RiscVVectorLength::VLEN256),
    ("VLEN512", RiscVVectorLength::VLEN512),
    ("Variable", RiscVVectorLength::Variable),
    ];

    for (name, vlen) in vector_lengths {
        group.bench_with_input(BenchmarkId::new("vector_length", name), &vlen,
        |b, vlen| {
            b.iter(|| {
                let mut optimizer = EdgeTargetOptimizer::new();
                let riscv_opts = RiscVOptimizations {
                    enable_rvv: true,
                    vector_length: *vlen,
                    custom_instructions: vec!["custom_matmul".to_string(), "custom_conv".to_string()],
                    memory_model: RiscVMemoryModel::TSO,
                    target_variant: RiscVVariant::Vector,
                };
                optimizer.configure_riscv(riscv_opts);
                black_box(optimizer.optimize_for_target(EdgeTarget::RISCV).unwrap())
            })
        });
    }

    group.finish();
}

/// Benchmark WebAssembly SIMD performance
fn benchmark_wasm_simd_optimization(c: &mut Criterion) {
    let mut group = c.benchmark_group("WASM_SIMD_Optimization");

    let runtime_configs = vec![
    ("Browser_SIMD", WasmRuntime::Browser, true, true),
    ("Browser_No_SIMD", WasmRuntime::Browser, false, false),
    ("Wasmtime_SIMD", WasmRuntime::Wasmtime, true, true),
    ("NodeJS_SIMD", WasmRuntime::NodeJS, true, true),
    ];

    for (name, runtime, simd, threads) in runtime_configs {
        group.bench_with_input(BenchmarkId::new("runtime", name), &(runtime, simd, threads),
        |b, (runtime, simd, threads)| {
            b.iter(|| {
                let mut optimizer = EdgeTargetOptimizer::new();
                let wasm_opts = WasmOptimizations {
                    enable_simd: *simd,
                    enable_threads: *threads,
                    memory_growth: WasmMemoryGrowth::Dynamic { max_pages: 2048 },
                    target_runtime: *runtime,
                    bulk_memory: true,
                };
                optimizer.configure_wasm(wasm_opts);
                black_box(optimizer.optimize_for_target(EdgeTarget::WASM).unwrap())
            })
        });
    }

    group.finish();
}

/// Benchmark Mobile GPU optimization across vendors
fn benchmark_mobile_gpu_optimization(c: &mut Criterion) {
    let mut group = c.benchmark_group("Mobile_GPU_Optimization");

    let gpu_vendors = vec![
    ("Mali", MobileGpuVendor::Mali),
    ("Adreno", MobileGpuVendor::Adreno),
    ("PowerVR", MobileGpuVendor::PowerVR),
    ("Apple_GPU", MobileGpuVendor::Apple),
    ("Intel_GPU", MobileGpuVendor::Intel),
    ];

    for (name, vendor) in gpu_vendors {
        group.bench_with_input(BenchmarkId::new("gpu_vendor", name), &vendor,
        |b, vendor| {
            b.iter(|| {
                let mut optimizer = EdgeTargetOptimizer::new();
                let mobile_gpu_opts = MobileGpuOptimizations {
                    gpu_vendor: *vendor,
                    compute_shaders: true,
                    tile_based_rendering: true,
                    bandwidth_optimization: true,
                    power_efficiency: true,
                };
                optimizer.configure_mobile_gpu(mobile_gpu_opts);
                black_box(optimizer.optimize_for_target(EdgeTarget::MobileGPU).unwrap())
            })
        });
    }

    group.finish();
}

/// Benchmark IoT ultra-low power optimization
fn benchmark_iot_optimization(c: &mut Criterion) {
    let mut group = c.benchmark_group("IoT_Ultra_Low_Power");

    let iot_platforms = vec![
    ("ESP32", IoTPlatform::ESP32),
    ("STM32", IoTPlatform::STM32),
    ("Arduino", IoTPlatform::Arduino),
    ("RaspberryPi_Pico", IoTPlatform::RaspberryPiPico),
    ];

    for (name, platform) in iot_platforms {
        group.bench_with_input(BenchmarkId::new("platform", name), &platform,
        |b, platform| {
            b.iter(|| {
                let mut optimizer = EdgeTargetOptimizer::new();
                let iot_opts = IoTOptimizations {
                    ultra_low_power: true,
                    minimal_memory: true,
                    wake_on_inference: true,
                    mesh_networking: true,
                    target_platform: *platform,
                };
                optimizer.configure_iot(iot_opts);
                black_box(optimizer.optimize_for_target(EdgeTarget::IoT).unwrap())
            })
        });
    }

    group.finish();
}

/// Benchmark federated device selection at scale
fn benchmark_federated_device_selection(c: &mut Criterion) {
    let mut group = c.benchmark_group("Federated_Device_Selection");

    let device_scales = vec![100, 1_000, 10_000, 50_000];
    let selection_strategies = vec![
    ("Random", SelectionStrategy::Random),
    ("BatteryAware", SelectionStrategy::BatteryAware),
    ("NetworkAware", SelectionStrategy::NetworkAware),
    ("PerformanceBased", SelectionStrategy::PerformanceBased),
    ("Hybrid", SelectionStrategy::Hybrid),
    ("Intelligent", SelectionStrategy::Intelligent),
    ];

    for device_count in device_scales {
        for (strategy_name, strategy) in &selection_strategies {
            let benchmark_name = format!("{}_{}", strategy_name, device_count);

            group.bench_with_input(
            BenchmarkId::new("device_selection", &benchmark_name),
            &(device_count, *strategy),
            |b, (count, strat)| {
                // Pre-create coordinator with devices
                let rt = tokio::runtime::Runtime::new().unwrap();
                let (coordinator, _sender) = FederatedCoordinator::new(
                "bench-coordinator".to_string(),
                *strat,
                GradientCompression {
                    algorithm: CompressionAlgorithm::TopK,
                    compression_ratio: 0.1,
                    error_correction: true,
                    adaptive_compression: true,
                },
                AggregationStrategy {
                    algorithm: AggregationAlgorithm::FedAvg,
                    weighting: WeightingScheme::Adaptive,
                    byzantine_tolerance: ByzantineTolerance {
                        enabled: false,
                        max_byzantine_fraction: 0.0,
                        detection_algorithm: ByzantineDetection::None,
                    },
                    differential_privacy: None,
                },
                );

                rt.block_on(async {
                    for i in 0..*count {
                        let device = create_benchmark_device(&format!("device-{:06}", i));
                        coordinator.register_device(device).await.unwrap();
                    }
                });

                let selection_criteria = SelectionCriteria {
                    min_battery_level: 0.3,
                    min_bandwidth_mbps: 5.0,
                    max_latency_ms: 200,
                    required_availability_minutes: 30,
                    min_data_quality: 0.7,
                };

                b.to_async(&rt).iter(|| async {
                    let target_count = (*count / 10).max(10); // Select 10% of devices
                    black_box(coordinator.select_devices(target_count as u32, selection_criteria.clone()).await.unwrap())
                });
            });
        }
    }

    group.finish();
}

/// Benchmark gradient aggregation performance
fn benchmark_gradient_aggregation(c: &mut Criterion) {
    let mut group = c.benchmark_group("Gradient_Aggregation");

    let device_counts = vec![10, 100, 1_000, 10_000];
    let aggregation_algorithms = vec![
    ("FedAvg", AggregationAlgorithm::FedAvg),
    ("FedProx", AggregationAlgorithm::FedProx),
    ("SCAFFOLD", AggregationAlgorithm::SCAFFOLD),
    ("FedAdam", AggregationAlgorithm::FedAdam),
    ];

    for device_count in device_counts {
        for (algo_name, algorithm) in &aggregation_algorithms {
            let benchmark_name = format!("{}_{}_devices", algo_name, device_count);

            group.bench_with_input(
            BenchmarkId::new("aggregation", &benchmark_name),
            &(device_count, *algorithm),
            |b, (count, algo)| {
                let rt = tokio::runtime::Runtime::new().unwrap();
                let (coordinator, _sender) = FederatedCoordinator::new(
                "bench-coordinator".to_string(),
                SelectionStrategy::Random,
                GradientCompression {
                    algorithm: CompressionAlgorithm::TopK,
                    compression_ratio: 0.1,
                    error_correction: false,
                    adaptive_compression: false,
                },
                AggregationStrategy {
                    algorithm: *algo,
                    weighting: WeightingScheme::Equal,
                    byzantine_tolerance: ByzantineTolerance {
                        enabled: false,
                        max_byzantine_fraction: 0.0,
                        detection_algorithm: ByzantineDetection::None,
                    },
                    differential_privacy: None,
                },
                );

                // Create dummy gradient data
                let mut device_gradients = HashMap::new();
                for i in 0..*count {
                    let device_id = format!("device-{:06}", i);
                    device_gradients.insert(device_id, vec![0u8; 1024]); // 1KB per device
                }

                let round_id = 123456789u64;

                b.to_async(&rt).iter(|| async {
                    black_box(coordinator.aggregate_gradients(round_id, device_gradients.clone()).await.unwrap())
                });
            });
        }
    }

    group.finish();
}

/// Benchmark end-to-end training round performance
fn benchmark_training_round_e2e(c: &mut Criterion) {
    let mut group = c.benchmark_group("Training_Round_E2E");
    group.measurement_time(Duration::from_secs(60)); // Longer measurement for E2E

    let device_counts = vec![100, 1_000, 5_000];

    for device_count in device_counts {
        group.bench_with_input(
        BenchmarkId::new("e2e_training", device_count.to_string()),
        &device_count,
        |b, count| {
            let rt = tokio::runtime::Runtime::new().unwrap();

            b.to_async(&rt).iter_custom(|iters| async move {
                let mut total_duration = Duration::from_secs(0);

                for _ in 0..iters {
                    let start = Instant::now();

                    // Create coordinator
                    let (coordinator, _sender) = FederatedCoordinator::new(
                    "bench-coordinator".to_string(),
                    SelectionStrategy::Intelligent,
                    GradientCompression {
                        algorithm: CompressionAlgorithm::TopK,
                        compression_ratio: 0.01,
                        error_correction: true,
                        adaptive_compression: true,
                    },
                    AggregationStrategy {
                        algorithm: AggregationAlgorithm::FedAvg,
                        weighting: WeightingScheme::Adaptive,
                        byzantine_tolerance: ByzantineTolerance {
                            enabled: true,
                            max_byzantine_fraction: 0.1,
                            detection_algorithm: ByzantineDetection::Krum,
                        },
                        differential_privacy: Some(DifferentialPrivacy {
                            epsilon: 1.0,
                            delta: 1e-5,
                            noise_mechanism: NoiseMechanism::Gaussian,
                            clipping_threshold: 1.0,
                        }),
                    },
                    );

                    // Register devices
                    for i in 0..*count {
                        let device = create_benchmark_device(&format!("device-{:06}", i));
                        coordinator.register_device(device).await.unwrap();
                    }

                    // Select devices
                    let selection_criteria = SelectionCriteria {
                        min_battery_level: 0.3,
                        min_bandwidth_mbps: 5.0,
                        max_latency_ms: 200,
                        required_availability_minutes: 30,
                        min_data_quality: 0.7,
                    };

                    let target_devices = (*count / 10).max(10) as u32;
                    let selection_result = coordinator.select_devices(target_devices, selection_criteria).await.unwrap();

                    // Start training round
                    let training_config = TrainingConfig {
                        local_epochs: 3,
                        local_batch_size: 16,
                        learning_rate: 0.001,
                        gradient_clipping: Some(1.0),
                        early_stopping_patience: Some(5),
                    };

                    let round_id = coordinator.start_training_round(
                    selection_result.selected_devices.clone(),
                    training_config,
                    Duration::from_secs(300),
                    ).await.unwrap();

                    // Simulate gradient aggregation
                    let mut device_gradients = HashMap::new();
                    for device_id in &selection_result.selected_devices {
                        device_gradients.insert(device_id.clone(), vec![0u8; 100]);
                    }

                    coordinator.aggregate_gradients(round_id, device_gradients).await.unwrap();

                    total_duration += start.elapsed();
                }

                total_duration
            });
        });
    }

    group.finish();
}

/// Performance validation benchmark - ensures targets are met
fn benchmark_performance_validation(c: &mut Criterion) {
    let mut group = c.benchmark_group("Performance_Validation");

    group.bench_function("ARM_NEON_Target_2.5x", |b| {
        b.iter(|| {
            let mut optimizer = EdgeTargetOptimizer::new();
            let arm_opts = ArmOptimizations {
                enable_neon: true,
                target_arch: ArmArchitecture::CortexA,
                ..Default::default()
            };
            optimizer.configure_arm(arm_opts);
            let metrics = optimizer.optimize_for_target(EdgeTarget::ARM).unwrap();

            assert!(
            metrics.performance_improvement >= 2.5,
            "ARM NEON target not met: {:.2}x < 2.5x",
            metrics.performance_improvement
            );

            black_box(metrics)
        })
    });

    group.bench_function("RISCV_RVV_Target_4.8x", |b| {
        b.iter(|| {
            let mut optimizer = EdgeTargetOptimizer::new();
            let riscv_opts = RiscVOptimizations {
                enable_rvv: true,
                vector_length: RiscVVectorLength::VLEN512,
                target_variant: RiscVVariant::Vector,
                ..Default::default()
            };
            optimizer.configure_riscv(riscv_opts);
            let metrics = optimizer.optimize_for_target(EdgeTarget::RISCV).unwrap();

            assert!(
            metrics.performance_improvement >= 4.8,
            "RISC-V RVV target not met: {:.2}x < 4.8x",
            metrics.performance_improvement
            );

            black_box(metrics)
        })
    });

    group.bench_function("WASM_SIMD_Target_1.8x", |b| {
        b.iter(|| {
            let mut optimizer = EdgeTargetOptimizer::new();
            let wasm_opts = WasmOptimizations {
                enable_simd: true,
                enable_threads: true,
                ..Default::default()
            };
            optimizer.configure_wasm(wasm_opts);
            let metrics = optimizer.optimize_for_target(EdgeTarget::WASM).unwrap();

            assert!(
            metrics.performance_improvement >= 1.8,
            "WASM SIMD target not met: {:.2}x < 1.8x",
            metrics.performance_improvement
            );

            black_box(metrics)
        })
    });

    group.bench_function("Mobile_GPU_Target_4.0x", |b| {
        b.iter(|| {
            let mut optimizer = EdgeTargetOptimizer::new();
            let mobile_gpu_opts = MobileGpuOptimizations {
                gpu_vendor: MobileGpuVendor::Apple,
                compute_shaders: true,
                tile_based_rendering: true,
                bandwidth_optimization: true,
                power_efficiency: true,
            };
            optimizer.configure_mobile_gpu(mobile_gpu_opts);
            let metrics = optimizer.optimize_for_target(EdgeTarget::MobileGPU).unwrap();

            assert!(
            metrics.performance_improvement >= 4.0,
            "Mobile GPU target not met: {:.2}x < 4.0x",
            metrics.performance_improvement
            );

            black_box(metrics)
        })
    });

    group.bench_function("IoT_Power_Target_10x", |b| {
        b.iter(|| {
            let mut optimizer = EdgeTargetOptimizer::new();
            let iot_opts = IoTOptimizations {
                ultra_low_power: true,
                minimal_memory: true,
                wake_on_inference: true,
                target_platform: IoTPlatform::ESP32,
                ..Default::default()
            };
            optimizer.configure_iot(iot_opts);
            let metrics = optimizer.optimize_for_target(EdgeTarget::IoT).unwrap();

            assert!(
            metrics.power_efficiency_gain >= 10.0,
            "IoT power efficiency target not met: {:.2}x < 10.0x",
            metrics.power_efficiency_gain
            );

            black_box(metrics)
        })
    });

    group.finish();
}

// Helper function to create benchmark test devices
fn create_benchmark_device(device_id: &str) -> FederatedDevice {
    use std::collections::HashSet;

    FederatedDevice {
        device_id: device_id.to_string(),
        device_type: EdgeDeviceType::StandardMobile,
        status: DeviceStatus::Available,
        network_info: NetworkInfo {
            connection_type: ConnectionType::WiFi,
            bandwidth_mbps: 50.0,
            latency_ms: 20,
            reliability: 0.95,
            data_plan: DataPlan {
                unlimited: true,
                monthly_allowance_gb: None,
                current_usage_gb: 0.0,
                cost_per_gb: None,
            },
        },
        power_status: PowerStatus {
            battery_level: 0.8,
            is_charging: false,
            power_source: PowerSource::Battery,
            estimated_battery_life_minutes: Some(240),
        },
        compute_capabilities: ComputeCapabilities {
            cpu_cores: 4,
            ram_mb: 4096,
            has_gpu: false,
            simd_support: true,
            estimated_flops: 1e9,
            memory_bandwidth_gbps: 10.0,
        },
        data_info: DataInfo {
            sample_count: 1000,
            quality_score: 0.9,
            privacy_level: PrivacyLevel::Personal,
            distribution: DataDistribution {
                distribution_type: "normal".to_string(),
                parameters: HashMap::new(),
            },
        },
        availability: AvailabilitySchedule {
            timezone_offset_hours: 0,
            available_hours: (0..24).collect(),
            preferred_duration_minutes: 30,
            blackout_periods: vec![],
        },
        performance_metrics: PerformanceMetrics {
            avg_training_time_seconds: 300.0,
            avg_upload_time_seconds: 10.0,
            accuracy_contribution: 0.85,
            reliability_score: 0.9,
            communication_efficiency: 0.8,
        },
        last_seen: SystemTime::now(),
    }
}

criterion_group!(
edge_benchmarks,
benchmark_arm_neon_optimization,
benchmark_riscv_vector_optimization,
benchmark_wasm_simd_optimization,
benchmark_mobile_gpu_optimization,
benchmark_iot_optimization,
benchmark_federated_device_selection,
benchmark_gradient_aggregation,
benchmark_training_round_e2e,
benchmark_performance_validation
);

criterion_main!(edge_benchmarks);