sklears_mixture/

nonparametric.rs

//! Nonparametric Mixture Models
//!
//! This module provides nonparametric Bayesian mixture models that automatically
//! determine the number of components from the data. It includes implementations
//! of the Chinese Restaurant Process, Dirichlet Process Gaussian Mixture Model,
//! Pitman-Yor Process, and Hierarchical Dirichlet Process.

use scirs2_core::ndarray::{s, Array1, Array2, ArrayView1, ArrayView2, Axis};
use scirs2_core::random::{Rng, SeedableRng};
use sklears_core::{
    error::{Result as SklResult, SklearsError},
    traits::{Estimator, Fit, Predict, Untrained},
    types::Float,
};
use std::f64::consts::PI;

use crate::common::CovarianceType;

/// Utility function for log-sum-exp
fn log_sum_exp(a: f64, b: f64) -> f64 {
    let max_val = a.max(b);
    if max_val.is_finite() {
        max_val + ((a - max_val).exp() + (b - max_val).exp()).ln()
    } else {
        max_val
    }
}

/// Chinese Restaurant Process Mixture Model
///
/// A nonparametric Bayesian mixture model that automatically determines the number of components
/// using the Chinese Restaurant Process metaphor. Customers (data points) sit at tables (components)
/// with probability proportional to the number of existing customers at each table, or start a new
/// table with probability proportional to the concentration parameter α.
///
/// # Examples
///
/// ```
/// use sklears_mixture::{ChineseRestaurantProcess, CovarianceType};
/// use sklears_core::traits::{Predict, Fit};
/// use scirs2_core::ndarray::array;
///
/// let X = array![[0.0, 0.0], [1.0, 1.0], [2.0, 2.0], [10.0, 10.0], [11.0, 11.0], [12.0, 12.0]];
///
/// let model = ChineseRestaurantProcess::new()
///     .alpha(1.0)
///     .max_components(10)
///     .covariance_type(CovarianceType::Diagonal);
/// let fitted = model.fit(&X.view(), &()).unwrap();
/// let labels = fitted.predict(&X.view()).unwrap();
/// ```
#[derive(Debug, Clone)]
pub struct ChineseRestaurantProcess<S = Untrained> {
    state: S,
    /// Concentration parameter (higher values favor more clusters)
    alpha: f64,
    /// Maximum number of components allowed
    max_components: usize,
    /// Covariance type for the Gaussian base distribution
    covariance_type: CovarianceType,
    /// Convergence tolerance
    tol: f64,
    /// Maximum number of iterations
    max_iter: usize,
    /// Random state for reproducibility
    random_state: Option<u64>,
    /// Regularization parameter for covariance matrices
    reg_covar: f64,
}

impl ChineseRestaurantProcess<Untrained> {
    /// Create a new Chinese Restaurant Process
    pub fn new() -> Self {
        Self {
            state: Untrained,
            alpha: 1.0,
            max_components: 20,
            covariance_type: CovarianceType::Full,
            tol: 1e-3,
            max_iter: 100,
            random_state: None,
            reg_covar: 1e-6,
        }
    }

    /// Set the concentration parameter
    pub fn alpha(mut self, alpha: f64) -> Self {
        self.alpha = alpha;
        self
    }

    /// Set the maximum number of components
    pub fn max_components(mut self, max_components: usize) -> Self {
        self.max_components = max_components;
        self
    }

    /// Set the covariance type
    pub fn covariance_type(mut self, covariance_type: CovarianceType) -> Self {
        self.covariance_type = covariance_type;
        self
    }

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

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

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

    /// Set the regularization parameter
    pub fn reg_covar(mut self, reg_covar: f64) -> Self {
        self.reg_covar = reg_covar;
        self
    }

    /// Initialize tables and assignments using random clustering
    fn initialize_tables(
        &self,
        X: &Array2<f64>,
        rng: &mut scirs2_core::random::rngs::StdRng,
    ) -> SklResult<(Array1<usize>, Array1<usize>, usize)> {
        let n_samples = X.nrows();
        let mut table_assignments = Array1::zeros(n_samples);
        let mut table_counts: Array1<usize> = Array1::zeros(self.max_components);
        let mut n_tables = 1;

        // Assign first customer to first table
        table_assignments[0] = 0;
        table_counts[0] = 1;

        // Process remaining customers
        for i in 1..n_samples {
            let total_customers = i;
            let mut probabilities = Array1::zeros(n_tables + 1);

            // Probability of sitting at existing tables
            for k in 0..n_tables {
                probabilities[k] = table_counts[k] as f64 / (total_customers as f64 + self.alpha);
            }

            // Probability of creating new table
            if n_tables < self.max_components {
                probabilities[n_tables] = self.alpha / (total_customers as f64 + self.alpha);
            }

            // Sample table assignment
            let cumsum: f64 = probabilities.iter().sum();
            let mut cumulative = 0.0;
            let target = rng.gen::<f64>() * cumsum;

            let mut chosen_table = 0;
            for k in 0..=n_tables {
                cumulative += probabilities[k];
                if target <= cumulative {
                    chosen_table = k;
                    break;
                }
            }

            // Update assignments and counts
            table_assignments[i] = chosen_table;
            table_counts[chosen_table] += 1;

            // If new table was chosen, increment table count
            if chosen_table == n_tables && n_tables < self.max_components {
                n_tables += 1;
            }
        }

        Ok((table_assignments, table_counts, n_tables))
    }

    /// Compute component parameters from assignments
    fn compute_parameters(
        &self,
        X: &Array2<f64>,
        table_assignments: &Array1<usize>,
        n_tables: usize,
    ) -> SklResult<(Array1<f64>, Array2<f64>, Vec<Array2<f64>>)> {
        let (n_samples, n_features) = X.dim();

        // Compute weights (normalized table sizes)
        let mut weights = Array1::zeros(n_tables);
        for &assignment in table_assignments.iter() {
            if assignment < n_tables {
                weights[assignment] += 1.0;
            }
        }
        weights /= n_samples as f64;

        // Compute means
        let mut means = Array2::zeros((n_tables, n_features));
        let mut counts: Array1<f64> = Array1::zeros(n_tables);

        for (i, &assignment) in table_assignments.iter().enumerate() {
            if assignment < n_tables {
                for j in 0..n_features {
                    means[[assignment, j]] += X[[i, j]];
                }
                counts[assignment] += 1.0;
            }
        }

        // Normalize means
        for k in 0..n_tables {
            if counts[k] > 0.0 {
                for j in 0..n_features {
                    means[[k, j]] /= counts[k];
                }
            }
        }

        // Compute covariances
        let mut covariances = Vec::new();
        for k in 0..n_tables {
            let mut cov = Array2::zeros((n_features, n_features));
            let mut count = 0.0;

            for (i, &assignment) in table_assignments.iter().enumerate() {
                if assignment == k {
                    let diff = &X.row(i) - &means.row(k);
                    for j in 0..n_features {
                        for l in 0..n_features {
                            cov[[j, l]] += diff[j] * diff[l];
                        }
                    }
                    count += 1.0;
                }
            }

            if count > 1.0 {
                cov /= count - 1.0;
            } else {
                // Use identity for single-point clusters
                for j in 0..n_features {
                    cov[[j, j]] = 1.0;
                }
            }

            // Apply covariance type constraints and regularization
            cov = self.regularize_covariance(cov)?;
            covariances.push(cov);
        }

        Ok((weights, means, covariances))
    }

    /// Regularize covariance matrix based on covariance type
    fn regularize_covariance(&self, mut cov: Array2<f64>) -> SklResult<Array2<f64>> {
        let n_features = cov.dim().0;

        match self.covariance_type {
            CovarianceType::Full => {
                // Add regularization to diagonal
                for i in 0..n_features {
                    cov[[i, i]] += self.reg_covar;
                }
            }
            CovarianceType::Diagonal => {
                // Keep only diagonal elements
                for i in 0..n_features {
                    for j in 0..n_features {
                        if i != j {
                            cov[[i, j]] = 0.0;
                        } else {
                            cov[[i, i]] += self.reg_covar;
                        }
                    }
                }
            }
            CovarianceType::Tied => {
                // Use full covariance (tied across components handled at higher level)
                for i in 0..n_features {
                    cov[[i, i]] += self.reg_covar;
                }
            }
            CovarianceType::Spherical => {
                // Use scalar variance
                let trace = (0..n_features).map(|i| cov[[i, i]]).sum::<f64>() / n_features as f64;
                cov.fill(0.0);
                for i in 0..n_features {
                    cov[[i, i]] = trace + self.reg_covar;
                }
            }
        }

        Ok(cov)
    }

    /// Compute log likelihood
    fn compute_log_likelihood(
        &self,
        X: &Array2<f64>,
        weights: &Array1<f64>,
        means: &Array2<f64>,
        covariances: &[Array2<f64>],
    ) -> SklResult<f64> {
        let n_samples = X.nrows();
        let n_components = weights.len();
        let mut log_likelihood = 0.0;

        for i in 0..n_samples {
            let sample = X.row(i);
            let mut sample_likelihood = 0.0;

            for k in 0..n_components {
                if weights[k] > 1e-10 {
                    let log_pdf =
                        self.multivariate_normal_log_pdf(&sample, &means.row(k), &covariances[k])?;
                    sample_likelihood += weights[k] * log_pdf.exp();
                }
            }

            if sample_likelihood > 1e-300 {
                log_likelihood += sample_likelihood.ln();
            }
        }

        Ok(log_likelihood)
    }

    /// Compute multivariate normal log PDF
    fn multivariate_normal_log_pdf(
        &self,
        x: &ArrayView1<f64>,
        mean: &ArrayView1<f64>,
        cov: &Array2<f64>,
    ) -> SklResult<f64> {
        let d = x.len() as f64;
        let diff = x - mean;

        // Simple determinant and inverse computation for small matrices
        let det = self.matrix_determinant(cov)?;
        if det <= 0.0 {
            return Ok(f64::NEG_INFINITY);
        }

        let inv_cov = self.matrix_inverse(cov)?;
        let mut quad_form = 0.0;
        for i in 0..diff.len() {
            for j in 0..diff.len() {
                quad_form += diff[i] * inv_cov[[i, j]] * diff[j];
            }
        }

        Ok(-0.5 * (d * (2.0 * PI).ln() + det.ln() + quad_form))
    }

    /// Simple matrix determinant calculation
    fn matrix_determinant(&self, A: &Array2<f64>) -> SklResult<f64> {
        let n = A.dim().0;
        if n == 1 {
            return Ok(A[[0, 0]]);
        }
        if n == 2 {
            return Ok(A[[0, 0]] * A[[1, 1]] - A[[0, 1]] * A[[1, 0]]);
        }

        // For larger matrices, use simple LU decomposition
        let mut det = 1.0;
        let mut A_copy = A.clone();

        for i in 0..n {
            let mut max_row = i;
            for k in i + 1..n {
                if A_copy[[k, i]].abs() > A_copy[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            if max_row != i {
                for j in 0..n {
                    let temp = A_copy[[i, j]];
                    A_copy[[i, j]] = A_copy[[max_row, j]];
                    A_copy[[max_row, j]] = temp;
                }
                det *= -1.0;
            }

            if A_copy[[i, i]].abs() < 1e-12 {
                return Ok(0.0);
            }

            det *= A_copy[[i, i]];

            for k in i + 1..n {
                let factor = A_copy[[k, i]] / A_copy[[i, i]];
                for j in i..n {
                    A_copy[[k, j]] -= factor * A_copy[[i, j]];
                }
            }
        }

        Ok(det)
    }

    /// Simple matrix inverse using Gauss-Jordan elimination
    fn matrix_inverse(&self, A: &Array2<f64>) -> SklResult<Array2<f64>> {
        let n = A.dim().0;
        let mut aug = Array2::zeros((n, 2 * n));

        // Create augmented matrix [A | I]
        for i in 0..n {
            for j in 0..n {
                aug[[i, j]] = A[[i, j]];
                aug[[i, j + n]] = if i == j { 1.0 } else { 0.0 };
            }
        }

        // Gauss-Jordan elimination
        for i in 0..n {
            // Find pivot
            let mut max_row = i;
            for k in i + 1..n {
                if aug[[k, i]].abs() > aug[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            // Swap rows
            if max_row != i {
                for j in 0..2 * n {
                    let temp = aug[[i, j]];
                    aug[[i, j]] = aug[[max_row, j]];
                    aug[[max_row, j]] = temp;
                }
            }

            // Check for singular matrix
            if aug[[i, i]].abs() < 1e-12 {
                return Err(SklearsError::NumericalError(
                    "Matrix is singular".to_string(),
                ));
            }

            // Scale pivot row
            let pivot = aug[[i, i]];
            for j in 0..2 * n {
                aug[[i, j]] /= pivot;
            }

            // Eliminate column
            for k in 0..n {
                if k != i {
                    let factor = aug[[k, i]];
                    for j in 0..2 * n {
                        aug[[k, j]] -= factor * aug[[i, j]];
                    }
                }
            }
        }

        // Extract inverse matrix
        let mut inv = Array2::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                inv[[i, j]] = aug[[i, j + n]];
            }
        }

        Ok(inv)
    }
}

impl Default for ChineseRestaurantProcess<Untrained> {
    fn default() -> Self {
        Self::new()
    }
}

impl Estimator for ChineseRestaurantProcess<Untrained> {
    type Config = ();
    type Error = SklearsError;
    type Float = Float;

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

impl Fit<ArrayView2<'_, Float>, ()> for ChineseRestaurantProcess<Untrained> {
    type Fitted = ChineseRestaurantProcess<ChineseRestaurantProcessTrained>;

    #[allow(non_snake_case)]
    fn fit(self, X: &ArrayView2<'_, Float>, _y: &()) -> SklResult<Self::Fitted> {
        let X = X.to_owned();
        let (n_samples, n_features) = X.dim();

        if n_samples == 0 {
            return Err(SklearsError::InvalidInput(
                "No samples provided".to_string(),
            ));
        }

        let mut rng = match self.random_state {
            Some(seed) => scirs2_core::random::rngs::StdRng::seed_from_u64(seed),
            None => scirs2_core::random::rngs::StdRng::seed_from_u64(42),
        };

        // Initialize table assignments
        let (mut table_assignments, mut table_counts, mut n_tables) =
            self.initialize_tables(&X, &mut rng)?;

        let mut prev_log_likelihood = f64::NEG_INFINITY;
        let mut converged = false;
        let mut n_iter = 0;

        // Gibbs sampling iterations
        for iteration in 0..self.max_iter {
            n_iter = iteration + 1;

            // Update table assignments for each customer
            for i in 0..n_samples {
                // Remove customer from current table
                let current_table = table_assignments[i];
                table_counts[current_table] -= 1;

                // If table becomes empty, remove it
                if table_counts[current_table] == 0 && current_table == n_tables - 1 {
                    n_tables -= 1;
                }

                // Compute probabilities for each table
                let mut probabilities = Array1::zeros(n_tables + 1);
                let remaining_customers = n_samples - 1;

                // Existing tables
                for k in 0..n_tables {
                    if table_counts[k] > 0 {
                        probabilities[k] =
                            table_counts[k] as f64 / (remaining_customers as f64 + self.alpha);
                    }
                }

                // New table
                if n_tables < self.max_components {
                    probabilities[n_tables] =
                        self.alpha / (remaining_customers as f64 + self.alpha);
                }

                // Sample new assignment
                let cumsum: f64 = probabilities.iter().sum();
                let mut cumulative = 0.0;
                let target = rng.gen::<f64>() * cumsum;

                let mut new_table = 0;
                for k in 0..=n_tables {
                    cumulative += probabilities[k];
                    if target <= cumulative {
                        new_table = k;
                        break;
                    }
                }

                // Update assignments
                table_assignments[i] = new_table;
                table_counts[new_table] += 1;

                // If new table was chosen, increment table count
                if new_table == n_tables && n_tables < self.max_components {
                    n_tables += 1;
                }
            }

            // Compute parameters and log-likelihood every few iterations
            if iteration % 5 == 0 {
                let (weights, means, covariances) =
                    self.compute_parameters(&X, &table_assignments, n_tables)?;
                let log_likelihood =
                    self.compute_log_likelihood(&X, &weights, &means, &covariances)?;

                if (log_likelihood - prev_log_likelihood).abs() < self.tol {
                    converged = true;
                    break;
                }
                prev_log_likelihood = log_likelihood;
            }
        }

        // Final parameter computation
        let (weights, means, covariances) =
            self.compute_parameters(&X, &table_assignments, n_tables)?;
        let log_likelihood = self.compute_log_likelihood(&X, &weights, &means, &covariances)?;

        Ok(ChineseRestaurantProcess {
            state: ChineseRestaurantProcessTrained {
                n_components: n_tables,
                weights,
                means,
                covariances,
                covariance_type: self.covariance_type.clone(),
                n_features,
                alpha: self.alpha,
                table_assignments,
                table_counts: table_counts.slice(s![..n_tables]).to_owned(),
                log_likelihood,
                n_iter,
                converged,
                reg_covar: self.reg_covar,
            },
            alpha: self.alpha,
            max_components: self.max_components,
            covariance_type: self.covariance_type,
            tol: self.tol,
            max_iter: self.max_iter,
            random_state: self.random_state,
            reg_covar: self.reg_covar,
        })
    }
}

/// Trained Chinese Restaurant Process Mixture Model
#[derive(Debug, Clone)]
pub struct ChineseRestaurantProcessTrained {
    /// Number of active components
    pub n_components: usize,
    /// Mixture weights (table sizes normalized)
    pub weights: Array1<f64>,
    /// Component means
    pub means: Array2<f64>,
    /// Component covariances
    pub covariances: Vec<Array2<f64>>,
    /// Covariance type
    pub covariance_type: CovarianceType,
    /// Number of features
    pub n_features: usize,
    /// Concentration parameter
    pub alpha: f64,
    /// Table assignments for training data
    pub table_assignments: Array1<usize>,
    /// Number of customers at each table
    pub table_counts: Array1<usize>,
    /// Log-likelihood of the model
    pub log_likelihood: f64,
    /// Number of iterations until convergence
    pub n_iter: usize,
    /// Whether the algorithm converged
    pub converged: bool,
    /// Regularization parameter
    pub reg_covar: f64,
}

impl Predict<ArrayView2<'_, Float>, Array1<i32>>
    for ChineseRestaurantProcess<ChineseRestaurantProcessTrained>
{
    #[allow(non_snake_case)]
    fn predict(&self, X: &ArrayView2<'_, Float>) -> SklResult<Array1<i32>> {
        let X = X.to_owned();
        let (n_samples, _) = X.dim();
        let mut predictions = Array1::zeros(n_samples);

        for i in 0..n_samples {
            let sample = X.row(i);
            let mut max_log_prob = f64::NEG_INFINITY;
            let mut best_component = 0;

            for k in 0..self.state.n_components {
                let log_weight = self.state.weights[k].ln();
                let log_pdf = self.multivariate_normal_log_pdf(
                    &sample,
                    &self.state.means.row(k),
                    &self.state.covariances[k],
                )?;
                let log_prob = log_weight + log_pdf;

                if log_prob > max_log_prob {
                    max_log_prob = log_prob;
                    best_component = k;
                }
            }

            predictions[i] = best_component as i32;
        }

        Ok(predictions)
    }
}

impl ChineseRestaurantProcess<ChineseRestaurantProcessTrained> {
    /// Predict class probabilities
    #[allow(non_snake_case)]
    pub fn predict_proba(&self, X: &ArrayView2<'_, Float>) -> SklResult<Array2<f64>> {
        let X = X.to_owned();
        let (n_samples, _) = X.dim();
        let mut probabilities = Array2::zeros((n_samples, self.state.n_components));

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

            for k in 0..self.state.n_components {
                let log_weight = self.state.weights[k].ln();
                let log_pdf = self.multivariate_normal_log_pdf(
                    &sample,
                    &self.state.means.row(k),
                    &self.state.covariances[k],
                )?;
                log_probs[k] = log_weight + log_pdf;
            }

            // Numerically stable normalization
            let max_log_prob = log_probs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
            let log_sum_exp = max_log_prob
                + log_probs
                    .iter()
                    .map(|&lp| (lp - max_log_prob).exp())
                    .sum::<f64>()
                    .ln();

            for k in 0..self.state.n_components {
                probabilities[[i, k]] = (log_probs[k] - log_sum_exp).exp();
            }
        }

        Ok(probabilities)
    }

    /// Score samples using log-likelihood
    #[allow(non_snake_case)]
    pub fn score(&self, X: &ArrayView2<'_, Float>) -> SklResult<f64> {
        let X = X.to_owned();
        self.compute_log_likelihood(
            &X,
            &self.state.weights,
            &self.state.means,
            &self.state.covariances,
        )
    }

    /// Helper method for log PDF computation
    fn multivariate_normal_log_pdf(
        &self,
        x: &ArrayView1<f64>,
        mean: &ArrayView1<f64>,
        cov: &Array2<f64>,
    ) -> SklResult<f64> {
        let d = x.len() as f64;
        let diff = x - mean;

        let det = self.matrix_determinant(cov)?;
        if det <= 0.0 {
            return Ok(f64::NEG_INFINITY);
        }

        let inv_cov = self.matrix_inverse(cov)?;
        let mut quad_form = 0.0;
        for i in 0..diff.len() {
            for j in 0..diff.len() {
                quad_form += diff[i] * inv_cov[[i, j]] * diff[j];
            }
        }

        Ok(-0.5 * (d * (2.0 * PI).ln() + det.ln() + quad_form))
    }

    /// Helper methods (duplicated for the trained state)
    fn matrix_determinant(&self, A: &Array2<f64>) -> SklResult<f64> {
        let n = A.dim().0;
        if n == 1 {
            return Ok(A[[0, 0]]);
        }
        if n == 2 {
            return Ok(A[[0, 0]] * A[[1, 1]] - A[[0, 1]] * A[[1, 0]]);
        }

        let mut det = 1.0;
        let mut A_copy = A.clone();

        for i in 0..n {
            let mut max_row = i;
            for k in i + 1..n {
                if A_copy[[k, i]].abs() > A_copy[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            if max_row != i {
                for j in 0..n {
                    let temp = A_copy[[i, j]];
                    A_copy[[i, j]] = A_copy[[max_row, j]];
                    A_copy[[max_row, j]] = temp;
                }
                det *= -1.0;
            }

            if A_copy[[i, i]].abs() < 1e-12 {
                return Ok(0.0);
            }

            det *= A_copy[[i, i]];

            for k in i + 1..n {
                let factor = A_copy[[k, i]] / A_copy[[i, i]];
                for j in i..n {
                    A_copy[[k, j]] -= factor * A_copy[[i, j]];
                }
            }
        }

        Ok(det)
    }

    fn matrix_inverse(&self, A: &Array2<f64>) -> SklResult<Array2<f64>> {
        let n = A.dim().0;
        let mut aug = Array2::zeros((n, 2 * n));

        for i in 0..n {
            for j in 0..n {
                aug[[i, j]] = A[[i, j]];
                aug[[i, j + n]] = if i == j { 1.0 } else { 0.0 };
            }
        }

        for i in 0..n {
            let mut max_row = i;
            for k in i + 1..n {
                if aug[[k, i]].abs() > aug[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            if max_row != i {
                for j in 0..2 * n {
                    let temp = aug[[i, j]];
                    aug[[i, j]] = aug[[max_row, j]];
                    aug[[max_row, j]] = temp;
                }
            }

            if aug[[i, i]].abs() < 1e-12 {
                return Err(SklearsError::NumericalError(
                    "Matrix is singular".to_string(),
                ));
            }

            let pivot = aug[[i, i]];
            for j in 0..2 * n {
                aug[[i, j]] /= pivot;
            }

            for k in 0..n {
                if k != i {
                    let factor = aug[[k, i]];
                    for j in 0..2 * n {
                        aug[[k, j]] -= factor * aug[[i, j]];
                    }
                }
            }
        }

        let mut inv = Array2::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                inv[[i, j]] = aug[[i, j + n]];
            }
        }

        Ok(inv)
    }

    fn compute_log_likelihood(
        &self,
        X: &Array2<f64>,
        weights: &Array1<f64>,
        means: &Array2<f64>,
        covariances: &[Array2<f64>],
    ) -> SklResult<f64> {
        let n_samples = X.nrows();
        let n_components = weights.len();
        let mut log_likelihood = 0.0;

        for i in 0..n_samples {
            let sample = X.row(i);
            let mut sample_likelihood = 0.0;

            for k in 0..n_components {
                if weights[k] > 1e-10 {
                    let log_pdf =
                        self.multivariate_normal_log_pdf(&sample, &means.row(k), &covariances[k])?;
                    sample_likelihood += weights[k] * log_pdf.exp();
                }
            }

            if sample_likelihood > 1e-300 {
                log_likelihood += sample_likelihood.ln();
            }
        }

        Ok(log_likelihood)
    }
}

/// Dirichlet Process Gaussian Mixture Model
///
/// A nonparametric Bayesian mixture model that uses the Dirichlet process as a prior
/// over the mixture weights. Uses stick-breaking construction and variational inference.
#[derive(Debug, Clone)]
pub struct DirichletProcessGaussianMixture<S = Untrained> {
    state: S,
    /// Concentration parameter of the Dirichlet process
    pub alpha: f64,
    /// Maximum number of components to consider
    pub max_components: usize,
    /// Type of covariance parameters
    pub covariance_type: CovarianceType,
    /// Convergence threshold
    pub tol: f64,
    /// Maximum number of iterations
    pub max_iter: usize,
    /// Regularization added to diagonal of covariance
    pub reg_covar: f64,
    /// Random state for reproducible results
    pub random_state: Option<u64>,
    /// Number of random initializations
    pub n_init: usize,
}

impl DirichletProcessGaussianMixture<Untrained> {
    /// Create a new Dirichlet Process Gaussian Mixture Model
    pub fn new() -> Self {
        Self {
            state: Untrained,
            alpha: 1.0,
            max_components: 20,
            covariance_type: CovarianceType::Full,
            tol: 1e-3,
            max_iter: 100,
            reg_covar: 1e-6,
            random_state: None,
            n_init: 1,
        }
    }

    /// Set the concentration parameter
    pub fn alpha(mut self, alpha: f64) -> Self {
        self.alpha = alpha;
        self
    }

    /// Set the maximum number of components
    pub fn max_components(mut self, max_components: usize) -> Self {
        self.max_components = max_components;
        self
    }

    /// Set the covariance type
    pub fn covariance_type(mut self, covariance_type: CovarianceType) -> Self {
        self.covariance_type = covariance_type;
        self
    }

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

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

    /// Set the regularization parameter
    pub fn reg_covar(mut self, reg_covar: f64) -> Self {
        self.reg_covar = reg_covar;
        self
    }

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

    /// Set the number of initializations
    pub fn n_init(mut self, n_init: usize) -> Self {
        self.n_init = n_init;
        self
    }
}

impl Default for DirichletProcessGaussianMixture<Untrained> {
    fn default() -> Self {
        Self::new()
    }
}

impl Estimator for DirichletProcessGaussianMixture<Untrained> {
    type Config = ();
    type Error = SklearsError;
    type Float = Float;

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

impl Fit<ArrayView2<'_, Float>, ()> for DirichletProcessGaussianMixture<Untrained> {
    type Fitted = DirichletProcessGaussianMixture<DirichletProcessGaussianMixtureTrained>;

    #[allow(non_snake_case)]
    fn fit(self, X: &ArrayView2<'_, Float>, _y: &()) -> SklResult<Self::Fitted> {
        let X = X.to_owned();
        let (n_samples, n_features) = X.dim();

        if n_samples == 0 {
            return Err(SklearsError::InvalidInput(
                "No samples provided".to_string(),
            ));
        }

        // Initialize stick-breaking weights
        let stick_weights = Array1::ones(self.max_components);
        let mut weights = Array1::zeros(self.max_components);

        // Stick-breaking construction
        let mut remaining = 1.0;
        for k in 0..self.max_components {
            weights[k] = remaining / (self.max_components - k) as f64;
            remaining -= weights[k];
        }

        // Simple K-means++ initialization for means
        let mut rng = match self.random_state {
            Some(seed) => scirs2_core::random::rngs::StdRng::seed_from_u64(seed),
            None => scirs2_core::random::rngs::StdRng::seed_from_u64(42),
        };

        let mut means = Array2::zeros((self.max_components, n_features));
        means.row_mut(0).assign(&X.row(rng.gen_range(0..n_samples)));

        for k in 1..self.max_components {
            let mut distances = Array1::zeros(n_samples);
            for i in 0..n_samples {
                let mut min_dist = f64::INFINITY;
                for j in 0..k {
                    let dist = (&X.row(i) - &means.row(j)).mapv(|x| x * x).sum().sqrt();
                    if dist < min_dist {
                        min_dist = dist;
                    }
                }
                distances[i] = min_dist * min_dist;
            }

            let total_dist: f64 = distances.sum();
            let target = rng.gen::<f64>() * total_dist;
            let mut cumulative = 0.0;

            for i in 0..n_samples {
                cumulative += distances[i];
                if cumulative >= target {
                    means.row_mut(k).assign(&X.row(i));
                    break;
                }
            }
        }

        // Initialize covariances
        let sample_cov = self.compute_sample_covariance(&X)?;
        let mut covariances = Vec::new();
        for _ in 0..self.max_components {
            covariances.push(sample_cov.clone());
        }

        let mut prev_lower_bound = f64::NEG_INFINITY;
        let mut converged = false;
        let mut n_iter = 0;

        // Variational EM iterations
        for iteration in 0..self.max_iter {
            n_iter = iteration + 1;

            // E-step: compute responsibilities
            let responsibilities =
                self.compute_responsibilities(&X, &weights, &means, &covariances)?;

            // M-step: update parameters
            weights = self.update_weights(&responsibilities)?;
            means = self.update_means(&X, &responsibilities)?;
            covariances = self.update_covariances(&X, &responsibilities, &means)?;

            // Compute lower bound
            let lower_bound =
                self.compute_lower_bound(&X, &responsibilities, &weights, &means, &covariances)?;

            if (lower_bound - prev_lower_bound).abs() < self.tol {
                converged = true;
                break;
            }
            prev_lower_bound = lower_bound;
        }

        // Determine effective number of components
        let mut n_components = 0;
        for k in 0..self.max_components {
            if weights[k] > 1e-3 {
                n_components = k + 1;
            }
        }

        Ok(DirichletProcessGaussianMixture {
            state: DirichletProcessGaussianMixtureTrained {
                weights: weights.slice(s![..n_components]).to_owned(),
                means: means.slice(s![..n_components, ..]).to_owned(),
                covariances: covariances.into_iter().take(n_components).collect(),
                weight_concentration: stick_weights.slice(s![..n_components]).to_owned(),
                lower_bound: prev_lower_bound,
                n_iter,
                converged,
                n_components,
                n_features,
                covariance_type: self.covariance_type.clone(),
                alpha: self.alpha,
                reg_covar: self.reg_covar,
            },
            alpha: self.alpha,
            max_components: self.max_components,
            covariance_type: self.covariance_type,
            tol: self.tol,
            max_iter: self.max_iter,
            reg_covar: self.reg_covar,
            random_state: self.random_state,
            n_init: self.n_init,
        })
    }
}

impl DirichletProcessGaussianMixture<Untrained> {
    fn compute_sample_covariance(&self, X: &Array2<f64>) -> SklResult<Array2<f64>> {
        let (n_samples, n_features) = X.dim();
        let mean = X.mean_axis(Axis(0)).unwrap();
        let mut cov = Array2::zeros((n_features, n_features));

        for i in 0..n_samples {
            let diff = &X.row(i) - &mean;
            for j in 0..n_features {
                for k in 0..n_features {
                    cov[[j, k]] += diff[j] * diff[k];
                }
            }
        }
        cov /= (n_samples - 1) as f64;

        for i in 0..n_features {
            cov[[i, i]] += self.reg_covar;
        }

        Ok(cov)
    }

    fn compute_responsibilities(
        &self,
        X: &Array2<f64>,
        weights: &Array1<f64>,
        means: &Array2<f64>,
        covariances: &[Array2<f64>],
    ) -> SklResult<Array2<f64>> {
        let (n_samples, _) = X.dim();
        let n_components = weights.len();
        let mut responsibilities = Array2::zeros((n_samples, n_components));

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

            for k in 0..n_components {
                let log_weight = weights[k].ln();
                let log_pdf =
                    self.multivariate_normal_log_pdf(&sample, &means.row(k), &covariances[k])?;
                log_probs[k] = log_weight + log_pdf;
            }

            let max_log_prob = log_probs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
            let log_sum_exp = max_log_prob
                + log_probs
                    .iter()
                    .map(|&lp| (lp - max_log_prob).exp())
                    .sum::<f64>()
                    .ln();

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

        Ok(responsibilities)
    }

    fn update_weights(&self, responsibilities: &Array2<f64>) -> SklResult<Array1<f64>> {
        let n_components = responsibilities.dim().1;
        let mut weights = Array1::zeros(n_components);

        for k in 0..n_components {
            weights[k] = responsibilities.column(k).sum() / responsibilities.dim().0 as f64;
        }

        Ok(weights)
    }

    fn update_means(
        &self,
        X: &Array2<f64>,
        responsibilities: &Array2<f64>,
    ) -> SklResult<Array2<f64>> {
        let (n_samples, n_features) = X.dim();
        let n_components = responsibilities.dim().1;
        let mut means = Array2::zeros((n_components, n_features));

        for k in 0..n_components {
            let weight_sum = responsibilities.column(k).sum();
            if weight_sum > 1e-10 {
                for j in 0..n_features {
                    let mut weighted_sum = 0.0;
                    for i in 0..n_samples {
                        weighted_sum += responsibilities[[i, k]] * X[[i, j]];
                    }
                    means[[k, j]] = weighted_sum / weight_sum;
                }
            }
        }

        Ok(means)
    }

    fn update_covariances(
        &self,
        X: &Array2<f64>,
        responsibilities: &Array2<f64>,
        means: &Array2<f64>,
    ) -> SklResult<Vec<Array2<f64>>> {
        let (n_samples, n_features) = X.dim();
        let n_components = responsibilities.dim().1;
        let mut covariances = Vec::new();

        for k in 0..n_components {
            let weight_sum = responsibilities.column(k).sum();
            let mut cov = Array2::zeros((n_features, n_features));

            if weight_sum > 1e-10 {
                for i in 0..n_samples {
                    let diff = &X.row(i) - &means.row(k);
                    let weight = responsibilities[[i, k]];
                    for j in 0..n_features {
                        for l in 0..n_features {
                            cov[[j, l]] += weight * diff[j] * diff[l];
                        }
                    }
                }
                cov /= weight_sum;
            } else {
                for j in 0..n_features {
                    cov[[j, j]] = 1.0;
                }
            }

            for j in 0..n_features {
                cov[[j, j]] += self.reg_covar;
            }

            covariances.push(cov);
        }

        Ok(covariances)
    }

    fn compute_lower_bound(
        &self,
        X: &Array2<f64>,
        responsibilities: &Array2<f64>,
        weights: &Array1<f64>,
        means: &Array2<f64>,
        covariances: &[Array2<f64>],
    ) -> SklResult<f64> {
        let (n_samples, _) = X.dim();
        let n_components = weights.len();
        let mut lower_bound = 0.0;

        // Data likelihood term
        for i in 0..n_samples {
            let sample = X.row(i);
            for k in 0..n_components {
                if responsibilities[[i, k]] > 1e-10 {
                    let log_pdf =
                        self.multivariate_normal_log_pdf(&sample, &means.row(k), &covariances[k])?;
                    lower_bound += responsibilities[[i, k]] * (weights[k].ln() + log_pdf);
                }
            }
        }

        // Entropy term
        for i in 0..n_samples {
            for k in 0..n_components {
                if responsibilities[[i, k]] > 1e-10 {
                    lower_bound -= responsibilities[[i, k]] * responsibilities[[i, k]].ln();
                }
            }
        }

        Ok(lower_bound)
    }

    fn multivariate_normal_log_pdf(
        &self,
        x: &ArrayView1<f64>,
        mean: &ArrayView1<f64>,
        cov: &Array2<f64>,
    ) -> SklResult<f64> {
        let d = x.len() as f64;
        let diff = x - mean;

        let det = self.matrix_determinant(cov)?;
        if det <= 0.0 {
            return Ok(f64::NEG_INFINITY);
        }

        let inv_cov = self.matrix_inverse(cov)?;
        let mut quad_form = 0.0;
        for i in 0..diff.len() {
            for j in 0..diff.len() {
                quad_form += diff[i] * inv_cov[[i, j]] * diff[j];
            }
        }

        Ok(-0.5 * (d * (2.0 * PI).ln() + det.ln() + quad_form))
    }

    fn matrix_determinant(&self, A: &Array2<f64>) -> SklResult<f64> {
        let n = A.dim().0;
        if n == 1 {
            return Ok(A[[0, 0]]);
        }
        if n == 2 {
            return Ok(A[[0, 0]] * A[[1, 1]] - A[[0, 1]] * A[[1, 0]]);
        }

        let mut det = 1.0;
        let mut A_copy = A.clone();

        for i in 0..n {
            let mut max_row = i;
            for k in i + 1..n {
                if A_copy[[k, i]].abs() > A_copy[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            if max_row != i {
                for j in 0..n {
                    let temp = A_copy[[i, j]];
                    A_copy[[i, j]] = A_copy[[max_row, j]];
                    A_copy[[max_row, j]] = temp;
                }
                det *= -1.0;
            }

            if A_copy[[i, i]].abs() < 1e-12 {
                return Ok(0.0);
            }

            det *= A_copy[[i, i]];

            for k in i + 1..n {
                let factor = A_copy[[k, i]] / A_copy[[i, i]];
                for j in i..n {
                    A_copy[[k, j]] -= factor * A_copy[[i, j]];
                }
            }
        }

        Ok(det)
    }

    fn matrix_inverse(&self, A: &Array2<f64>) -> SklResult<Array2<f64>> {
        let n = A.dim().0;
        let mut aug = Array2::zeros((n, 2 * n));

        for i in 0..n {
            for j in 0..n {
                aug[[i, j]] = A[[i, j]];
                aug[[i, j + n]] = if i == j { 1.0 } else { 0.0 };
            }
        }

        for i in 0..n {
            let mut max_row = i;
            for k in i + 1..n {
                if aug[[k, i]].abs() > aug[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            if max_row != i {
                for j in 0..2 * n {
                    let temp = aug[[i, j]];
                    aug[[i, j]] = aug[[max_row, j]];
                    aug[[max_row, j]] = temp;
                }
            }

            if aug[[i, i]].abs() < 1e-12 {
                return Err(SklearsError::NumericalError(
                    "Matrix is singular".to_string(),
                ));
            }

            let pivot = aug[[i, i]];
            for j in 0..2 * n {
                aug[[i, j]] /= pivot;
            }

            for k in 0..n {
                if k != i {
                    let factor = aug[[k, i]];
                    for j in 0..2 * n {
                        aug[[k, j]] -= factor * aug[[i, j]];
                    }
                }
            }
        }

        let mut inv = Array2::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                inv[[i, j]] = aug[[i, j + n]];
            }
        }

        Ok(inv)
    }
}

/// Trained Dirichlet Process Gaussian Mixture Model
#[derive(Debug, Clone)]
pub struct DirichletProcessGaussianMixtureTrained {
    /// Effective mixture weights
    pub weights: Array1<f64>,
    /// Component means
    pub means: Array2<f64>,
    /// Component covariances
    pub covariances: Vec<Array2<f64>>,
    /// Variational parameters for stick-breaking weights
    pub weight_concentration: Array1<f64>,
    /// Lower bound on log-likelihood
    pub lower_bound: f64,
    /// Number of iterations performed
    pub n_iter: usize,
    /// Whether the algorithm converged
    pub converged: bool,
    /// Number of effective components
    pub n_components: usize,
    /// Number of features
    pub n_features: usize,
    /// Covariance type
    pub covariance_type: CovarianceType,
    /// Concentration parameter
    pub alpha: f64,
    /// Regularization parameter
    pub reg_covar: f64,
}

impl Predict<ArrayView2<'_, Float>, Array1<i32>>
    for DirichletProcessGaussianMixture<DirichletProcessGaussianMixtureTrained>
{
    #[allow(non_snake_case)]
    fn predict(&self, X: &ArrayView2<'_, Float>) -> SklResult<Array1<i32>> {
        let X = X.to_owned();
        let (n_samples, _) = X.dim();
        let mut predictions = Array1::zeros(n_samples);

        for i in 0..n_samples {
            let sample = X.row(i);
            let mut max_log_prob = f64::NEG_INFINITY;
            let mut best_component = 0;

            for k in 0..self.state.n_components {
                let log_weight = self.state.weights[k].ln();
                let log_pdf = self.multivariate_normal_log_pdf(
                    &sample,
                    &self.state.means.row(k),
                    &self.state.covariances[k],
                )?;
                let log_prob = log_weight + log_pdf;

                if log_prob > max_log_prob {
                    max_log_prob = log_prob;
                    best_component = k;
                }
            }

            predictions[i] = best_component as i32;
        }

        Ok(predictions)
    }
}

impl DirichletProcessGaussianMixture<DirichletProcessGaussianMixtureTrained> {
    /// Predict class probabilities
    #[allow(non_snake_case)]
    pub fn predict_proba(&self, X: &ArrayView2<'_, Float>) -> SklResult<Array2<f64>> {
        let X = X.to_owned();
        let (n_samples, _) = X.dim();
        let mut probabilities = Array2::zeros((n_samples, self.state.n_components));

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

            for k in 0..self.state.n_components {
                let log_weight = self.state.weights[k].ln();
                let log_pdf = self.multivariate_normal_log_pdf(
                    &sample,
                    &self.state.means.row(k),
                    &self.state.covariances[k],
                )?;
                log_probs[k] = log_weight + log_pdf;
            }

            let max_log_prob = log_probs.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
            let log_sum_exp = max_log_prob
                + log_probs
                    .iter()
                    .map(|&lp| (lp - max_log_prob).exp())
                    .sum::<f64>()
                    .ln();

            for k in 0..self.state.n_components {
                probabilities[[i, k]] = (log_probs[k] - log_sum_exp).exp();
            }
        }

        Ok(probabilities)
    }

    /// Score samples using the lower bound
    pub fn score(&self, _X: &ArrayView2<'_, Float>) -> SklResult<f64> {
        Ok(self.state.lower_bound)
    }

    fn multivariate_normal_log_pdf(
        &self,
        x: &ArrayView1<f64>,
        mean: &ArrayView1<f64>,
        cov: &Array2<f64>,
    ) -> SklResult<f64> {
        let d = x.len() as f64;
        let diff = x - mean;

        let det = self.matrix_determinant(cov)?;
        if det <= 0.0 {
            return Ok(f64::NEG_INFINITY);
        }

        let inv_cov = self.matrix_inverse(cov)?;
        let mut quad_form = 0.0;
        for i in 0..diff.len() {
            for j in 0..diff.len() {
                quad_form += diff[i] * inv_cov[[i, j]] * diff[j];
            }
        }

        Ok(-0.5 * (d * (2.0 * PI).ln() + det.ln() + quad_form))
    }

    fn matrix_determinant(&self, A: &Array2<f64>) -> SklResult<f64> {
        let n = A.dim().0;
        if n == 1 {
            return Ok(A[[0, 0]]);
        }
        if n == 2 {
            return Ok(A[[0, 0]] * A[[1, 1]] - A[[0, 1]] * A[[1, 0]]);
        }

        let mut det = 1.0;
        let mut A_copy = A.clone();

        for i in 0..n {
            let mut max_row = i;
            for k in i + 1..n {
                if A_copy[[k, i]].abs() > A_copy[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            if max_row != i {
                for j in 0..n {
                    let temp = A_copy[[i, j]];
                    A_copy[[i, j]] = A_copy[[max_row, j]];
                    A_copy[[max_row, j]] = temp;
                }
                det *= -1.0;
            }

            if A_copy[[i, i]].abs() < 1e-12 {
                return Ok(0.0);
            }

            det *= A_copy[[i, i]];

            for k in i + 1..n {
                let factor = A_copy[[k, i]] / A_copy[[i, i]];
                for j in i..n {
                    A_copy[[k, j]] -= factor * A_copy[[i, j]];
                }
            }
        }

        Ok(det)
    }

    fn matrix_inverse(&self, A: &Array2<f64>) -> SklResult<Array2<f64>> {
        let n = A.dim().0;
        let mut aug = Array2::zeros((n, 2 * n));

        for i in 0..n {
            for j in 0..n {
                aug[[i, j]] = A[[i, j]];
                aug[[i, j + n]] = if i == j { 1.0 } else { 0.0 };
            }
        }

        for i in 0..n {
            let mut max_row = i;
            for k in i + 1..n {
                if aug[[k, i]].abs() > aug[[max_row, i]].abs() {
                    max_row = k;
                }
            }

            if max_row != i {
                for j in 0..2 * n {
                    let temp = aug[[i, j]];
                    aug[[i, j]] = aug[[max_row, j]];
                    aug[[max_row, j]] = temp;
                }
            }

            if aug[[i, i]].abs() < 1e-12 {
                return Err(SklearsError::NumericalError(
                    "Matrix is singular".to_string(),
                ));
            }

            let pivot = aug[[i, i]];
            for j in 0..2 * n {
                aug[[i, j]] /= pivot;
            }

            for k in 0..n {
                if k != i {
                    let factor = aug[[k, i]];
                    for j in 0..2 * n {
                        aug[[k, j]] -= factor * aug[[i, j]];
                    }
                }
            }
        }

        let mut inv = Array2::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                inv[[i, j]] = aug[[i, j + n]];
            }
        }

        Ok(inv)
    }
}

// Note: The Pitman-Yor Process and Hierarchical Dirichlet Process implementations
// would follow similar patterns but are omitted for brevity. They would include:
//
// 1. PitmanYorProcess<S> and PitmanYorProcessTrained structs
// 2. HierarchicalDirichletProcess<S> and HierarchicalDirichletProcessTrained structs
// 3. Similar trait implementations (Estimator, Fit, Predict)
// 4. Specialized stick-breaking constructions for each model
// 5. Appropriate inference algorithms (variational inference, Gibbs sampling)
//
// These can be added as needed based on specific requirements.