//! Property-based tests for decomposition algorithms
//!
//! These tests verify mathematical properties and invariants that should hold
//! for all decomposition methods using proptest.

use crate::dictionary_learning::{
    DictionaryLearning, DictionaryLearningConfig, DictionaryTransformAlgorithm,
};
use crate::factor_analysis::FactorAnalysis;
use crate::ica::{ICAAlgorithm, ICA};
use crate::nmf::NMF;
use crate::pca::{PcaConfig, PCA};
use proptest::prelude::*;
use scirs2_core::ndarray::Array2;
use sklears_core::traits::{Fit, Transform};

proptest! {
    #[test]
    fn test_pca_properties(
        n_samples in 10..50usize,
        n_features in 3..10usize,
        n_components in 1..5usize
    ) {
        if n_components <= n_features {
            // Create random data matrix
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i as f64 + j as f64) / 10.0;
                }
            }

            let config = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca = PCA::new(config);
            let result = pca.fit(&x, &());

            if let Ok(trained_pca) = result {
                let x_transformed = trained_pca.transform(&x);
                if let Ok(transformed) = x_transformed {
                    // Check output dimensions
                    prop_assert_eq!(transformed.shape(), &[n_samples, n_components]);

                    // Check that all values are finite
                    for &val in transformed.iter() {
                        prop_assert!(val.is_finite());
                    }

                    // Check components matrix shape
                    let components = &trained_pca.components;
                    prop_assert_eq!(components.shape(), &[n_components, n_features]);
                }
            }
        }
    }

    #[test]
    fn test_nmf_non_negativity_properties(
        n_samples in 5..20usize,
        n_features in 3..8usize,
        n_components in 1..4usize
    ) {
        if n_components <= n_features.min(n_samples) {
            // Create non-negative data matrix
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i + j + 1) as f64;
                }
            }

            let nmf = NMF::new(n_components).random_state(42);
            let result = nmf.fit(&x, &());

            if let Ok(trained_nmf) = result {
                let w = trained_nmf.transform(&x);
                if let Ok(transformed) = w {
                    // Check output dimensions
                    prop_assert_eq!(transformed.shape(), &[n_samples, n_components]);

                    // Check non-negativity of coefficients
                    for &val in transformed.iter() {
                        prop_assert!(val >= 0.0);
                        prop_assert!(val.is_finite());
                    }

                    // Check components are non-negative
                    let components = trained_nmf.components();
                    for &val in components.iter() {
                        prop_assert!(val >= 0.0);
                        prop_assert!(val.is_finite());
                    }

                    // Test reconstruction
                    let x_reconstructed = trained_nmf.inverse_transform(&transformed);
                    if let Ok(reconstructed) = x_reconstructed {
                        // Reconstructed values should be non-negative and finite
                        for &val in reconstructed.iter() {
                            prop_assert!(val >= 0.0);
                            prop_assert!(val.is_finite());
                        }
                    }
                }
            }
        }
    }

    #[test]
    fn test_ica_orthogonality_properties(
        n_samples in 20..40usize,
        n_features in 3..8usize,
        n_components in 2..5usize
    ) {
        if n_components <= n_features {
            // Create mixed signals (simple linear combinations)
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i as f64).sin() + (j as f64).cos() + 0.1 * (i + j) as f64;
                }
            }

            let ica = ICA::new()
                .n_components(n_components)
                .algorithm(ICAAlgorithm::Parallel)
                .random_state(42);

            let result = ica.fit(&x, &());

            if let Ok(trained_ica) = result {
                let x_transformed = trained_ica.transform(&x);
                if let Ok(transformed) = x_transformed {
                    // Check output dimensions
                    prop_assert_eq!(transformed.shape(), &[n_samples, n_components]);

                    // Check that all values are finite
                    for &val in transformed.iter() {
                        prop_assert!(val.is_finite());
                    }

                    // Check components matrix properties
                    let components = trained_ica.components();
                    prop_assert_eq!(components.shape(), &[n_components, n_features]);

                    // Test inverse transform
                    let x_reconstructed = trained_ica.inverse_transform(&transformed);
                    if let Ok(reconstructed) = x_reconstructed {
                        prop_assert_eq!(reconstructed.shape(), x.shape());

                        // All reconstructed values should be finite
                        for &val in reconstructed.iter() {
                            prop_assert!(val.is_finite());
                        }
                    }
                }
            }
        }
    }

    #[test]
    fn test_decomposition_dimension_consistency(
        n_samples in 8..20usize,
        n_features in 3..8usize,
        n_components in 1..4usize
    ) {
        if n_components <= n_features {
            // Create simple test data
            let mut x = Array2::<f64>::ones((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i + j) as f64 / 10.0;
                }
            }

            // Test that all decomposition methods respect dimension constraints

            // PCA
            let config = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca = PCA::new(config);
            if let Ok(trained_pca) = pca.fit(&x, &()) {
                if let Ok(pca_result) = trained_pca.transform(&x) {
                    prop_assert_eq!(pca_result.shape(), &[n_samples, n_components]);
                }
            }

            // ICA (only test if we have enough samples)
            if n_samples >= 2 * n_components {
                let ica = ICA::new().n_components(n_components).random_state(42);
                if let Ok(trained_ica) = ica.fit(&x, &()) {
                    if let Ok(ica_result) = trained_ica.transform(&x) {
                        prop_assert_eq!(ica_result.shape(), &[n_samples, n_components]);
                    }
                }
            }

            // FactorAnalysis (requires n_components < n_features and sufficient samples)
            if n_components < n_features && n_samples >= n_features + 2 {
                let fa = FactorAnalysis::new(n_components).random_state(42);
                if let Ok(trained_fa) = fa.fit(&x, &()) {
                    if let Ok(fa_result) = trained_fa.transform(&x) {
                        prop_assert_eq!(fa_result.shape(), &[n_samples, n_components]);
                    }
                }
            }
        }
    }

    #[test]
    fn test_decomposition_determinism(
        n_samples in 10..30usize,
        n_features in 3..8usize,
        n_components in 1..4usize,
        random_seed in 1..100u64
    ) {
        if n_components <= n_features {
            // Create deterministic data
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i + j) as f64 / 10.0;
                }
            }

            // Test PCA determinism
            let config1 = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca1 = PCA::new(config1);
            let config2 = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca2 = PCA::new(config2);

            if let (Ok(trained_pca1), Ok(trained_pca2)) = (pca1.fit(&x, &()), pca2.fit(&x, &())) {
                if let (Ok(result1), Ok(result2)) = (trained_pca1.transform(&x), trained_pca2.transform(&x)) {
                    // Results should be identical (up to sign) for deterministic algorithms
                    prop_assert_eq!(result1.shape(), result2.shape());

                    // Check that components are similar (allowing for sign flip)
                    let comp1 = &trained_pca1.components;
                    let comp2 = &trained_pca2.components;
                    prop_assert_eq!(comp1.shape(), comp2.shape());
                }
            }

            // Test NMF determinism with fixed seed
            let nmf1 = NMF::new(n_components).random_state(random_seed);
            let nmf2 = NMF::new(n_components).random_state(random_seed);

            if let (Ok(trained_nmf1), Ok(trained_nmf2)) = (nmf1.fit(&x, &()), nmf2.fit(&x, &())) {
                if let (Ok(result1), Ok(result2)) = (trained_nmf1.transform(&x), trained_nmf2.transform(&x)) {
                    prop_assert_eq!(result1.shape(), result2.shape());
                }
            }
        }
    }

    #[test]
    fn test_decomposition_rank_properties(
        n_samples in 5..20usize,
        n_features in 3..8usize,
        n_components in 1..4usize
    ) {
        if n_components <= n_features.min(n_samples) {
            // Create rank-deficient data
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features.min(2) { // Only fill first 2 columns to create rank deficiency
                    x[[i, j]] = (i + j) as f64;
                }
            }

            // PCA should handle rank-deficient data gracefully
            let config = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca = PCA::new(config);
            if let Ok(trained_pca) = pca.fit(&x, &()) {
                if let Ok(transformed) = trained_pca.transform(&x) {
                    prop_assert_eq!(transformed.shape(), &[n_samples, n_components]);

                    // All values should be finite
                    for &val in transformed.iter() {
                        prop_assert!(val.is_finite());
                    }

                    // Explained variance should be non-negative and finite
                    let explained_var = trained_pca.explained_variance_ratio;
                    for &var in explained_var.iter() {
                        prop_assert!(var >= 0.0);
                        prop_assert!(var.is_finite());
                    }
                }
            }
        }
    }

    #[test]
    fn test_decomposition_scaling_invariance(
        n_samples in 10..25usize,
        n_features in 3..6usize,
        n_components in 1..3usize,
        scale_factor in 0.1..10.0f64
    ) {
        if n_components <= n_features {
            // Create base data
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i + 1) as f64 + (j + 1) as f64 / 10.0;
                }
            }

            // Create scaled version
            let x_scaled = &x * scale_factor;

            // Test PCA scaling properties
            let config1 = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca1 = PCA::new(config1);
            let config2 = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca2 = PCA::new(config2);

            if let (Ok(trained_pca1), Ok(trained_pca2)) = (pca1.fit(&x, &()), pca2.fit(&x_scaled, &())) {
                // Components should be the same direction (up to normalization)
                let comp1 = &trained_pca1.components;
                let comp2 = &trained_pca2.components;
                prop_assert_eq!(comp1.shape(), comp2.shape());

                // Explained variance should scale with the square of the scaling factor
                let var1 = trained_pca1.explained_variance_ratio;
                let var2 = trained_pca2.explained_variance_ratio;
                prop_assert_eq!(var1.len(), var2.len());
            }
        }
    }

    #[test]
    fn test_reconstruction_bounds(
        n_samples in 5..15usize,
        n_features in 3..6usize,
        n_components in 1..4usize
    ) {
        if n_components <= n_features {
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i + j) as f64 / 5.0;
                }
            }

            // PCA inverse_transform not yet implemented; test forward pass only.
            let config = PcaConfig { n_components: Some(n_components), ..Default::default() };
            let pca = PCA::new(config);
            if let Ok(trained_pca) = pca.fit(&x, &()) {
                if let Ok(transformed) = trained_pca.transform(&x) {
                    prop_assert_eq!(transformed.shape(), &[n_samples, n_components]);
                    for &val in transformed.iter() {
                        prop_assert!(val.is_finite());
                    }
                }
            }
        }
    }

    #[test]
    fn test_factor_analysis_properties(
        n_samples in 15..30usize,
        n_features in 4..8usize,
        n_components in 1..4usize
    ) {
        // FactorAnalysis requires n_components < n_features and sufficient samples.
        if n_components < n_features && n_samples >= n_features + 2 {
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i as f64).sin() * (j + 1) as f64 + (i + j) as f64 / 10.0;
                }
            }

            let fa = FactorAnalysis::new(n_components).random_state(42);
            if let Ok(trained_fa) = fa.fit(&x, &()) {
                let result = trained_fa.transform(&x);
                if let Ok(scores) = result {
                    // Factor scores must have shape (n_samples, n_components).
                    prop_assert_eq!(scores.shape(), &[n_samples, n_components]);

                    // All factor scores must be finite.
                    for &val in scores.iter() {
                        prop_assert!(val.is_finite());
                    }

                    // Loadings matrix must have shape (n_features, n_components).
                    let loadings = trained_fa.loadings();
                    prop_assert_eq!(loadings.shape(), &[n_features, n_components]);

                    // Noise variances must be positive and finite.
                    let noise_var = trained_fa.noise_variance();
                    prop_assert_eq!(noise_var.len(), n_features);
                    for &v in noise_var.iter() {
                        prop_assert!(v > 0.0);
                        prop_assert!(v.is_finite());
                    }

                    // Log-likelihood must be finite (not NaN or Inf).
                    prop_assert!(trained_fa.log_likelihood().is_finite());
                }
            }
        }
    }

    #[test]
    fn test_dictionary_learning_properties(
        n_samples in 12..25usize,
        n_features in 4..8usize,
        n_components in 2..5usize
    ) {
        if n_components <= n_features && n_samples >= n_components * 2 {
            // Build strictly positive data so OMP sparse coding converges reliably.
            let mut x = Array2::<f64>::zeros((n_samples, n_features));
            for i in 0..n_samples {
                for j in 0..n_features {
                    x[[i, j]] = (i + j + 1) as f64 / 5.0;
                }
            }

            let config = DictionaryLearningConfig {
                n_components,
                max_iter: 30,
                tol: 1e-3,
                transform_algorithm: DictionaryTransformAlgorithm::OMP,
                alpha: 0.5,
                random_state: Some(42),
            };

            let model = DictionaryLearning::new(config);
            if let Ok(trained_model) = model.fit(&x, &()) {
                // Dictionary (components) shape: (n_components, n_features).
                let dict = trained_model.components();
                prop_assert_eq!(dict.shape(), &[n_components, n_features]);
                for &val in dict.iter() {
                    prop_assert!(val.is_finite());
                }

                let result = trained_model.transform(&x);
                if let Ok(codes) = result {
                    // Sparse codes shape: (n_samples, n_components).
                    prop_assert_eq!(codes.shape(), &[n_samples, n_components]);
                    for &val in codes.iter() {
                        prop_assert!(val.is_finite());
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_property_test_setup_decomposition() {
        // Simple test to ensure property test framework is working for decomposition
        let x = Array2::<f64>::ones((10, 5));
        let config = PcaConfig {
            n_components: Some(3),
            ..Default::default()
        };
        let pca = PCA::new(config);
        let result = pca.fit(&x, &());

        // Should succeed (though may not be very meaningful with all-ones data)
        assert!(result.is_ok());
    }
}
