use anyhow::Result;
use rtx_compress::quantization::vector_quantization::{
    CodebookInitialization, DistanceMetric, VQConfig, VectorQuantizer,
};
use rtx_tensor::{DType, Device, Tensor};

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_codebook_initialization() -> Result<()> {
    let device = Device::try_default()?;

    // Create test data with some structure
    let data = Tensor::randn(&[1000, 128], &device)?;

    let config = VQConfig {
        codebook_size: 256,
        vector_dim: 128,
        max_iterations: 100,
        tolerance: 1e-6,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let mut vq = VectorQuantizer::new(config);

    // Test different initialization methods
    let init_methods = vec![
        CodebookInitialization::Random,
        CodebookInitialization::KMeansPlusPlus,
        CodebookInitialization::FromData,
    ];

    for init_method in init_methods {
        let mut vq_test = VectorQuantizer::new(VQConfig {
            initialization: init_method.clone(),
            ..vq.config().clone()
        });

        vq_test.fit(&data)?;

        // Verify codebook properties
        let codebook = vq_test.codebook();
        assert_eq!(codebook.shape(), &[256, 128]);

        // Test encoding
        let codes = vq_test.encode(&data)?;
        assert_eq!(codes.shape(), &[1000]);

        // All codes should be valid indices
        let max_code = codes.max()?.to_scalar::<f32>()? as usize;
        // For min, negate and use max
        let neg_codes = codes.mul_scalar(-1.0)?;
        let min_code = (neg_codes.max()?.to_scalar::<f32>()? * -1.0) as usize;

        assert!(max_code < 256, "All codes should be < codebook_size");
        assert!(min_code < 256, "Min code should be valid");

        println!(
            "Init method {:?}: max_code={}, min_code={}",
            init_method, max_code, min_code
        );
    }

    Ok(())
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_distance_metrics() -> Result<()> {
    let device = Device::try_default()?;

    let data = Tensor::randn(&[500, 64], &device)?;

    let distance_metrics = vec![
        DistanceMetric::Euclidean,
        DistanceMetric::Cosine,
        DistanceMetric::Manhattan,
    ];

    let mut results = vec![];

    for distance_metric in distance_metrics {
        let config = VQConfig {
            codebook_size: 128,
            vector_dim: 64,
            max_iterations: 50,
            tolerance: 1e-5,
            initialization: CodebookInitialization::KMeansPlusPlus,
            distance_metric: distance_metric.clone(),
        };

        let mut vq = VectorQuantizer::new(config);
        vq.fit(&data)?;

        let codes = vq.encode(&data)?;
        let reconstructed = vq.decode(&codes)?;

        // Compute reconstruction error using the same distance metric
        let reconstruction_error = match distance_metric {
            DistanceMetric::Euclidean => {
                let diff = (data.clone() - reconstructed.clone())?;
                let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
                let squared = diff.pow(&two)?;
                let sum_per_sample = squared.sum(Some(1))?;
                sum_per_sample.mean(&[], false)?.to_scalar::<f32>()?
            }
            DistanceMetric::Cosine => {
                // 1 - cosine similarity
                let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
                let data_norm = data.pow(&two)?.sum(Some(1))?.sqrt()?;
                let recon_norm = reconstructed.pow(&two)?.sum(Some(1))?.sqrt()?;
                let dot_product = (data.clone() * reconstructed.clone())?.sum(Some(1))?;
                let cosine_sim = (dot_product / (data_norm * recon_norm)?)?;
                let one_minus_cosine = cosine_sim.mul_scalar(-1.0)?.add_scalar(1.0)?;
                one_minus_cosine.mean(&[], false)?.to_scalar::<f32>()?
            }
            DistanceMetric::Manhattan => {
                let diff = (data.clone() - reconstructed.clone())?;
                let abs_diff = diff.abs()?;
                let sum_per_sample = abs_diff.sum(Some(1))?;
                sum_per_sample.mean(&[], false)?.to_scalar::<f32>()?
            }
        };

        results.push((distance_metric.clone(), reconstruction_error));

        println!(
            "Distance metric {:?}: reconstruction error = {:.6}",
            distance_metric, reconstruction_error
        );
    }

    // All metrics should produce reasonable results
    for (_, error) in &results {
        assert!(
            *error > 0.0 && *error < 100.0,
            "Reconstruction error should be reasonable"
        );
    }

    Ok(())
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_online_updates() -> Result<()> {
    let device = Device::try_default()?;

    // Initial training data
    let initial_data = Tensor::randn(&[1000, 32], &device)?;

    let config = VQConfig {
        codebook_size: 64,
        vector_dim: 32,
        max_iterations: 30,
        tolerance: 1e-4,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let mut vq = VectorQuantizer::new(config);
    vq.fit(&initial_data)?;

    let initial_codebook = vq.codebook().clone();

    // Enable online updates
    vq.enable_online_updates(true, 0.01); // learning rate = 0.01

    // Stream new data and update codebook online
    for batch in 0..10 {
        let new_data = Tensor::randn(&[100, 32], &device)?;
        vq.update_online(&new_data)?;
    }

    let updated_codebook = vq.codebook().clone();

    // Verify codebook has been updated
    let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
    let codebook_change = (initial_codebook - updated_codebook)?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;

    assert!(
        codebook_change > 1e-4,
        "Codebook should change with online updates, change = {}",
        codebook_change
    );

    // Test encoding with updated codebook
    let test_data = Tensor::randn(&[50, 32], &device)?;
    let codes = vq.encode(&test_data)?;
    let reconstructed = vq.decode(&codes)?;

    let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
    let mse = (test_data.clone() - reconstructed)?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;
    assert!(
        mse < 1.0,
        "Updated VQ should still provide reasonable reconstruction"
    );

    Ok(())
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_hierarchical_quantization() -> Result<()> {
    let device = Device::try_default()?;

    let data = Tensor::randn(&[2000, 256], &device)?;

    // Create two-level hierarchical VQ
    let coarse_config = VQConfig {
        codebook_size: 64, // Coarse level
        vector_dim: 256,
        max_iterations: 50,
        tolerance: 1e-5,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let fine_config = VQConfig {
        codebook_size: 256, // Fine level
        vector_dim: 256,
        max_iterations: 50,
        tolerance: 1e-5,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let mut coarse_vq = VectorQuantizer::new(coarse_config);
    let mut fine_vq = VectorQuantizer::new(fine_config);

    // First level: coarse quantization
    coarse_vq.fit(&data)?;
    let coarse_codes = coarse_vq.encode(&data)?;
    let coarse_reconstructed = coarse_vq.decode(&coarse_codes)?;

    // Second level: quantize residuals
    let residuals = (data.clone() - coarse_reconstructed.clone())?;
    fine_vq.fit(&residuals)?;
    let fine_codes = fine_vq.encode(&residuals)?;
    let fine_reconstructed = fine_vq.decode(&fine_codes)?;

    // Final reconstruction
    let final_reconstructed = (coarse_reconstructed.clone() + fine_reconstructed)?;

    // Verify hierarchical VQ improves reconstruction
    let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
    let coarse_mse = (data.clone() - coarse_reconstructed.clone())?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;
    let hierarchical_mse = (data.clone() - final_reconstructed)?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;

    assert!(
        hierarchical_mse < coarse_mse * 0.8,
        "Hierarchical VQ should improve reconstruction: coarse={:.6}, hierarchical={:.6}",
        coarse_mse,
        hierarchical_mse
    );

    // Verify code storage efficiency
    let total_codes_per_vector = 2; // One coarse + one fine code
    let bits_per_code = (64_f32.log2().ceil() as usize) + (256_f32.log2().ceil() as usize);
    let compression_ratio = (256 * 32) as f32 / bits_per_code as f32; // f32 = 32 bits

    println!(
        "Hierarchical VQ compression ratio: {:.2}x",
        compression_ratio
    );
    assert!(
        compression_ratio > 8.0,
        "Should achieve significant compression"
    );

    Ok(())
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_codebook_pruning() -> Result<()> {
    let device = Device::try_default()?;

    // Create data with clusters of different sizes
    let mut data_parts: Vec<Tensor> = vec![];

    // Large cluster
    let cluster1 = (Tensor::randn(&[800, 64], &device)?
        + Tensor::from_slice(&[2.0], &[1, 1], &device)?.broadcast_to(&[800, 64])?)?;
    data_parts.push(cluster1);

    // Medium cluster
    let cluster2 = (Tensor::randn(&[150, 64], &device)?
        + Tensor::from_slice(&[-2.0], &[1, 1], &device)?.broadcast_to(&[150, 64])?)?;
    data_parts.push(cluster2);

    // Small cluster (should be pruned)
    let cluster3 = (Tensor::randn(&[50, 64], &device)?
        + Tensor::from_slice(&[0.0, 4.0], &[1, 2], &device)?.broadcast_to(&[50, 64])?)?;
    data_parts.push(cluster3);

    let data = Tensor::cat(&data_parts, 0)?;

    let config = VQConfig {
        codebook_size: 128,
        vector_dim: 64,
        max_iterations: 100,
        tolerance: 1e-6,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let mut vq = VectorQuantizer::new(config);
    vq.fit(&data)?;

    // Analyze codebook usage
    let codes = vq.encode(&data)?;
    let usage_stats = vq.analyze_codebook_usage(&codes)?;

    assert_eq!(usage_stats.len(), 128);

    // Find underutilized codewords (< 1% of data)
    let min_usage_threshold = 0.01;
    let underutilized: Vec<_> = usage_stats
        .iter()
        .enumerate()
        .filter(|(_, usage)| **usage < min_usage_threshold)
        .collect();

    println!("Underutilized codewords: {}", underutilized.len());

    // Prune underutilized codewords
    let pruned_vq = vq.prune_codebook(min_usage_threshold)?;

    assert!(
        pruned_vq.codebook_size() < vq.codebook_size(),
        "Pruned codebook should be smaller"
    );

    // Test that pruned VQ still works
    let pruned_codes = pruned_vq.encode(&data)?;
    let pruned_reconstructed = pruned_vq.decode(&pruned_codes)?;

    let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
    let original_mse = (data.clone() - vq.decode(&codes)?)?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;
    let pruned_mse = (data.clone() - pruned_reconstructed)?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;

    // Pruned VQ should have similar quality (within 20%)
    assert!(
        pruned_mse < original_mse * 1.2,
        "Pruned VQ quality should be similar: original={:.6}, pruned={:.6}",
        original_mse,
        pruned_mse
    );

    Ok(())
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_adaptive_codebook_size() -> Result<()> {
    let device = Device::try_default()?;

    // Test data with varying complexity
    let simple_data = (Tensor::randn(&[1000, 32], &device)? * 0.1)?; // Low variance
    let complex_data = (Tensor::randn(&[1000, 32], &device)? * 2.0)?; // High variance

    let config = VQConfig {
        codebook_size: 0, // Will be determined adaptively
        vector_dim: 32,
        max_iterations: 50,
        tolerance: 1e-5,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let mut vq_simple = VectorQuantizer::new(config.clone());
    let mut vq_complex = VectorQuantizer::new(config);

    // Enable adaptive sizing
    let target_distortion = 0.01;
    vq_simple.enable_adaptive_sizing(target_distortion, 512)?; // max 512 codewords
    vq_complex.enable_adaptive_sizing(target_distortion, 512)?;

    // Fit with adaptive sizing
    vq_simple.fit(&simple_data)?;
    vq_complex.fit(&complex_data)?;

    println!("Simple data codebook size: {}", vq_simple.codebook_size());
    println!("Complex data codebook size: {}", vq_complex.codebook_size());

    // Complex data should need more codewords
    assert!(
        vq_complex.codebook_size() > vq_simple.codebook_size(),
        "Complex data should require larger codebook"
    );

    // Both should meet distortion target
    let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
    let simple_codes = vq_simple.encode(&simple_data)?;
    let simple_reconstructed = vq_simple.decode(&simple_codes)?;
    let simple_mse = (simple_data.clone() - simple_reconstructed)?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;

    let complex_codes = vq_complex.encode(&complex_data)?;
    let complex_reconstructed = vq_complex.decode(&complex_codes)?;
    let complex_mse = (complex_data.clone() - complex_reconstructed)?
        .pow(&two)?
        .mean(&[], false)?
        .to_scalar::<f32>()?;

    assert!(
        simple_mse <= (target_distortion * 1.1) as f32,
        "Simple VQ should meet distortion target: {:.6}",
        simple_mse
    );
    assert!(
        complex_mse <= (target_distortion * 1.1) as f32,
        "Complex VQ should meet distortion target: {:.6}",
        complex_mse
    );

    Ok(())
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_batch_operations() -> Result<()> {
    let device = Device::try_default()?;

    let config = VQConfig {
        codebook_size: 128,
        vector_dim: 64,
        max_iterations: 30,
        tolerance: 1e-4,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let mut vq = VectorQuantizer::new(config);

    // Train on batch
    let train_data = Tensor::randn(&[2000, 64], &device)?;
    vq.fit(&train_data)?;

    // Test batch encoding with different shapes
    let test_cases = vec![
        Tensor::randn(&[100, 64], &device)?,     // 2D batch
        Tensor::randn(&[10, 10, 64], &device)?,  // 3D batch
        Tensor::randn(&[5, 4, 5, 64], &device)?, // 4D batch
    ];

    for (i, test_data) in test_cases.iter().enumerate() {
        let original_shape = test_data.shape().to_vec();

        // Encode
        let codes = vq.encode(test_data)?;

        // Verify code shape (last dimension should be removed)
        let expected_code_shape = original_shape[..original_shape.len() - 1].to_vec();
        assert_eq!(
            codes.shape(),
            expected_code_shape.as_slice(),
            "Case {}: code shape mismatch",
            i
        );

        // Decode
        let reconstructed = vq.decode(&codes)?;

        // Verify reconstruction shape matches original
        assert_eq!(
            reconstructed.shape(),
            test_data.shape(),
            "Case {}: reconstruction shape mismatch",
            i
        );

        // Verify reasonable reconstruction quality
        let two = Tensor::from_data(vec![2.0], vec![1], &device)?;
        let mse = (test_data.clone() - reconstructed)?
            .pow(&two)?
            .mean(&[], false)?
            .to_scalar::<f32>()?;
        assert!(
            mse < 2.0,
            "Case {}: reconstruction quality poor: MSE={:.6}",
            i,
            mse
        );
    }

    Ok(())
}

#[test]
#[ignore = "Pre-existing Metal device randn issue"]
fn test_vq_serialization() -> Result<()> {
    let device = Device::try_default()?;

    let data = Tensor::randn(&[500, 128], &device)?;

    let config = VQConfig {
        codebook_size: 256,
        vector_dim: 128,
        max_iterations: 50,
        tolerance: 1e-5,
        initialization: CodebookInitialization::KMeansPlusPlus,
        distance_metric: DistanceMetric::Euclidean,
    };

    let mut vq = VectorQuantizer::new(config);
    vq.fit(&data)?;

    // Test serialization
    let serialized = vq.serialize()?;

    // Deserialize into new VQ
    let vq2 = VectorQuantizer::deserialize(&serialized)?;

    // Verify both VQs produce identical results
    let codes1 = vq.encode(&data)?;
    let codes2 = vq2.encode(&data)?;

    let code_diff = (codes1.to_dtype(DType::F32)? - codes2.to_dtype(DType::F32)?)?
        .abs()?
        .max()?
        .to_scalar::<f32>()?;

    assert!(
        code_diff < 1e-6,
        "Serialized VQ should produce identical codes"
    );

    // Verify codebooks are identical
    let codebook_diff = (vq.codebook().clone() - vq2.codebook().clone())?
        .abs()?
        .max()?
        .to_scalar::<f32>()?;

    assert!(
        codebook_diff < 1e-6,
        "Serialized codebook should be identical"
    );

    Ok(())
}