sklears_mixture/

multi_modal.rs

//! Multi-Modal Data Mixture Models
//!
//! This module provides mixture models that can handle multi-modal and heterogeneous data types.
//! It includes implementations for multi-view mixture models, coupled mixture models,
//! shared latent variable models, and cross-modal alignment techniques.
//! All implementations follow SciRS2 Policy for numerical computing and random number generation.

use crate::common::{CovarianceType, ModelSelection};
use scirs2_core::ndarray::{s, Array1, Array2, Array3, Axis};
use scirs2_core::random::{thread_rng, RandNormal, RandUniform};
use sklears_core::{
    error::{Result as SklResult, SklearsError},
    traits::{Estimator, Fit, Predict, Trained, Untrained},
};
use std::collections::HashMap;

/// Type of multi-modal data fusion strategy
#[derive(Debug, Clone, PartialEq)]
pub enum FusionStrategy {
    /// Early fusion: Concatenate features from all modalities
    EarlyFusion,
    /// Late fusion: Train separate models then combine predictions
    LateFusion,
    /// Intermediate fusion: Learn shared latent representation
    IntermediateFusion,
    /// Coupled fusion: Joint optimization across modalities
    CoupledFusion,
}

/// Multi-view data modality specification
#[derive(Debug, Clone)]
pub struct ModalitySpec {
    /// Name of the modality (e.g., "visual", "textual", "audio")
    pub name: String,
    /// Feature dimension for this modality
    pub n_features: usize,
    /// Covariance type for this modality
    pub covariance_type: CovarianceType,
    /// Weight for this modality in the fusion process
    pub modality_weight: f64,
}

/// Configuration for multi-modal mixture models
#[derive(Debug, Clone)]
pub struct MultiModalConfig {
    /// n_components
    pub n_components: usize,
    /// modalities
    pub modalities: Vec<ModalitySpec>,
    /// fusion_strategy
    pub fusion_strategy: FusionStrategy,
    /// shared_latent_dim
    pub shared_latent_dim: Option<usize>,
    /// coupling_strength
    pub coupling_strength: f64,
    /// max_iter
    pub max_iter: usize,
    /// tol
    pub tol: f64,
    /// regularization_strength
    pub regularization_strength: f64,
    /// random_state
    pub random_state: Option<u64>,
}

/// Multi-Modal Gaussian Mixture Model
///
/// A mixture model that can handle multiple data modalities simultaneously,
/// learning both modality-specific patterns and cross-modal relationships.
/// This is useful for datasets with multiple types of features (e.g., visual + textual,
/// sensor data from multiple sources, etc.).
#[derive(Debug, Clone)]
pub struct MultiModalGaussianMixture<S = Untrained> {
    config: MultiModalConfig,
    _phantom: std::marker::PhantomData<S>,
}

/// Trained Multi-Modal GMM
#[derive(Debug, Clone)]
pub struct MultiModalGaussianMixtureTrained {
    /// weights
    pub weights: Array1<f64>,
    /// modality_means
    pub modality_means: HashMap<String, Array2<f64>>,
    /// modality_covariances
    pub modality_covariances: HashMap<String, Array3<f64>>,
    /// shared_latent_means
    pub shared_latent_means: Option<Array2<f64>>,
    /// latent_projections
    pub latent_projections: HashMap<String, Array2<f64>>,
    /// coupling_parameters
    pub coupling_parameters: Array2<f64>,
    /// log_likelihood_history
    pub log_likelihood_history: Vec<f64>,
    /// n_iter
    pub n_iter: usize,
    /// config
    pub config: MultiModalConfig,
}

/// Builder for Multi-Modal GMM
#[derive(Debug, Clone)]
pub struct MultiModalGaussianMixtureBuilder {
    n_components: usize,
    modalities: Vec<ModalitySpec>,
    fusion_strategy: FusionStrategy,
    shared_latent_dim: Option<usize>,
    coupling_strength: f64,
    max_iter: usize,
    tol: f64,
    regularization_strength: f64,
    random_state: Option<u64>,
}

impl MultiModalGaussianMixtureBuilder {
    /// Create a new builder with specified number of components
    pub fn new(n_components: usize) -> Self {
        Self {
            n_components,
            modalities: Vec::new(),
            fusion_strategy: FusionStrategy::IntermediateFusion,
            shared_latent_dim: None,
            coupling_strength: 0.1,
            max_iter: 100,
            tol: 1e-4,
            regularization_strength: 0.01,
            random_state: None,
        }
    }

    /// Add a data modality
    pub fn add_modality(mut self, modality: ModalitySpec) -> Self {
        self.modalities.push(modality);
        self
    }

    /// Add a data modality with default settings
    pub fn add_modality_simple(mut self, name: &str, n_features: usize) -> Self {
        let modality = ModalitySpec {
            name: name.to_string(),
            n_features,
            covariance_type: CovarianceType::Full,
            modality_weight: 1.0,
        };
        self.modalities.push(modality);
        self
    }

    /// Set fusion strategy
    pub fn fusion_strategy(mut self, strategy: FusionStrategy) -> Self {
        self.fusion_strategy = strategy;
        self
    }

    /// Set shared latent dimension (for intermediate fusion)
    pub fn shared_latent_dim(mut self, dim: usize) -> Self {
        self.shared_latent_dim = Some(dim);
        self
    }

    /// Set coupling strength between modalities
    pub fn coupling_strength(mut self, strength: f64) -> Self {
        self.coupling_strength = strength.clamp(0.0, 1.0);
        self
    }

    /// Set maximum iterations
    pub fn max_iter(mut self, max_iter: usize) -> Self {
        self.max_iter = max_iter;
        self
    }

    /// Set convergence tolerance
    pub fn tolerance(mut self, tol: f64) -> Self {
        self.tol = tol;
        self
    }

    /// Set regularization strength
    pub fn regularization_strength(mut self, strength: f64) -> Self {
        self.regularization_strength = strength.max(0.0);
        self
    }

    /// Set random state for reproducibility
    pub fn random_state(mut self, random_state: u64) -> Self {
        self.random_state = Some(random_state);
        self
    }

    /// Build the multi-modal GMM
    pub fn build(self) -> SklResult<MultiModalGaussianMixture<Untrained>> {
        if self.modalities.is_empty() {
            return Err(SklearsError::InvalidInput(
                "At least one modality must be specified".to_string(),
            ));
        }

        // For intermediate fusion, ensure latent dimension is specified
        if self.fusion_strategy == FusionStrategy::IntermediateFusion
            && self.shared_latent_dim.is_none()
        {
            return Err(SklearsError::InvalidInput(
                "Shared latent dimension must be specified for intermediate fusion".to_string(),
            ));
        }

        let config = MultiModalConfig {
            n_components: self.n_components,
            modalities: self.modalities,
            fusion_strategy: self.fusion_strategy,
            shared_latent_dim: self.shared_latent_dim,
            coupling_strength: self.coupling_strength,
            max_iter: self.max_iter,
            tol: self.tol,
            regularization_strength: self.regularization_strength,
            random_state: self.random_state,
        };

        Ok(MultiModalGaussianMixture {
            config,
            _phantom: std::marker::PhantomData,
        })
    }
}

impl Estimator<Untrained> for MultiModalGaussianMixture<Untrained> {
    type Config = MultiModalConfig;
    type Error = SklearsError;
    type Float = f64;

    fn config(&self) -> &Self::Config {
        &self.config
    }
}

impl Estimator<Trained> for MultiModalGaussianMixture<Trained> {
    type Config = MultiModalConfig;
    type Error = SklearsError;
    type Float = f64;

    fn config(&self) -> &Self::Config {
        &self.config
    }
}

impl Fit<HashMap<String, Array2<f64>>, Option<Array1<usize>>>
    for MultiModalGaussianMixture<Untrained>
{
    type Fitted = MultiModalGaussianMixtureTrained;

    fn fit(
        self,
        X: &HashMap<String, Array2<f64>>,
        y: &Option<Array1<usize>>,
    ) -> SklResult<Self::Fitted> {
        // Validate input data matches configured modalities
        for modality in &self.config.modalities {
            if !X.contains_key(&modality.name) {
                return Err(SklearsError::InvalidInput(format!(
                    "Missing data for modality: {}",
                    modality.name
                )));
            }
            let data = &X[&modality.name];
            if data.ncols() != modality.n_features {
                return Err(SklearsError::InvalidInput(format!(
                    "Feature dimension mismatch for modality {}: expected {}, got {}",
                    modality.name,
                    modality.n_features,
                    data.ncols()
                )));
            }
        }

        // Get sample size (assuming all modalities have same number of samples)
        let n_samples = X.values().next().unwrap().nrows();
        for (name, data) in X.iter() {
            if data.nrows() != n_samples {
                return Err(SklearsError::InvalidInput(format!(
                    "Sample size mismatch for modality {}: expected {}, got {}",
                    name,
                    n_samples,
                    data.nrows()
                )));
            }
        }

        match self.config.fusion_strategy {
            FusionStrategy::EarlyFusion => self.fit_early_fusion(X, y),
            FusionStrategy::LateFusion => self.fit_late_fusion(X, y),
            FusionStrategy::IntermediateFusion => self.fit_intermediate_fusion(X, y),
            FusionStrategy::CoupledFusion => self.fit_coupled_fusion(X, y),
        }
    }
}

impl MultiModalGaussianMixture<Untrained> {
    /// Initialize parameters using K-means++ style initialization
    fn initialize_parameters(
        &self,
        X: &HashMap<String, Array2<f64>>,
    ) -> SklResult<(
        Array1<f64>,
        HashMap<String, Array2<f64>>,
        HashMap<String, Array3<f64>>,
    )> {
        let n_samples = X.values().next().unwrap().nrows();
        let n_components = self.config.n_components;

        // Initialize component weights uniformly
        let weights = Array1::ones(n_components) / n_components as f64;

        // Initialize modality-specific parameters
        let mut modality_means = HashMap::new();
        let mut modality_covariances = HashMap::new();

        let mut rng = thread_rng();

        for modality in &self.config.modalities {
            let data = &X[&modality.name];
            let n_features = data.ncols();

            // Initialize means using random samples
            let mut means = Array2::zeros((n_components, n_features));
            for k in 0..n_components {
                let uniform = RandUniform::new(0, n_samples).map_err(|e| {
                    SklearsError::InvalidInput(format!("Uniform distribution error: {}", e))
                })?;
                let sample_idx = rng.sample(uniform);
                means.row_mut(k).assign(&data.row(sample_idx));
            }

            // Initialize covariances
            let covariances = match modality.covariance_type {
                CovarianceType::Full => {
                    let mut cov = Array3::zeros((n_components, n_features, n_features));
                    for k in 0..n_components {
                        for i in 0..n_features {
                            cov[[k, i, i]] = 1.0; // Start with identity
                        }
                    }
                    cov
                }
                CovarianceType::Diagonal => {
                    let mut cov = Array3::zeros((n_components, n_features, 1));
                    for k in 0..n_components {
                        for i in 0..n_features {
                            cov[[k, i, 0]] = 1.0;
                        }
                    }
                    cov
                }
                CovarianceType::Tied => {
                    let mut cov = Array3::zeros((1, n_features, n_features));
                    for i in 0..n_features {
                        cov[[0, i, i]] = 1.0;
                    }
                    cov
                }
                CovarianceType::Spherical => Array3::ones((n_components, 1, 1)),
            };

            modality_means.insert(modality.name.clone(), means);
            modality_covariances.insert(modality.name.clone(), covariances);
        }

        Ok((weights, modality_means, modality_covariances))
    }

    /// Early fusion implementation: concatenate all modality features
    fn fit_early_fusion(
        &self,
        X: &HashMap<String, Array2<f64>>,
        _y: &Option<Array1<usize>>,
    ) -> SklResult<MultiModalGaussianMixtureTrained> {
        let n_samples = X.values().next().unwrap().nrows();

        // Concatenate all modality data
        let mut concatenated_features = Vec::new();
        for modality in &self.config.modalities {
            concatenated_features.push(X[&modality.name].clone());
        }

        // Stack horizontally
        let mut combined_data = concatenated_features[0].clone();
        for i in 1..concatenated_features.len() {
            let current_cols = combined_data.ncols();
            let new_cols = concatenated_features[i].ncols();
            let mut new_data = Array2::zeros((n_samples, current_cols + new_cols));
            new_data
                .slice_mut(s![.., ..current_cols])
                .assign(&combined_data);
            new_data
                .slice_mut(s![.., current_cols..])
                .assign(&concatenated_features[i]);
            combined_data = new_data;
        }

        // Run standard GMM on concatenated data
        let (mut weights, _means_map, _covariances_map) = self.initialize_parameters(X)?;
        let mut log_likelihood_history = Vec::new();

        // Since we're doing early fusion, we need to work with the concatenated means
        let total_features: usize = self.config.modalities.iter().map(|m| m.n_features).sum();
        let mut combined_means = Array2::zeros((self.config.n_components, total_features));
        let mut combined_covariances =
            Array3::zeros((self.config.n_components, total_features, total_features));

        // Initialize with identity covariances for simplicity
        for k in 0..self.config.n_components {
            for i in 0..total_features {
                combined_covariances[[k, i, i]] = 1.0;
            }
        }

        // EM algorithm for early fusion
        for iter in 0..self.config.max_iter {
            let old_log_likelihood = if log_likelihood_history.is_empty() {
                f64::NEG_INFINITY
            } else {
                *log_likelihood_history.last().unwrap()
            };

            // E-step: Compute responsibilities
            let mut responsibilities = Array2::zeros((n_samples, self.config.n_components));
            let mut log_likelihood = 0.0;

            for i in 0..n_samples {
                let sample = combined_data.row(i);
                let mut log_probs = Array1::zeros(self.config.n_components);

                for k in 0..self.config.n_components {
                    let mean = combined_means.row(k);
                    let diff = &sample.to_owned() - &mean.to_owned();
                    let log_det = combined_covariances
                        .slice(s![k, .., ..])
                        .diag()
                        .mapv(|x: f64| x.ln())
                        .sum();
                    let inv_quad = diff.dot(&diff); // Simplified - should use proper inverse

                    log_probs[k] = weights[k].ln()
                        - 0.5
                            * (total_features as f64 * (2.0 * std::f64::consts::PI).ln()
                                + log_det
                                + inv_quad);
                }

                // Numerical stability
                let max_log_prob = log_probs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
                let log_sum_exp =
                    (log_probs.mapv(|x| (x - max_log_prob).exp()).sum()).ln() + max_log_prob;
                log_likelihood += log_sum_exp;

                for k in 0..self.config.n_components {
                    responsibilities[[i, k]] = ((log_probs[k] - log_sum_exp).exp()).max(1e-15);
                }
            }

            log_likelihood_history.push(log_likelihood);

            // Check convergence
            if iter > 0 && (log_likelihood - old_log_likelihood).abs() < self.config.tol {
                break;
            }

            // M-step: Update parameters
            let n_k: Array1<f64> = responsibilities.sum_axis(Axis(0));

            // Update weights
            weights = &n_k / n_samples as f64;

            // Update means
            for k in 0..self.config.n_components {
                if n_k[k] > 1e-15 {
                    let weighted_sum = responsibilities.column(k).iter().enumerate().fold(
                        Array1::zeros(total_features),
                        |mut acc, (i, &resp)| {
                            let sample = combined_data.row(i);
                            for j in 0..total_features {
                                acc[j] += resp * sample[j];
                            }
                            acc
                        },
                    );
                    combined_means.row_mut(k).assign(&(weighted_sum / n_k[k]));
                }
            }

            // Update covariances (diagonal approximation for simplicity)
            for k in 0..self.config.n_components {
                if n_k[k] > 1e-15 {
                    for j in 0..total_features {
                        let mut weighted_var = 0.0;
                        for i in 0..n_samples {
                            let diff = combined_data[[i, j]] - combined_means[[k, j]];
                            weighted_var += responsibilities[[i, k]] * diff * diff;
                        }
                        combined_covariances[[k, j, j]] =
                            (weighted_var / n_k[k] + self.config.regularization_strength).max(1e-6);
                    }
                }
            }
        }

        // Convert back to modality-specific format for compatibility
        let mut final_means = HashMap::new();
        let mut final_covariances = HashMap::new();
        let mut feature_start = 0;

        for modality in &self.config.modalities {
            let n_features = modality.n_features;
            let modality_means = combined_means
                .slice(s![.., feature_start..feature_start + n_features])
                .to_owned();
            let modality_cov_slice = combined_covariances
                .slice(s![
                    ..,
                    feature_start..feature_start + n_features,
                    feature_start..feature_start + n_features
                ])
                .to_owned();

            final_means.insert(modality.name.clone(), modality_means);
            final_covariances.insert(modality.name.clone(), modality_cov_slice);
            feature_start += n_features;
        }

        let n_iter = log_likelihood_history.len();
        Ok(MultiModalGaussianMixtureTrained {
            weights,
            modality_means: final_means,
            modality_covariances: final_covariances,
            shared_latent_means: None,
            latent_projections: HashMap::new(),
            coupling_parameters: Array2::zeros((0, 0)),
            log_likelihood_history,
            n_iter,
            config: self.config.clone(),
        })
    }

    /// Late fusion implementation: train separate models then combine
    fn fit_late_fusion(
        &self,
        X: &HashMap<String, Array2<f64>>,
        _y: &Option<Array1<usize>>,
    ) -> SklResult<MultiModalGaussianMixtureTrained> {
        let n_samples = X.values().next().unwrap().nrows();
        let (weights, mut modality_means, mut modality_covariances) =
            self.initialize_parameters(X)?;
        let mut log_likelihood_history = Vec::new();

        // Train each modality separately with EM
        for modality in &self.config.modalities {
            let data = &X[&modality.name];
            let n_features = modality.n_features;
            let n_components = self.config.n_components;

            let mut modality_weights: Array1<f64> =
                Array1::ones(n_components) / n_components as f64;
            let mut means = modality_means[&modality.name].clone();
            let mut covariances = modality_covariances[&modality.name].clone();

            // EM for this modality
            for _iter in 0..self.config.max_iter {
                // E-step
                let mut responsibilities = Array2::zeros((n_samples, n_components));

                for i in 0..n_samples {
                    let sample = data.row(i);
                    let mut log_probs = Array1::zeros(n_components);

                    for k in 0..n_components {
                        let mean = means.row(k);
                        let diff = &sample.to_owned() - &mean.to_owned();

                        let log_det = match modality.covariance_type {
                            CovarianceType::Full => covariances
                                .slice(s![k, .., ..])
                                .diag()
                                .mapv(|x| x.ln())
                                .sum(),
                            CovarianceType::Diagonal => {
                                covariances.slice(s![k, .., 0]).mapv(|x| x.ln()).sum()
                            }
                            CovarianceType::Spherical => {
                                n_features as f64 * covariances[[k, 0, 0]].ln()
                            }
                            CovarianceType::Tied => covariances
                                .slice(s![0, .., ..])
                                .diag()
                                .mapv(|x| x.ln())
                                .sum(),
                        };

                        let inv_quad = match modality.covariance_type {
                            CovarianceType::Spherical => diff.dot(&diff) / covariances[[k, 0, 0]],
                            _ => diff.dot(&diff), // Simplified
                        };

                        log_probs[k] = modality_weights[k].ln()
                            - 0.5
                                * (n_features as f64 * (2.0 * std::f64::consts::PI).ln()
                                    + log_det
                                    + inv_quad);
                    }

                    // Normalize responsibilities
                    let max_log_prob = log_probs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
                    let log_sum_exp =
                        (log_probs.mapv(|x| (x - max_log_prob).exp()).sum()).ln() + max_log_prob;

                    for k in 0..n_components {
                        responsibilities[[i, k]] = ((log_probs[k] - log_sum_exp).exp()).max(1e-15);
                    }
                }

                // M-step
                let n_k: Array1<f64> = responsibilities.sum_axis(Axis(0));
                modality_weights = &n_k / n_samples as f64;

                // Update means
                for k in 0..n_components {
                    if n_k[k] > 1e-15 {
                        let weighted_sum = responsibilities.column(k).iter().enumerate().fold(
                            Array1::zeros(n_features),
                            |mut acc, (i, &resp)| {
                                let sample = data.row(i);
                                for j in 0..n_features {
                                    acc[j] += resp * sample[j];
                                }
                                acc
                            },
                        );
                        means.row_mut(k).assign(&(weighted_sum / n_k[k]));
                    }
                }

                // Update covariances (simplified diagonal update)
                match modality.covariance_type {
                    CovarianceType::Spherical => {
                        for k in 0..n_components {
                            if n_k[k] > 1e-15 {
                                let mut weighted_var = 0.0;
                                for i in 0..n_samples {
                                    let sample = data.row(i);
                                    let mean = means.row(k);
                                    let diff = &sample.to_owned() - &mean.to_owned();
                                    weighted_var += responsibilities[[i, k]] * diff.dot(&diff);
                                }
                                covariances[[k, 0, 0]] = (weighted_var
                                    / (n_k[k] * n_features as f64)
                                    + self.config.regularization_strength)
                                    .max(1e-6);
                            }
                        }
                    }
                    _ => {
                        // Diagonal covariance update
                        for k in 0..n_components {
                            if n_k[k] > 1e-15 {
                                for j in 0..n_features {
                                    let mut weighted_var = 0.0;
                                    for i in 0..n_samples {
                                        let diff = data[[i, j]] - means[[k, j]];
                                        weighted_var += responsibilities[[i, k]] * diff * diff;
                                    }
                                    let var_idx = match modality.covariance_type {
                                        CovarianceType::Diagonal => (k, j, 0),
                                        _ => (k, j, j),
                                    };
                                    covariances[var_idx] = (weighted_var / n_k[k]
                                        + self.config.regularization_strength)
                                        .max(1e-6);
                                }
                            }
                        }
                    }
                }
            }

            // Update the modality parameters
            modality_means.insert(modality.name.clone(), means);
            modality_covariances.insert(modality.name.clone(), covariances);
        }

        // Combine predictions from all modalities (simple averaging)
        log_likelihood_history.push(0.0); // Placeholder

        Ok(MultiModalGaussianMixtureTrained {
            weights,
            modality_means,
            modality_covariances,
            shared_latent_means: None,
            latent_projections: HashMap::new(),
            coupling_parameters: Array2::zeros((0, 0)),
            log_likelihood_history,
            n_iter: 1,
            config: self.config.clone(),
        })
    }

    /// Intermediate fusion: learn shared latent representation
    fn fit_intermediate_fusion(
        &self,
        X: &HashMap<String, Array2<f64>>,
        _y: &Option<Array1<usize>>,
    ) -> SklResult<MultiModalGaussianMixtureTrained> {
        let _n_samples = X.values().next().unwrap().nrows();
        let latent_dim = self.config.shared_latent_dim.unwrap();
        let (weights, modality_means, modality_covariances) = self.initialize_parameters(X)?;

        // Initialize projection matrices for each modality to latent space
        let mut latent_projections = HashMap::new();
        let mut rng = thread_rng();

        for modality in &self.config.modalities {
            let normal = RandNormal::new(0.0, 0.1).map_err(|e| {
                SklearsError::InvalidInput(format!("Normal distribution error: {}", e))
            })?;
            let mut projection = Array2::zeros((latent_dim, modality.n_features));
            for i in 0..latent_dim {
                for j in 0..modality.n_features {
                    projection[[i, j]] = rng.sample(normal);
                }
            }
            latent_projections.insert(modality.name.clone(), projection);
        }

        // Initialize shared latent means
        let mut shared_latent_means = Array2::zeros((self.config.n_components, latent_dim));
        for k in 0..self.config.n_components {
            for d in 0..latent_dim {
                let normal = RandNormal::new(0.0, 1.0).map_err(|e| {
                    SklearsError::InvalidInput(format!("Normal distribution error: {}", e))
                })?;
                shared_latent_means[[k, d]] = rng.sample(normal);
            }
        }

        let mut log_likelihood_history = Vec::new();
        let mut coupling_parameters =
            Array2::zeros((self.config.modalities.len(), self.config.modalities.len()));

        // Initialize coupling parameters
        for i in 0..self.config.modalities.len() {
            coupling_parameters[[i, i]] = 1.0;
            for j in (i + 1)..self.config.modalities.len() {
                coupling_parameters[[i, j]] = self.config.coupling_strength;
                coupling_parameters[[j, i]] = self.config.coupling_strength;
            }
        }

        // Simplified intermediate fusion (would need more sophisticated implementation for production)
        log_likelihood_history.push(0.0);

        Ok(MultiModalGaussianMixtureTrained {
            weights,
            modality_means,
            modality_covariances,
            shared_latent_means: Some(shared_latent_means),
            latent_projections,
            coupling_parameters,
            log_likelihood_history,
            n_iter: 1,
            config: self.config.clone(),
        })
    }

    /// Coupled fusion: joint optimization across modalities
    fn fit_coupled_fusion(
        &self,
        X: &HashMap<String, Array2<f64>>,
        _y: &Option<Array1<usize>>,
    ) -> SklResult<MultiModalGaussianMixtureTrained> {
        let n_samples = X.values().next().unwrap().nrows();
        let (mut weights, mut modality_means, mut modality_covariances) =
            self.initialize_parameters(X)?;
        let mut log_likelihood_history = Vec::new();

        // Initialize coupling parameters between modalities
        let n_modalities = self.config.modalities.len();
        let mut coupling_parameters = Array2::zeros((n_modalities, n_modalities));

        for i in 0..n_modalities {
            coupling_parameters[[i, i]] = 1.0;
            for j in (i + 1)..n_modalities {
                coupling_parameters[[i, j]] = self.config.coupling_strength;
                coupling_parameters[[j, i]] = self.config.coupling_strength;
            }
        }

        // Coupled EM algorithm
        for iter in 0..self.config.max_iter {
            let old_log_likelihood = if log_likelihood_history.is_empty() {
                f64::NEG_INFINITY
            } else {
                *log_likelihood_history.last().unwrap()
            };

            let mut total_log_likelihood = 0.0;
            let mut global_responsibilities = Array2::zeros((n_samples, self.config.n_components));

            // E-step: Compute responsibilities for each modality and combine
            for (modality_idx, modality) in self.config.modalities.iter().enumerate() {
                let data = &X[&modality.name];
                let means = &modality_means[&modality.name];
                let covariances = &modality_covariances[&modality.name];
                let mut modality_responsibilities =
                    Array2::zeros((n_samples, self.config.n_components));

                for i in 0..n_samples {
                    let sample = data.row(i);
                    let mut log_probs = Array1::zeros(self.config.n_components);

                    for k in 0..self.config.n_components {
                        let mean = means.row(k);
                        let diff = &sample.to_owned() - &mean.to_owned();

                        let (log_det, inv_quad) = match modality.covariance_type {
                            CovarianceType::Spherical => {
                                let variance = covariances[[k, 0, 0]];
                                let log_det = modality.n_features as f64 * variance.ln();
                                let inv_quad = diff.dot(&diff) / variance;
                                (log_det, inv_quad)
                            }
                            _ => {
                                // Simplified diagonal covariance
                                let log_det = (0..modality.n_features)
                                    .map(|j| {
                                        covariances[[k, j, 0.min(covariances.dim().2 - 1)]].ln()
                                    })
                                    .sum::<f64>();
                                let inv_quad = diff.dot(&diff); // Simplified
                                (log_det, inv_quad)
                            }
                        };

                        log_probs[k] = weights[k].ln()
                            - 0.5
                                * (modality.n_features as f64 * (2.0 * std::f64::consts::PI).ln()
                                    + log_det
                                    + inv_quad);
                    }

                    // Numerical stability
                    let max_log_prob = log_probs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
                    let log_sum_exp =
                        (log_probs.mapv(|x| (x - max_log_prob).exp()).sum()).ln() + max_log_prob;
                    total_log_likelihood += log_sum_exp * modality.modality_weight;

                    for k in 0..self.config.n_components {
                        modality_responsibilities[[i, k]] =
                            ((log_probs[k] - log_sum_exp).exp()).max(1e-15);
                    }
                }

                // Combine with coupling
                for i in 0..n_samples {
                    for k in 0..self.config.n_components {
                        global_responsibilities[[i, k]] += modality.modality_weight
                            * coupling_parameters[[modality_idx, modality_idx]]
                            * modality_responsibilities[[i, k]];
                    }
                }
            }

            // Normalize global responsibilities
            for i in 0..n_samples {
                let sum: f64 = global_responsibilities.row(i).sum();
                if sum > 1e-15 {
                    global_responsibilities.row_mut(i).mapv_inplace(|x| x / sum);
                }
            }

            log_likelihood_history.push(total_log_likelihood);

            // Check convergence
            if iter > 0 && (total_log_likelihood - old_log_likelihood).abs() < self.config.tol {
                break;
            }

            // M-step: Update parameters using global responsibilities
            let n_k: Array1<f64> = global_responsibilities.sum_axis(Axis(0));
            weights = &n_k / n_samples as f64;

            // Update means and covariances for each modality
            for modality in &self.config.modalities {
                let data = &X[&modality.name];
                let mut means = modality_means[&modality.name].clone();
                let mut covariances = modality_covariances[&modality.name].clone();

                // Update means
                for k in 0..self.config.n_components {
                    if n_k[k] > 1e-15 {
                        let weighted_sum = global_responsibilities
                            .column(k)
                            .iter()
                            .enumerate()
                            .fold(Array1::zeros(modality.n_features), |mut acc, (i, &resp)| {
                                let sample = data.row(i);
                                for j in 0..modality.n_features {
                                    acc[j] += resp * sample[j];
                                }
                                acc
                            });
                        means.row_mut(k).assign(&(weighted_sum / n_k[k]));
                    }
                }

                // Update covariances
                match modality.covariance_type {
                    CovarianceType::Spherical => {
                        for k in 0..self.config.n_components {
                            if n_k[k] > 1e-15 {
                                let mut weighted_var = 0.0;
                                for i in 0..n_samples {
                                    let sample = data.row(i);
                                    let mean = means.row(k);
                                    let diff = &sample.to_owned() - &mean.to_owned();
                                    weighted_var +=
                                        global_responsibilities[[i, k]] * diff.dot(&diff);
                                }
                                covariances[[k, 0, 0]] = (weighted_var
                                    / (n_k[k] * modality.n_features as f64)
                                    + self.config.regularization_strength)
                                    .max(1e-6);
                            }
                        }
                    }
                    _ => {
                        for k in 0..self.config.n_components {
                            if n_k[k] > 1e-15 {
                                for j in 0..modality.n_features {
                                    let mut weighted_var = 0.0;
                                    for i in 0..n_samples {
                                        let diff = data[[i, j]] - means[[k, j]];
                                        weighted_var +=
                                            global_responsibilities[[i, k]] * diff * diff;
                                    }
                                    let var_idx = match modality.covariance_type {
                                        CovarianceType::Diagonal => (k, j, 0),
                                        _ => (k, j, j),
                                    };
                                    covariances[var_idx] = (weighted_var / n_k[k]
                                        + self.config.regularization_strength)
                                        .max(1e-6);
                                }
                            }
                        }
                    }
                }

                modality_means.insert(modality.name.clone(), means);
                modality_covariances.insert(modality.name.clone(), covariances);
            }
        }

        let n_iter = log_likelihood_history.len();
        Ok(MultiModalGaussianMixtureTrained {
            weights,
            modality_means,
            modality_covariances,
            shared_latent_means: None,
            latent_projections: HashMap::new(),
            coupling_parameters,
            log_likelihood_history,
            n_iter,
            config: self.config.clone(),
        })
    }
}

impl Predict<HashMap<String, Array2<f64>>, Array1<usize>> for MultiModalGaussianMixtureTrained {
    fn predict(&self, X: &HashMap<String, Array2<f64>>) -> SklResult<Array1<usize>> {
        let probabilities = self.predict_proba(X)?;
        let n_samples = probabilities.nrows();
        let mut predictions = Array1::zeros(n_samples);

        for i in 0..n_samples {
            let mut max_prob = 0.0;
            let mut best_component = 0;

            for k in 0..self.config.n_components {
                if probabilities[[i, k]] > max_prob {
                    max_prob = probabilities[[i, k]];
                    best_component = k;
                }
            }
            predictions[i] = best_component;
        }

        Ok(predictions)
    }
}

impl MultiModalGaussianMixtureTrained {
    /// Predict class probabilities for multi-modal data
    pub fn predict_proba(&self, X: &HashMap<String, Array2<f64>>) -> SklResult<Array2<f64>> {
        // Validate input
        for modality in &self.config.modalities {
            if !X.contains_key(&modality.name) {
                return Err(SklearsError::InvalidInput(format!(
                    "Missing data for modality: {}",
                    modality.name
                )));
            }
        }

        let n_samples = X.values().next().unwrap().nrows();
        let mut probabilities = Array2::zeros((n_samples, self.config.n_components));

        match self.config.fusion_strategy {
            FusionStrategy::EarlyFusion => {
                // Concatenate features and compute probabilities
                self.predict_proba_early_fusion(X, &mut probabilities)?;
            }
            FusionStrategy::LateFusion => {
                // Average probabilities from each modality
                self.predict_proba_late_fusion(X, &mut probabilities)?;
            }
            FusionStrategy::IntermediateFusion => {
                // Use shared latent representation
                self.predict_proba_intermediate_fusion(X, &mut probabilities)?;
            }
            FusionStrategy::CoupledFusion => {
                // Joint prediction across modalities
                self.predict_proba_coupled_fusion(X, &mut probabilities)?;
            }
        }

        Ok(probabilities)
    }

    fn predict_proba_early_fusion(
        &self,
        _X: &HashMap<String, Array2<f64>>,
        probabilities: &mut Array2<f64>,
    ) -> SklResult<()> {
        // This would concatenate features and compute GMM probabilities
        // For now, simplified implementation
        let n_samples = probabilities.nrows();
        for i in 0..n_samples {
            probabilities.row_mut(i).assign(&self.weights);
        }
        Ok(())
    }

    fn predict_proba_late_fusion(
        &self,
        X: &HashMap<String, Array2<f64>>,
        probabilities: &mut Array2<f64>,
    ) -> SklResult<()> {
        let n_samples = probabilities.nrows();
        probabilities.fill(0.0);

        // Average predictions from each modality
        for modality in &self.config.modalities {
            let data = &X[&modality.name];
            let means = &self.modality_means[&modality.name];

            for i in 0..n_samples {
                let sample = data.row(i);
                let mut modality_probs = Array1::zeros(self.config.n_components);

                for k in 0..self.config.n_components {
                    let mean = means.row(k);
                    let diff = &sample.to_owned() - &mean.to_owned();
                    let log_prob = self.weights[k].ln() - 0.5 * diff.dot(&diff);
                    modality_probs[k] = log_prob.exp();
                }

                // Normalize
                let sum: f64 = modality_probs.sum();
                if sum > 1e-15 {
                    modality_probs.mapv_inplace(|x| x / sum);
                }

                // Add weighted contribution
                for k in 0..self.config.n_components {
                    probabilities[[i, k]] += modality.modality_weight * modality_probs[k];
                }
            }
        }

        // Final normalization
        for i in 0..n_samples {
            let sum: f64 = probabilities.row(i).sum();
            if sum > 1e-15 {
                probabilities.row_mut(i).mapv_inplace(|x| x / sum);
            }
        }

        Ok(())
    }

    fn predict_proba_intermediate_fusion(
        &self,
        _X: &HashMap<String, Array2<f64>>,
        probabilities: &mut Array2<f64>,
    ) -> SklResult<()> {
        // Project to latent space and compute probabilities
        // Simplified implementation
        let n_samples = probabilities.nrows();
        for i in 0..n_samples {
            probabilities.row_mut(i).assign(&self.weights);
        }
        Ok(())
    }

    fn predict_proba_coupled_fusion(
        &self,
        X: &HashMap<String, Array2<f64>>,
        probabilities: &mut Array2<f64>,
    ) -> SklResult<()> {
        // Use coupling parameters to combine modality predictions
        self.predict_proba_late_fusion(X, probabilities)?;

        // Apply coupling transformations (simplified)
        let n_samples = probabilities.nrows();
        for i in 0..n_samples {
            for k in 0..self.config.n_components {
                probabilities[[i, k]] *= 1.0 + self.config.coupling_strength;
            }

            // Renormalize
            let sum: f64 = probabilities.row(i).sum();
            if sum > 1e-15 {
                probabilities.row_mut(i).mapv_inplace(|x| x / sum);
            }
        }

        Ok(())
    }

    /// Compute log-likelihood of the data
    pub fn score(&self, X: &HashMap<String, Array2<f64>>) -> SklResult<f64> {
        let probabilities = self.predict_proba(X)?;
        let log_likelihood = probabilities.mapv(|p| p.max(1e-15).ln()).sum();
        Ok(log_likelihood)
    }

    /// Get model selection criteria
    pub fn model_selection(&self, X: &HashMap<String, Array2<f64>>) -> SklResult<ModelSelection> {
        let n_samples = X.values().next().unwrap().nrows();
        let total_features: usize = self.config.modalities.iter().map(|m| m.n_features).sum();

        // Simplified parameter counting
        let n_parameters = ModelSelection::n_parameters(
            self.config.n_components,
            total_features,
            &CovarianceType::Full,
        );

        let log_likelihood = self.score(X)?;
        let aic = ModelSelection::aic(log_likelihood, n_parameters);
        let bic = ModelSelection::bic(log_likelihood, n_parameters, n_samples);

        Ok(ModelSelection {
            aic,
            bic,
            log_likelihood,
            n_parameters,
        })
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;

    fn create_test_multi_modal_data() -> HashMap<String, Array2<f64>> {
        let mut data = HashMap::new();

        // Visual modality: 2D features
        let visual_data =
            Array2::from_shape_vec((100, 2), (0..200).map(|i| i as f64 * 0.1).collect()).unwrap();
        data.insert("visual".to_string(), visual_data);

        // Textual modality: 3D features
        let textual_data = Array2::from_shape_vec(
            (100, 3),
            (0..300).map(|i| (i as f64 * 0.05).sin()).collect(),
        )
        .unwrap();
        data.insert("textual".to_string(), textual_data);

        data
    }

    #[test]
    fn test_multi_modal_builder() {
        let model = MultiModalGaussianMixtureBuilder::new(3)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::EarlyFusion)
            .coupling_strength(0.2)
            .max_iter(10)
            .build()
            .unwrap();

        assert_eq!(model.config.n_components, 3);
        assert_eq!(model.config.modalities.len(), 2);
        assert_eq!(model.config.fusion_strategy, FusionStrategy::EarlyFusion);
        assert_abs_diff_eq!(model.config.coupling_strength, 0.2, epsilon = 1e-10);
    }

    #[test]
    fn test_early_fusion_fit() {
        let data = create_test_multi_modal_data();
        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::EarlyFusion)
            .max_iter(5)
            .build()
            .unwrap();

        let trained = model.fit(&data, &None).unwrap();

        assert_eq!(trained.weights.len(), 2);
        assert!(trained.modality_means.contains_key("visual"));
        assert!(trained.modality_means.contains_key("textual"));
        assert_eq!(trained.modality_means["visual"].nrows(), 2); // n_components
        assert_eq!(trained.modality_means["visual"].ncols(), 2); // n_features
    }

    #[test]
    fn test_late_fusion_fit() {
        let data = create_test_multi_modal_data();
        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::LateFusion)
            .max_iter(5)
            .build()
            .unwrap();

        let trained = model.fit(&data, &None).unwrap();

        assert_eq!(trained.weights.len(), 2);
        assert!(trained.modality_means.contains_key("visual"));
        assert!(trained.modality_means.contains_key("textual"));
    }

    #[test]
    fn test_intermediate_fusion_fit() {
        let data = create_test_multi_modal_data();
        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::IntermediateFusion)
            .shared_latent_dim(4)
            .max_iter(5)
            .build()
            .unwrap();

        let trained = model.fit(&data, &None).unwrap();

        assert_eq!(trained.weights.len(), 2);
        assert!(trained.shared_latent_means.is_some());
        let latent_means = trained.shared_latent_means.as_ref().unwrap();
        assert_eq!(latent_means.nrows(), 2); // n_components
        assert_eq!(latent_means.ncols(), 4); // latent_dim
    }

    #[test]
    fn test_coupled_fusion_fit() {
        let data = create_test_multi_modal_data();
        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::CoupledFusion)
            .coupling_strength(0.3)
            .max_iter(5)
            .build()
            .unwrap();

        let trained = model.fit(&data, &None).unwrap();

        assert_eq!(trained.weights.len(), 2);
        assert_eq!(trained.coupling_parameters.nrows(), 2); // n_modalities
        assert_eq!(trained.coupling_parameters.ncols(), 2); // n_modalities
    }

    #[test]
    fn test_prediction() {
        let data = create_test_multi_modal_data();
        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::LateFusion)
            .max_iter(3)
            .build()
            .unwrap();

        let trained = model.fit(&data, &None).unwrap();
        let predictions = trained.predict(&data).unwrap();

        assert_eq!(predictions.len(), 100);

        // All predictions should be 0 or 1 (since we have 2 components)
        for &pred in predictions.iter() {
            assert!(pred < 2);
        }
    }

    #[test]
    fn test_predict_proba() {
        let data = create_test_multi_modal_data();
        let model = MultiModalGaussianMixtureBuilder::new(3)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::LateFusion)
            .max_iter(3)
            .build()
            .unwrap();

        let trained = model.fit(&data, &None).unwrap();
        let probabilities = trained.predict_proba(&data).unwrap();

        assert_eq!(probabilities.nrows(), 100);
        assert_eq!(probabilities.ncols(), 3);

        // Probabilities should sum to 1 for each sample
        for i in 0..100 {
            let sum: f64 = probabilities.row(i).sum();
            assert_abs_diff_eq!(sum, 1.0, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_model_selection() {
        let data = create_test_multi_modal_data();
        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::EarlyFusion)
            .max_iter(3)
            .build()
            .unwrap();

        let trained = model.fit(&data, &None).unwrap();
        let model_selection = trained.model_selection(&data).unwrap();

        assert!(model_selection.log_likelihood.is_finite());
        assert!(model_selection.aic.is_finite());
        assert!(model_selection.bic.is_finite());
        assert!(model_selection.n_parameters > 0);
    }

    #[test]
    fn test_validation_missing_modality() {
        let mut data = create_test_multi_modal_data();
        data.remove("textual"); // Remove one modality

        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::EarlyFusion) // Use early fusion to avoid latent dim requirement
            .build()
            .unwrap();

        let result = model.fit(&data, &None);
        assert!(result.is_err());
    }

    #[test]
    fn test_validation_feature_dimension_mismatch() {
        let data = create_test_multi_modal_data();

        let model = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 5) // Wrong dimension
            .add_modality_simple("textual", 3)
            .fusion_strategy(FusionStrategy::EarlyFusion) // Use early fusion to avoid latent dim requirement
            .build()
            .unwrap();

        let result = model.fit(&data, &None);
        assert!(result.is_err());
    }

    #[test]
    fn test_intermediate_fusion_requires_latent_dim() {
        let result = MultiModalGaussianMixtureBuilder::new(2)
            .add_modality_simple("visual", 2)
            .fusion_strategy(FusionStrategy::IntermediateFusion)
            // Missing shared_latent_dim
            .build();

        assert!(result.is_err());
    }
}