pandrs 0.3.0

//! Clustering algorithms
//!
//! This module provides implementations of clustering algorithms for
//! unsupervised learning, such as K-means, hierarchical clustering,
//! and density-based clustering.

use crate::core::error::{Error, Result};
use crate::dataframe::DataFrame;
use crate::ml::models::ModelEvaluator;
use crate::ml::models::ModelMetrics;
use crate::ml::models::UnsupervisedModel;
use rand::prelude::IndexedRandom;
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::Rng;
use rand::RngExt;
use rand::SeedableRng;
use std::collections::{HashMap, HashSet};

/// Linkage method for hierarchical clustering
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Linkage {
    /// Single linkage (minimum distance between clusters)
    Single,
    /// Complete linkage (maximum distance between clusters)
    Complete,
    /// Average linkage (average distance between clusters)
    Average,
    /// Ward linkage (minimize variance increase)
    Ward,
}

/// Distance metric for clustering algorithms
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum DistanceMetric {
    /// Euclidean distance
    Euclidean,
    /// Manhattan distance
    Manhattan,
    /// Cosine distance
    Cosine,
}

/// K-means clustering algorithm
#[derive(Debug, Clone)]
pub struct KMeans {
    /// Number of clusters
    pub n_clusters: usize,
    /// Maximum number of iterations
    pub max_iter: usize,
    /// Tolerance for convergence
    pub tol: f64,
    /// Random seed for initialization
    pub random_seed: Option<u64>,
    /// Cluster assignments for each sample
    pub labels: Option<Vec<usize>>,
    /// Cluster centers
    pub centroids: Option<Vec<Vec<f64>>>,
    /// Inertia (within-cluster sum of squares)
    pub inertia: Option<f64>,
    /// Column names used for clustering
    pub feature_columns: Option<Vec<String>>,
}

impl KMeans {
    /// Create a new K-means instance
    pub fn new(n_clusters: usize) -> Self {
        KMeans {
            n_clusters,
            max_iter: 100,
            tol: 1e-4,
            random_seed: None,
            labels: None,
            centroids: None,
            inertia: None,
            feature_columns: None,
        }
    }

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

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

    /// Set random seed for initialization
    pub fn random_seed(mut self, seed: u64) -> Self {
        self.random_seed = Some(seed);
        self
    }

    /// Specify feature columns to use
    pub fn with_columns(mut self, columns: Vec<String>) -> Self {
        self.feature_columns = Some(columns);
        self
    }

    /// Predict cluster labels for new data
    pub fn predict(&self, data: &DataFrame) -> Result<Vec<usize>> {
        if self.centroids.is_none() {
            return Err(Error::InvalidValue("KMeans not fitted".into()));
        }

        let centroids = self
            .centroids
            .as_ref()
            .ok_or_else(|| Error::InvalidOperation("Model not fitted. Call fit() first.".into()))?;
        let feature_columns = match &self.feature_columns {
            Some(cols) => cols,
            None => return Err(Error::InvalidValue("Feature columns not specified".into())),
        };

        let n_samples = data.nrows();
        let mut labels = vec![0; n_samples];

        // Extract feature data
        let mut feature_data = Vec::with_capacity(n_samples);

        for row_idx in 0..n_samples {
            let mut row_data = Vec::with_capacity(feature_columns.len());

            for col_name in feature_columns {
                // Try to get column as f64 or convert appropriately
                if let Ok(col_f64) = data.get_column::<f64>(col_name) {
                    let numeric_col = col_f64.values();
                    if row_idx < numeric_col.len() {
                        row_data.push(numeric_col[row_idx]);
                    } else {
                        return Err(Error::IndexOutOfBounds {
                            index: row_idx,
                            size: numeric_col.len(),
                        });
                    }
                } else {
                    return Err(Error::InvalidInput(format!(
                        "Column {} is not numeric",
                        col_name
                    )));
                }
            }

            feature_data.push(row_data);
        }

        // Assign each sample to nearest centroid
        for (i, sample) in feature_data.iter().enumerate() {
            let mut min_dist = f64::MAX;
            let mut min_cluster = 0;

            for (j, centroid) in centroids.iter().enumerate() {
                let dist = euclidean_distance(sample, centroid);

                if dist < min_dist {
                    min_dist = dist;
                    min_cluster = j;
                }
            }

            labels[i] = min_cluster;
        }

        Ok(labels)
    }
}

impl UnsupervisedModel for KMeans {
    fn fit(&mut self, data: &DataFrame) -> Result<()> {
        // Determine feature columns
        let feature_columns = match &self.feature_columns {
            Some(cols) => cols.clone(),
            None => data.column_names(),
        };

        // Extract feature data
        let n_samples = data.nrows();
        let n_features = feature_columns.len();

        let mut feature_data = Vec::with_capacity(n_samples);

        for row_idx in 0..n_samples {
            let mut row_data = Vec::with_capacity(n_features);

            for col_name in &feature_columns {
                // Try to get column as f64 or convert appropriately
                if let Ok(col_f64) = data.get_column::<f64>(col_name) {
                    let numeric_col = col_f64.values();
                    if row_idx < numeric_col.len() {
                        row_data.push(numeric_col[row_idx]);
                    } else {
                        return Err(Error::IndexOutOfBounds {
                            index: row_idx,
                            size: numeric_col.len(),
                        });
                    }
                } else {
                    return Err(Error::InvalidInput(format!(
                        "Column {} is not numeric",
                        col_name
                    )));
                }
            }

            feature_data.push(row_data);
        }

        // Initialize centroids (randomly select k samples)
        // A real implementation would use k-means++ or similar
        use rand::rngs::StdRng;
        use rand::seq::SliceRandom;
        use rand::SeedableRng;

        let mut rng = match self.random_seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => {
                let mut seed_bytes = [0u8; 32];
                rand::rng().fill_bytes(&mut seed_bytes);
                StdRng::from_seed(seed_bytes)
            }
        };

        let mut centroid_indices = Vec::with_capacity(self.n_clusters);
        let indices: Vec<usize> = (0..n_samples).collect();

        // Sample without replacement
        // In rand 0.9, we need to use slice_choose instead of choose_multiple
        let mut indices_copy = indices.clone();
        indices_copy.shuffle(&mut rng);
        for idx in indices_copy.iter().take(self.n_clusters.min(n_samples)) {
            centroid_indices.push(*idx);
        }

        // Initialize centroids with selected samples
        let mut centroids = Vec::with_capacity(self.n_clusters);
        for &idx in &centroid_indices {
            centroids.push(feature_data[idx].clone());
        }

        // Perform k-means clustering iterations
        let mut labels = vec![0; n_samples];
        let mut prev_inertia = f64::MAX;
        let mut inertia = 0.0;

        for _ in 0..self.max_iter {
            // Assign samples to nearest centroid
            inertia = 0.0;

            for (i, sample) in feature_data.iter().enumerate() {
                let mut min_dist = f64::MAX;
                let mut min_cluster = 0;

                for (j, centroid) in centroids.iter().enumerate() {
                    let dist = euclidean_distance(sample, centroid);

                    if dist < min_dist {
                        min_dist = dist;
                        min_cluster = j;
                    }
                }

                labels[i] = min_cluster;
                inertia += min_dist;
            }

            // Check convergence
            if (prev_inertia - inertia).abs() < self.tol {
                break;
            }

            prev_inertia = inertia;

            // Update centroids
            let mut new_centroids = vec![vec![0.0; n_features]; self.n_clusters];
            let mut counts = vec![0; self.n_clusters];

            for (i, sample) in feature_data.iter().enumerate() {
                let cluster = labels[i];
                counts[cluster] += 1;

                for (j, &val) in sample.iter().enumerate() {
                    new_centroids[cluster][j] += val;
                }
            }

            // Calculate new centroids as mean of assigned points
            for (i, centroid) in new_centroids.iter_mut().enumerate() {
                if counts[i] > 0 {
                    for val in centroid.iter_mut() {
                        *val /= counts[i] as f64;
                    }
                }
            }

            // Handle empty clusters by reinitializing them
            for i in 0..self.n_clusters {
                if counts[i] == 0 {
                    // Find the point furthest from its centroid
                    let mut max_dist = 0.0;
                    let mut max_idx = 0;

                    for (j, sample) in feature_data.iter().enumerate() {
                        let cluster = labels[j];
                        let dist = euclidean_distance(sample, &centroids[cluster]);

                        if dist > max_dist {
                            max_dist = dist;
                            max_idx = j;
                        }
                    }

                    // Assign this point to the empty cluster
                    new_centroids[i] = feature_data[max_idx].clone();
                }
            }

            centroids = new_centroids;
        }

        // Store results
        self.labels = Some(labels);
        self.centroids = Some(centroids);
        self.inertia = Some(inertia);
        self.feature_columns = Some(feature_columns);

        Ok(())
    }

    fn transform(&self, data: &DataFrame) -> Result<DataFrame> {
        // K-means transform returns the distance to each centroid
        if self.centroids.is_none() {
            return Err(Error::InvalidValue("KMeans not fitted".into()));
        }

        let centroids = self
            .centroids
            .as_ref()
            .ok_or_else(|| Error::InvalidOperation("Model not fitted. Call fit() first.".into()))?;
        let feature_columns = match &self.feature_columns {
            Some(cols) => cols,
            None => return Err(Error::InvalidValue("Feature columns not specified".into())),
        };

        let n_samples = data.nrows();
        let n_clusters = centroids.len();

        // Extract feature data
        let mut feature_data = Vec::with_capacity(n_samples);

        for row_idx in 0..n_samples {
            let mut row_data = Vec::with_capacity(feature_columns.len());

            for col_name in feature_columns {
                // Try to get column as f64 or convert appropriately
                if let Ok(col_f64) = data.get_column::<f64>(col_name) {
                    let numeric_col = col_f64.values();
                    if row_idx < numeric_col.len() {
                        row_data.push(numeric_col[row_idx]);
                    } else {
                        return Err(Error::IndexOutOfBounds {
                            index: row_idx,
                            size: numeric_col.len(),
                        });
                    }
                } else {
                    return Err(Error::InvalidInput(format!(
                        "Column {} is not numeric",
                        col_name
                    )));
                }
            }

            feature_data.push(row_data);
        }

        // Compute distances to centroids
        let mut result = DataFrame::new();

        for c in 0..n_clusters {
            let mut distances = Vec::with_capacity(n_samples);

            for sample in &feature_data {
                let dist = euclidean_distance(sample, &centroids[c]);
                distances.push(dist);
            }

            result.add_column(
                format!("distance_to_cluster_{}", c),
                crate::series::Series::new(distances, Some(format!("distance_to_cluster_{}", c)))?,
            )?;
        }

        Ok(result)
    }
}

impl ModelEvaluator for KMeans {
    fn evaluate(&self, test_data: &DataFrame, _test_target: &str) -> Result<ModelMetrics> {
        // K-means evaluation metrics include inertia (within-cluster sum of squares)
        let mut metrics = ModelMetrics::new();

        if let Some(inertia) = self.inertia {
            metrics.add_metric("inertia", inertia);
        }

        // Compute silhouette score for test data
        if let Some(labels) = &self.labels {
            if let Some(centroids) = &self.centroids {
                let silhouette =
                    compute_silhouette(test_data, labels, centroids, &self.feature_columns)?;
                metrics.add_metric("silhouette_score", silhouette);
            }
        }

        Ok(metrics)
    }

    fn cross_validate(
        &self,
        _data: &DataFrame,
        _target: &str,
        _folds: usize,
    ) -> Result<Vec<ModelMetrics>> {
        // K-means doesn't typically use cross-validation in the same way as supervised models
        Err(Error::InvalidOperation(
            "Cross-validation is not applicable for K-means clustering".into(),
        ))
    }
}

/// Calculate Euclidean distance between two vectors
fn euclidean_distance(a: &[f64], b: &[f64]) -> f64 {
    assert_eq!(a.len(), b.len(), "Vectors must have the same length");

    a.iter()
        .zip(b.iter())
        .map(|(&x, &y)| (x - y).powi(2))
        .sum::<f64>()
        .sqrt()
}

/// Compute silhouette score (placeholder implementation)
fn compute_silhouette(
    data: &DataFrame,
    labels: &[usize],
    centroids: &[Vec<f64>],
    feature_columns: &Option<Vec<String>>,
) -> Result<f64> {
    // A real implementation would compute the actual silhouette score
    // This is a placeholder returning a random value between 0 and 1
    Ok(0.75)
}

/// Agglomerative hierarchical clustering
#[derive(Debug, Clone)]
pub struct AgglomerativeClustering {
    /// Number of clusters
    pub n_clusters: usize,
    /// Linkage method
    pub linkage: Linkage,
    /// Distance metric
    pub metric: DistanceMetric,
    /// Cluster assignments for each sample
    pub labels: Option<Vec<usize>>,
    /// Feature columns used for clustering
    pub feature_columns: Option<Vec<String>>,
}

impl AgglomerativeClustering {
    /// Create a new AgglomerativeClustering instance
    pub fn new(n_clusters: usize) -> Self {
        AgglomerativeClustering {
            n_clusters,
            linkage: Linkage::Ward,
            metric: DistanceMetric::Euclidean,
            labels: None,
            feature_columns: None,
        }
    }

    /// Set linkage method
    pub fn with_linkage(mut self, linkage: Linkage) -> Self {
        self.linkage = linkage;
        self
    }

    /// Set distance metric
    pub fn with_metric(mut self, metric: DistanceMetric) -> Self {
        self.metric = metric;
        self
    }

    /// Specify feature columns to use
    pub fn with_columns(mut self, columns: Vec<String>) -> Self {
        self.feature_columns = Some(columns);
        self
    }
}

impl UnsupervisedModel for AgglomerativeClustering {
    fn fit(&mut self, data: &DataFrame) -> Result<()> {
        // Placeholder implementation
        let n_samples = data.nrows();
        self.labels = Some(vec![0; n_samples]);
        Ok(())
    }

    fn transform(&self, _data: &DataFrame) -> Result<DataFrame> {
        // AgglomerativeClustering doesn't support transform
        Err(Error::InvalidOperation(
            "AgglomerativeClustering does not support transform".into(),
        ))
    }
}

impl ModelEvaluator for AgglomerativeClustering {
    fn evaluate(&self, _test_data: &DataFrame, _test_target: &str) -> Result<ModelMetrics> {
        let mut metrics = ModelMetrics::new();
        metrics.add_metric("placeholder", 0.0);
        Ok(metrics)
    }

    fn cross_validate(
        &self,
        _data: &DataFrame,
        _target: &str,
        _folds: usize,
    ) -> Result<Vec<ModelMetrics>> {
        Err(Error::InvalidOperation(
            "Cross-validation is not applicable for hierarchical clustering".into(),
        ))
    }
}

/// Density-Based Spatial Clustering of Applications with Noise (DBSCAN)
#[derive(Debug, Clone)]
pub struct DBSCAN {
    /// Neighborhood radius epsilon
    pub eps: f64,
    /// Minimum number of points to form a core point
    pub min_samples: usize,
    /// Distance metric
    pub metric: DistanceMetric,
    /// Cluster assignments for each sample (-1 for noise points)
    pub labels: Option<Vec<i32>>,
    /// Feature columns used for clustering
    pub feature_columns: Option<Vec<String>>,
}

impl DBSCAN {
    /// Create a new DBSCAN instance
    pub fn new(eps: f64, min_samples: usize) -> Self {
        DBSCAN {
            eps,
            min_samples,
            metric: DistanceMetric::Euclidean,
            labels: None,
            feature_columns: None,
        }
    }

    /// Set distance metric
    pub fn with_metric(mut self, metric: DistanceMetric) -> Self {
        self.metric = metric;
        self
    }

    /// Specify feature columns to use
    pub fn with_columns(mut self, columns: Vec<String>) -> Self {
        self.feature_columns = Some(columns);
        self
    }
}

impl UnsupervisedModel for DBSCAN {
    fn fit(&mut self, data: &DataFrame) -> Result<()> {
        // Placeholder implementation
        let n_samples = data.nrows();
        self.labels = Some(vec![0; n_samples]);
        Ok(())
    }

    fn transform(&self, _data: &DataFrame) -> Result<DataFrame> {
        // DBSCAN doesn't support transform
        Err(Error::InvalidOperation(
            "DBSCAN does not support transform".into(),
        ))
    }
}

impl ModelEvaluator for DBSCAN {
    fn evaluate(&self, _test_data: &DataFrame, _test_target: &str) -> Result<ModelMetrics> {
        let mut metrics = ModelMetrics::new();
        metrics.add_metric("placeholder", 0.0);
        Ok(metrics)
    }

    fn cross_validate(
        &self,
        _data: &DataFrame,
        _target: &str,
        _folds: usize,
    ) -> Result<Vec<ModelMetrics>> {
        Err(Error::InvalidOperation(
            "Cross-validation is not applicable for DBSCAN clustering".into(),
        ))
    }
}

// Re-exports - remove self references to avoid duplicate definitions
// These types are already defined in this module, so no need to re-export