use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use rtx_compress::{
    checkpoint::{CheckpointCompressor, CheckpointFormat, CompressionConfig},
    kv_cache::{CompressedKVCache, CompressionMethod, KVCacheConfig},
    quantization::{
        mixed_precision::{MixedPrecisionOptimizer, OptimizationObjective, PrecisionConfig},
        product_quantization::{PQConfig, ProductQuantizer},
        vector_quantization::{CodebookInitialization, DistanceMetric, VQConfig, VectorQuantizer},
    },
};
use rtx_tensor::{DType, Device, Tensor};
use std::collections::HashMap;

fn bench_product_quantization(c: &mut Criterion) {
    let device = Device::try_default().unwrap();

    let mut group = c.benchmark_group("product_quantization");

    // Test different data sizes
    let sizes = vec![
        ("small", 1000, 128),
        ("medium", 10000, 256),
        ("large", 100000, 512),
    ];

    for (name, num_vectors, dim) in sizes {
        let data = Tensor::randn(&[num_vectors, dim], DType::F32, &device).unwrap();

        // Benchmark codebook learning
        group.bench_with_input(BenchmarkId::new("fit", name), &data, |b, data| {
            b.iter(|| {
                let config = PQConfig {
                    num_subquantizers: 8,
                    codebook_size: 256,
                    max_iterations: 20, // Reduced for benchmarking
                    tolerance: 1e-4,
                };
                let mut pq = ProductQuantizer::new(config);
                black_box(pq.fit(data).unwrap());
            });
        });

        // Pre-train a quantizer for encoding/decoding benchmarks
        let config = PQConfig {
            num_subquantizers: 8,
            codebook_size: 256,
            max_iterations: 10,
            tolerance: 1e-3,
        };
        let mut pq = ProductQuantizer::new(config);
        pq.fit(&data).unwrap();

        // Benchmark encoding
        group.bench_with_input(
            BenchmarkId::new("encode", name),
            &(&data, &pq),
            |b, (data, pq)| {
                b.iter(|| {
                    black_box(pq.encode(data).unwrap());
                });
            },
        );

        // Benchmark decoding
        let codes = pq.encode(&data).unwrap();
        group.bench_with_input(
            BenchmarkId::new("decode", name),
            &(&codes, &pq),
            |b, (codes, pq)| {
                b.iter(|| {
                    black_box(pq.decode(codes).unwrap());
                });
            },
        );
    }

    group.finish();
}

fn bench_vector_quantization(c: &mut Criterion) {
    let device = Device::try_default().unwrap();

    let mut group = c.benchmark_group("vector_quantization");

    let data = Tensor::randn(&[50000, 256], DType::F32, &device).unwrap();

    // Test different codebook sizes
    let codebook_sizes = vec![128, 256, 512, 1024];

    for codebook_size in codebook_sizes {
        let config = VQConfig {
            codebook_size,
            vector_dim: 256,
            max_iterations: 20,
            tolerance: 1e-4,
            initialization: CodebookInitialization::KMeansPlusPlus,
            distance_metric: DistanceMetric::Euclidean,
        };

        group.bench_with_input(
            BenchmarkId::new("kmeans_fit", codebook_size),
            &config,
            |b, config| {
                b.iter(|| {
                    let mut vq = VectorQuantizer::new(config.clone());
                    black_box(vq.fit(&data).unwrap());
                });
            },
        );

        // Pre-train for encoding benchmarks
        let mut vq = VectorQuantizer::new(config);
        vq.fit(&data).unwrap();

        group.bench_with_input(
            BenchmarkId::new("encode", codebook_size),
            &(&data, &vq),
            |b, (data, vq)| {
                b.iter(|| {
                    black_box(vq.encode(data).unwrap());
                });
            },
        );
    }

    group.finish();
}

fn bench_kv_cache_compression(c: &mut Criterion) {
    let device = Device::try_default().unwrap();

    let mut group = c.benchmark_group("kv_cache");

    // Test different sequence lengths
    let seq_lengths = vec![256, 512, 1024, 2048];

    for seq_len in seq_lengths {
        let keys = Tensor::randn(&[4, 12, seq_len, 64], DType::F32, &device).unwrap();
        let values = Tensor::randn(&[4, 12, seq_len, 64], DType::F32, &device).unwrap();

        // Product Quantization method
        let pq_config = KVCacheConfig {
            compression_method: CompressionMethod::ProductQuantization {
                num_subquantizers: 8,
                codebook_size: 256,
            },
            compression_ratio_target: 4.0,
            quality_threshold: 0.90,
        };

        group.bench_with_input(
            BenchmarkId::new("pq_insert", seq_len),
            &(&keys, &values, &pq_config),
            |b, (keys, values, config)| {
                b.iter_with_setup(
                    || CompressedKVCache::new(config.clone()),
                    |mut cache| {
                        black_box(cache.insert(0, keys, values).unwrap());
                    },
                );
            },
        );

        // Vector Quantization method
        let vq_config = KVCacheConfig {
            compression_method: CompressionMethod::VectorQuantization {
                codebook_size: 1024,
                update_frequency: 100,
            },
            compression_ratio_target: 3.0,
            quality_threshold: 0.88,
        };

        group.bench_with_input(
            BenchmarkId::new("vq_insert", seq_len),
            &(&keys, &values, &vq_config),
            |b, (keys, values, config)| {
                b.iter_with_setup(
                    || CompressedKVCache::new(config.clone()),
                    |mut cache| {
                        black_box(cache.insert(0, keys, values).unwrap());
                    },
                );
            },
        );

        // Benchmark retrieval
        let mut cache = CompressedKVCache::new(pq_config);
        cache.insert(0, &keys, &values).unwrap();

        group.bench_with_input(
            BenchmarkId::new("retrieve", seq_len),
            &(&cache, seq_len),
            |b, (cache, seq_len)| {
                b.iter(|| {
                    black_box(cache.get(0, 0, *seq_len).unwrap());
                });
            },
        );
    }

    group.finish();
}

fn bench_checkpoint_compression(c: &mut Criterion) {
    let device = Device::try_default().unwrap();

    let mut group = c.benchmark_group("checkpoint_compression");

    // Create model of different sizes
    let model_sizes = vec![
        ("small", vec![(512, 1024), (256, 512)]),
        ("medium", vec![(2048, 4096), (1024, 2048), (512, 1024)]),
        (
            "large",
            vec![(4096, 8_192), (2048, 4096), (1024, 2048), (512, 1024)],
        ),
    ];

    for (size_name, layer_dims) in model_sizes {
        let mut state_dict = HashMap::new();
        for (i, (in_dim, out_dim)) in layer_dims.iter().enumerate() {
            state_dict.insert(
                format!("layer_{}.weight", i),
                Tensor::randn(&[*out_dim, *in_dim], DType::F32, &device).unwrap(),
            );
        }

        // Test different compression formats
        let formats = vec![
            ("lz4", CheckpointFormat::Lz4, 1),
            ("zstd_fast", CheckpointFormat::Zstd, 1),
            ("zstd_balanced", CheckpointFormat::Zstd, 6),
            ("zstd_best", CheckpointFormat::Zstd, 22),
        ];

        for (format_name, format, compression_level) in formats {
            let config = CompressionConfig {
                format,
                compression_level,
                quantization_bits: 16, // fp16
                exclude_patterns: vec![],
            };

            let compressor = CheckpointCompressor::new(config);

            // Benchmark save
            group.bench_with_input(
                BenchmarkId::new(format!("save_{}_{}", size_name, format_name), ""),
                &(&state_dict, &compressor),
                |b, (state_dict, compressor)| {
                    b.iter(|| {
                        black_box(compressor.save(state_dict).unwrap());
                    });
                },
            );

            // Pre-compress for load benchmark
            let compressed_data = compressor.save(&state_dict).unwrap();

            // Benchmark load
            group.bench_with_input(
                BenchmarkId::new(format!("load_{}_{}", size_name, format_name), ""),
                &(&compressed_data, &compressor),
                |b, (compressed_data, compressor)| {
                    b.iter(|| {
                        black_box(compressor.load(compressed_data).unwrap());
                    });
                },
            );
        }
    }

    group.finish();
}

fn bench_mixed_precision_optimization(c: &mut Criterion) {
    let device = Device::try_default().unwrap();

    let mut group = c.benchmark_group("mixed_precision");

    // Create different sized model architectures
    let architectures = vec![
        ("transformer_small", 6, 512, 2048), // 6 layers, 512 hidden, 2048 ffn
        ("transformer_base", 12, 768, 3072), // 12 layers, 768 hidden, 3072 ffn
        ("transformer_large", 24, 1024, 4096), // 24 layers, 1024 hidden, 4096 ffn
    ];

    for (arch_name, num_layers, hidden_dim, ffn_dim) in architectures {
        let mut layers = HashMap::new();

        for i in 0..num_layers {
            layers.insert(
                format!("layers.{}.attention.query", i),
                Tensor::randn(&[hidden_dim, hidden_dim], DType::F32, &device).unwrap(),
            );
            layers.insert(
                format!("layers.{}.attention.key", i),
                Tensor::randn(&[hidden_dim, hidden_dim], DType::F32, &device).unwrap(),
            );
            layers.insert(
                format!("layers.{}.attention.value", i),
                Tensor::randn(&[hidden_dim, hidden_dim], DType::F32, &device).unwrap(),
            );
            layers.insert(
                format!("layers.{}.ffn.up_proj", i),
                Tensor::randn(&[hidden_dim, ffn_dim], DType::F32, &device).unwrap(),
            );
            layers.insert(
                format!("layers.{}.ffn.down_proj", i),
                Tensor::randn(&[ffn_dim, hidden_dim], DType::F32, &device).unwrap(),
            );
        }

        let calibration_data = Tensor::randn(&[100, hidden_dim], DType::F32, &device).unwrap();

        // Benchmark sensitivity analysis
        let config = PrecisionConfig {
            precision_bits: vec![4, 8, 12, 16],
            sensitivity_threshold: 0.02,
            performance_weight: 0.6,
            quality_weight: 0.4,
        };

        group.bench_with_input(
            BenchmarkId::new("sensitivity_analysis", arch_name),
            &(&layers, &calibration_data, &config),
            |b, (layers, calibration_data, config)| {
                b.iter(|| {
                    let mut optimizer = MixedPrecisionOptimizer::new(config.clone());
                    black_box(
                        optimizer
                            .analyze_sensitivity(layers, calibration_data)
                            .unwrap(),
                    );
                });
            },
        );

        // Benchmark full optimization
        let objective = OptimizationObjective {
            target_compression_ratio: 4.0,
            max_quality_loss: 0.05,
            memory_constraint_mb: Some(1000),
        };

        group.bench_with_input(
            BenchmarkId::new("full_optimization", arch_name),
            &(&layers, &calibration_data, &config, &objective),
            |b, (layers, calibration_data, config, objective)| {
                b.iter(|| {
                    let mut optimizer = MixedPrecisionOptimizer::new(config.clone());
                    black_box(
                        optimizer
                            .optimize(layers, calibration_data, objective.clone())
                            .unwrap(),
                    );
                });
            },
        );
    }

    group.finish();
}

fn bench_compression_throughput(c: &mut Criterion) {
    let device = Device::try_default().unwrap();

    let mut group = c.benchmark_group("throughput");
    group.throughput(criterion::Throughput::Bytes(1024 * 1024 * 4)); // 1M f32 values

    let data = Tensor::randn(&[1024, 1024], DType::F32, &device).unwrap();

    // Benchmark different compression algorithms
    let algorithms = vec![
        ("pq_8x256", |data: &Tensor| {
            let config = PQConfig {
                num_subquantizers: 8,
                codebook_size: 256,
                max_iterations: 10,
                tolerance: 1e-3,
            };
            let mut pq = ProductQuantizer::new(config);
            pq.fit(data).unwrap();
            pq.encode(data).unwrap();
        }),
        ("vq_1024", |data: &Tensor| {
            let config = VQConfig {
                codebook_size: 1024,
                vector_dim: 1024,
                max_iterations: 10,
                tolerance: 1e-3,
                initialization: CodebookInitialization::Random,
                distance_metric: DistanceMetric::Euclidean,
            };
            let mut vq = VectorQuantizer::new(config);
            vq.fit(data).unwrap();
            vq.encode(data).unwrap();
        }),
    ];

    for (name, algorithm) in algorithms {
        group.bench_function(name, |b| {
            b.iter(|| {
                black_box(algorithm(&data));
            });
        });
    }

    group.finish();
}

fn bench_memory_efficiency(c: &mut Criterion) {
    let device = Device::try_default().unwrap();

    let mut group = c.benchmark_group("memory_efficiency");

    // Test compression with different memory constraints
    let memory_limits = vec![100, 200, 500, 1000]; // MB

    for memory_limit in memory_limits {
        let data_size_mb = memory_limit * 2; // 2x memory pressure
        let num_elements = (data_size_mb * 1024 * 1024) / 4; // f32 = 4 bytes
        let data = Tensor::randn(&[num_elements], DType::F32, &device).unwrap();

        group.bench_with_input(
            BenchmarkId::new("constrained_compression", memory_limit),
            &(&data, memory_limit),
            |b, (data, memory_limit)| {
                b.iter(|| {
                    // Simulate memory-constrained compression
                    let config = PQConfig {
                        num_subquantizers: 8,
                        codebook_size: 128, // Smaller codebook for memory constraint
                        max_iterations: 5,  // Fewer iterations
                        tolerance: 1e-2,    // Less strict tolerance
                    };

                    let mut pq = ProductQuantizer::new(config);

                    // In real implementation, this would respect memory limits
                    let reshaped = data.reshape(&[num_elements / 128, 128]).unwrap();
                    pq.fit(&reshaped).unwrap();
                    black_box(pq.encode(&reshaped).unwrap());
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_product_quantization,
    bench_vector_quantization,
    bench_kv_cache_compression,
    bench_checkpoint_compression,
    bench_mixed_precision_optimization,
    bench_compression_throughput,
    bench_memory_efficiency
);

criterion_main!(benches);