leptos-helios 0.8.1

//! Clustering Analysis Module

use super::MLError;
use crate::DataFrame;
// use polars::prelude::*; // Currently unused
// use std::collections::HashMap; // Currently unused

/// Cluster analyzer for grouping similar data points
pub struct ClusterAnalyzer {
    /// Number of clusters to create
    pub n_clusters: usize,
    /// Clustering algorithm to use
    pub algorithm: ClusteringAlgorithm,
    /// Cluster centers (for K-means)
    pub centers: Option<Vec<Vec<f64>>>,
    /// Cluster assignments
    pub assignments: Option<Vec<usize>>,
    /// Clustering results
    pub results: Option<ClusteringResults>,
}

/// Clustering algorithms supported
#[derive(Debug, Clone, PartialEq)]
pub enum ClusteringAlgorithm {
    /// K-Means clustering
    KMeans,
    /// DBSCAN clustering
    DBSCAN,
    /// Hierarchical clustering
    Hierarchical,
    /// Gaussian Mixture Model
    GaussianMixture,
}

/// Clustering results
#[derive(Debug, Clone)]
pub struct ClusteringResults {
    /// Cluster assignments for each data point
    pub assignments: Vec<usize>,
    /// Cluster centers (for centroid-based algorithms)
    pub centers: Option<Vec<Vec<f64>>>,
    /// Number of clusters found
    pub n_clusters: usize,
    /// Silhouette score (clustering quality metric)
    pub silhouette_score: Option<f64>,
    /// Inertia (for K-means)
    pub inertia: Option<f64>,
    /// Cluster statistics
    pub cluster_stats: Vec<ClusterStats>,
}

/// Statistics for each cluster
#[derive(Debug, Clone)]
pub struct ClusterStats {
    /// Cluster ID
    pub cluster_id: usize,
    /// Number of points in cluster
    pub size: usize,
    /// Center of the cluster
    pub center: Vec<f64>,
    /// Average distance to center
    pub avg_distance: f64,
    /// Maximum distance to center
    pub max_distance: f64,
    /// Minimum distance to center
    pub min_distance: f64,
}

/// Clustering configuration
#[derive(Debug, Clone)]
pub struct ClusteringConfig {
    /// Number of clusters
    pub n_clusters: usize,
    /// Algorithm to use
    pub algorithm: ClusteringAlgorithm,
    /// Maximum iterations (for iterative algorithms)
    pub max_iterations: usize,
    /// Tolerance for convergence
    pub tolerance: f64,
    /// Random seed for reproducibility
    pub random_seed: Option<u64>,
}

impl ClusterAnalyzer {
    /// Create a new cluster analyzer
    pub fn new(n_clusters: usize, algorithm: ClusteringAlgorithm) -> Self {
        Self {
            n_clusters,
            algorithm,
            centers: None,
            assignments: None,
            results: None,
        }
    }

    /// Perform clustering on the data
    pub fn cluster(&mut self, data: &DataFrame) -> Result<ClusteringResults, MLError> {
        if data.height() < self.n_clusters {
            return Err(MLError::InsufficientData(
                format!(
                    "Need at least {} data points for {} clusters",
                    self.n_clusters, self.n_clusters
                ),
            ));
        }

        let points = self.extract_points(data)?;

        let results = match self.algorithm {
            ClusteringAlgorithm::KMeans => self.kmeans_clustering(&points)?,
            ClusteringAlgorithm::DBSCAN => self.dbscan_clustering(&points)?,
            ClusteringAlgorithm::Hierarchical => self.hierarchical_clustering(&points)?,
            ClusteringAlgorithm::GaussianMixture => self.gaussian_mixture_clustering(&points)?,
        };

        self.results = Some(results.clone());
        Ok(results)
    }

    /// Extract data points from DataFrame
    fn extract_points(&self, data: &DataFrame) -> Result<Vec<Vec<f64>>, MLError> {
        let mut points = Vec::new();
        let columns = data.get_columns();

        // Find numeric columns
        let mut numeric_columns = Vec::new();
        for column in columns {
            if let Ok(series) = column.f64() {
                numeric_columns.push(series);
            }
        }

        if numeric_columns.is_empty() {
            return Err(MLError::InvalidData(
                "No numeric columns found for clustering".to_string(),
            ));
        }

        let n_rows = numeric_columns[0].len();
        for i in 0..n_rows {
            let mut point = Vec::new();
            for series in &numeric_columns {
                if let Some(value) = series.get(i) {
                    point.push(value);
                } else {
                    return Err(MLError::InvalidData(
                        "Missing values in numeric columns".to_string(),
                    ));
                }
            }
            points.push(point);
        }

        Ok(points)
    }

    /// Perform K-means clustering
    fn kmeans_clustering(&mut self, points: &[Vec<f64>]) -> Result<ClusteringResults, MLError> {
        let n_points = points.len();
        let n_features = points[0].len();
        let max_iterations = 100;
        let tolerance = 1e-4;

        // Initialize cluster centers randomly
        let mut centers = self.initialize_centers(points, self.n_clusters);
        let mut assignments = vec![0; n_points];
        let mut prev_inertia = f64::INFINITY;

        for _iteration in 0..max_iterations {
            // Assign points to nearest centers
            for (i, point) in points.iter().enumerate() {
                let mut min_distance = f64::INFINITY;
                let mut best_cluster = 0;

                for (j, center) in centers.iter().enumerate() {
                    let distance = self.euclidean_distance(point, center);
                    if distance < min_distance {
                        min_distance = distance;
                        best_cluster = j;
                    }
                }
                assignments[i] = best_cluster;
            }

            // Update centers
            let mut new_centers = vec![vec![0.0; n_features]; self.n_clusters];
            let mut cluster_counts = vec![0; self.n_clusters];

            for (i, point) in points.iter().enumerate() {
                let cluster = assignments[i];
                for (j, &value) in point.iter().enumerate() {
                    new_centers[cluster][j] += value;
                }
                cluster_counts[cluster] += 1;
            }

            for i in 0..self.n_clusters {
                if cluster_counts[i] > 0 {
                    for j in 0..n_features {
                        new_centers[i][j] /= cluster_counts[i] as f64;
                    }
                }
            }

            // Calculate inertia
            let inertia = self.calculate_inertia(points, &assignments, &new_centers);

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

            centers = new_centers;
            prev_inertia = inertia;
        }

        self.centers = Some(centers.clone());
        self.assignments = Some(assignments.clone());

        // Calculate cluster statistics
        let cluster_stats = self.calculate_cluster_stats(points, &assignments, &centers);
        let silhouette_score = self.calculate_silhouette_score(points, &assignments);

        Ok(ClusteringResults {
            assignments,
            centers: Some(centers),
            n_clusters: self.n_clusters,
            silhouette_score: Some(silhouette_score),
            inertia: Some(prev_inertia),
            cluster_stats,
        })
    }

    /// Perform DBSCAN clustering
    fn dbscan_clustering(&mut self, points: &[Vec<f64>]) -> Result<ClusteringResults, MLError> {
        let eps = 0.5; // Distance threshold
        let min_points = 3; // Minimum points for core point

        let mut assignments = vec![usize::MAX; points.len()]; // -1 for noise
        let mut cluster_id = 0;

        for i in 0..points.len() {
            if assignments[i] != usize::MAX {
                continue; // Already processed
            }

            let neighbors = self.find_neighbors(points, i, eps);
            if neighbors.len() < min_points {
                assignments[i] = usize::MAX; // Noise point
                continue;
            }

            // Start new cluster
            assignments[i] = cluster_id;
            let mut seed_set = neighbors;

            while let Some(point_idx) = seed_set.pop() {
                if assignments[point_idx] == usize::MAX {
                    assignments[point_idx] = cluster_id;
                }
                if assignments[point_idx] != usize::MAX {
                    continue;
                }

                assignments[point_idx] = cluster_id;
                let point_neighbors = self.find_neighbors(points, point_idx, eps);
                if point_neighbors.len() >= min_points {
                    seed_set.extend(point_neighbors);
                }
            }

            cluster_id += 1;
        }

        let n_clusters = cluster_id;
        let cluster_stats = self.calculate_dbscan_cluster_stats(points, &assignments, n_clusters);

        Ok(ClusteringResults {
            assignments,
            centers: None,
            n_clusters,
            silhouette_score: None,
            inertia: None,
            cluster_stats,
        })
    }

    /// Perform hierarchical clustering
    fn hierarchical_clustering(&mut self, points: &[Vec<f64>]) -> Result<ClusteringResults, MLError> {
        // Simplified hierarchical clustering using single linkage
        let n_points = points.len();
        let mut assignments = vec![0; n_points];

        // Start with each point as its own cluster
        let mut clusters: Vec<Vec<usize>> = (0..n_points).map(|i| vec![i]).collect();

        while clusters.len() > self.n_clusters {
            // Find closest pair of clusters
            let mut min_distance = f64::INFINITY;
            let mut best_pair = (0, 1);

            for i in 0..clusters.len() {
                for j in (i + 1)..clusters.len() {
                    let distance = self.cluster_distance(&clusters[i], &clusters[j], points);
                    if distance < min_distance {
                        min_distance = distance;
                        best_pair = (i, j);
                    }
                }
            }

            // Merge clusters
            let (i, j) = best_pair;
            let mut merged_cluster = clusters[i].clone();
            merged_cluster.extend(clusters[j].clone());
            clusters[i] = merged_cluster;
            clusters.remove(j);
        }

        // Assign cluster IDs
        for (cluster_id, cluster) in clusters.iter().enumerate() {
            for &point_idx in cluster {
                assignments[point_idx] = cluster_id;
            }
        }

        let cluster_stats = self.calculate_hierarchical_cluster_stats(points, &assignments, clusters.len());

        Ok(ClusteringResults {
            assignments,
            centers: None,
            n_clusters: clusters.len(),
            silhouette_score: None,
            inertia: None,
            cluster_stats,
        })
    }

    /// Perform Gaussian Mixture Model clustering
    fn gaussian_mixture_clustering(&mut self, points: &[Vec<f64>]) -> Result<ClusteringResults, MLError> {
        // Simplified GMM implementation
        let n_points = points.len();
        let n_features = points[0].len();
        let max_iterations = 50;

        // Initialize means randomly
        let mut means = self.initialize_centers(points, self.n_clusters);
        let covariances = vec![vec![vec![1.0; n_features]; n_features]; self.n_clusters];
        let mut weights = vec![1.0 / self.n_clusters as f64; self.n_clusters];
        let mut assignments = vec![0; n_points];

        for _ in 0..max_iterations {
            // E-step: Calculate responsibilities
            let mut responsibilities = vec![vec![0.0; self.n_clusters]; n_points];

            for (i, point) in points.iter().enumerate() {
                let mut total_prob = 0.0;
                for k in 0..self.n_clusters {
                    let prob = self.gaussian_probability(point, &means[k], &covariances[k]) * weights[k];
                    responsibilities[i][k] = prob;
                    total_prob += prob;
                }

                // Normalize
                if total_prob > 0.0 {
                    for k in 0..self.n_clusters {
                        responsibilities[i][k] /= total_prob;
                    }
                }
            }

            // M-step: Update parameters
            for k in 0..self.n_clusters {
                let mut sum_resp = 0.0;
                for i in 0..n_points {
                    sum_resp += responsibilities[i][k];
                }

                if sum_resp > 0.0 {
                    // Update mean
                    for j in 0..n_features {
                        means[k][j] = 0.0;
                        for i in 0..n_points {
                            means[k][j] += responsibilities[i][k] * points[i][j];
                        }
                        means[k][j] /= sum_resp;
                    }

                    // Update weight
                    weights[k] = sum_resp / n_points as f64;
                }
            }

            // Assign points to clusters
            for (i, _point) in points.iter().enumerate() {
                let mut max_prob = 0.0;
                let mut best_cluster = 0;
                for k in 0..self.n_clusters {
                    if responsibilities[i][k] > max_prob {
                        max_prob = responsibilities[i][k];
                        best_cluster = k;
                    }
                }
                assignments[i] = best_cluster;
            }
        }

        let cluster_stats = self.calculate_cluster_stats(points, &assignments, &means);

        Ok(ClusteringResults {
            assignments,
            centers: Some(means),
            n_clusters: self.n_clusters,
            silhouette_score: None,
            inertia: None,
            cluster_stats,
        })
    }

    /// Initialize cluster centers randomly
    fn initialize_centers(&self, points: &[Vec<f64>], n_clusters: usize) -> Vec<Vec<f64>> {
        let n_features = points[0].len();
        let mut centers = Vec::new();

        for _ in 0..n_clusters {
            let mut center = Vec::new();
            for j in 0..n_features {
                let min_val = points.iter().map(|p| p[j]).fold(f64::INFINITY, f64::min);
                let max_val = points.iter().map(|p| p[j]).fold(f64::NEG_INFINITY, f64::max);
                let random_val = min_val + (max_val - min_val) * 0.5; // Simplified random
                center.push(random_val);
            }
            centers.push(center);
        }

        centers
    }

    /// Calculate Euclidean distance between two points
    fn euclidean_distance(&self, a: &[f64], b: &[f64]) -> f64 {
        a.iter()
            .zip(b.iter())
            .map(|(x, y)| (x - y).powi(2))
            .sum::<f64>()
            .sqrt()
    }

    /// Calculate inertia (sum of squared distances to centers)
    fn calculate_inertia(&self, points: &[Vec<f64>], assignments: &[usize], centers: &[Vec<f64>]) -> f64 {
        let mut inertia = 0.0;
        for (i, point) in points.iter().enumerate() {
            let cluster = assignments[i];
            let distance = self.euclidean_distance(point, &centers[cluster]);
            inertia += distance * distance;
        }
        inertia
    }

    /// Find neighbors within eps distance
    fn find_neighbors(&self, points: &[Vec<f64>], point_idx: usize, eps: f64) -> Vec<usize> {
        let mut neighbors = Vec::new();
        let point = &points[point_idx];

        for (i, other_point) in points.iter().enumerate() {
            if i != point_idx {
                let distance = self.euclidean_distance(point, other_point);
                if distance <= eps {
                    neighbors.push(i);
                }
            }
        }

        neighbors
    }

    /// Calculate distance between two clusters (single linkage)
    fn cluster_distance(&self, cluster1: &[usize], cluster2: &[usize], points: &[Vec<f64>]) -> f64 {
        let mut min_distance = f64::INFINITY;

        for &i in cluster1 {
            for &j in cluster2 {
                let distance = self.euclidean_distance(&points[i], &points[j]);
                if distance < min_distance {
                    min_distance = distance;
                }
            }
        }

        min_distance
    }

    /// Calculate Gaussian probability
    fn gaussian_probability(&self, point: &[f64], mean: &[f64], covariance: &[Vec<f64>]) -> f64 {
        // Simplified 1D Gaussian for now
        let diff = point[0] - mean[0];
        let variance = covariance[0][0];
        (-0.5 * diff * diff / variance).exp() / (2.0 * std::f64::consts::PI * variance).sqrt()
    }

    /// Calculate cluster statistics
    fn calculate_cluster_stats(
        &self,
        points: &[Vec<f64>],
        assignments: &[usize],
        centers: &[Vec<f64>],
    ) -> Vec<ClusterStats> {
        let mut cluster_stats = Vec::new();

        for cluster_id in 0..self.n_clusters {
            let cluster_points: Vec<&Vec<f64>> = points
                .iter()
                .zip(assignments.iter())
                .filter(|(_, &assignment)| assignment == cluster_id)
                .map(|(point, _)| point)
                .collect();

            if cluster_points.is_empty() {
                continue;
            }

            let size = cluster_points.len();
            let center = centers[cluster_id].clone();

            let mut distances = Vec::new();
            for point in &cluster_points {
                let distance = self.euclidean_distance(point, &center);
                distances.push(distance);
            }

            distances.sort_by(|a, b| a.partial_cmp(b).unwrap());

            cluster_stats.push(ClusterStats {
                cluster_id,
                size,
                center,
                avg_distance: distances.iter().sum::<f64>() / size as f64,
                max_distance: *distances.last().unwrap(),
                min_distance: *distances.first().unwrap(),
            });
        }

        cluster_stats
    }

    /// Calculate DBSCAN cluster statistics
    fn calculate_dbscan_cluster_stats(
        &self,
        points: &[Vec<f64>],
        assignments: &[usize],
        n_clusters: usize,
    ) -> Vec<ClusterStats> {
        let mut cluster_stats = Vec::new();

        for cluster_id in 0..n_clusters {
            let cluster_points: Vec<&Vec<f64>> = points
                .iter()
                .zip(assignments.iter())
                .filter(|(_, &assignment)| assignment == cluster_id)
                .map(|(point, _)| point)
                .collect();

            if cluster_points.is_empty() {
                continue;
            }

            let size = cluster_points.len();

            // Calculate center
            let n_features = cluster_points[0].len();
            let mut center = vec![0.0; n_features];
            for point in &cluster_points {
                for (j, &value) in point.iter().enumerate() {
                    center[j] += value;
                }
            }
            for j in 0..n_features {
                center[j] /= size as f64;
            }

            // Calculate distances
            let mut distances = Vec::new();
            for point in &cluster_points {
                let distance = self.euclidean_distance(point, &center);
                distances.push(distance);
            }

            distances.sort_by(|a, b| a.partial_cmp(b).unwrap());

            cluster_stats.push(ClusterStats {
                cluster_id,
                size,
                center,
                avg_distance: distances.iter().sum::<f64>() / size as f64,
                max_distance: *distances.last().unwrap(),
                min_distance: *distances.first().unwrap(),
            });
        }

        cluster_stats
    }

    /// Calculate hierarchical cluster statistics
    fn calculate_hierarchical_cluster_stats(
        &self,
        points: &[Vec<f64>],
        assignments: &[usize],
        n_clusters: usize,
    ) -> Vec<ClusterStats> {
        self.calculate_dbscan_cluster_stats(points, assignments, n_clusters)
    }

    /// Calculate silhouette score
    fn calculate_silhouette_score(&self, points: &[Vec<f64>], assignments: &[usize]) -> f64 {
        let mut total_silhouette = 0.0;
        let mut count = 0;

        for (i, point) in points.iter().enumerate() {
            let cluster = assignments[i];
            
            // Calculate average distance to points in same cluster
            let mut intra_cluster_dist = 0.0;
            let mut intra_cluster_count = 0;
            
            for (j, other_point) in points.iter().enumerate() {
                if i != j && assignments[j] == cluster {
                    intra_cluster_dist += self.euclidean_distance(point, other_point);
                    intra_cluster_count += 1;
                }
            }
            
            if intra_cluster_count == 0 {
                continue;
            }
            
            intra_cluster_dist /= intra_cluster_count as f64;
            
            // Calculate minimum average distance to other clusters
            let mut min_inter_cluster_dist = f64::INFINITY;
            
            for other_cluster in 0..self.n_clusters {
                if other_cluster == cluster {
                    continue;
                }
                
                let mut inter_cluster_dist = 0.0;
                let mut inter_cluster_count = 0;
                
                for (j, other_point) in points.iter().enumerate() {
                    if assignments[j] == other_cluster {
                        inter_cluster_dist += self.euclidean_distance(point, other_point);
                        inter_cluster_count += 1;
                    }
                }
                
                if inter_cluster_count > 0 {
                    inter_cluster_dist /= inter_cluster_count as f64;
                    if inter_cluster_dist < min_inter_cluster_dist {
                        min_inter_cluster_dist = inter_cluster_dist;
                    }
                }
            }
            
            if min_inter_cluster_dist != f64::INFINITY {
                let silhouette = (min_inter_cluster_dist - intra_cluster_dist) / 
                               min_inter_cluster_dist.max(intra_cluster_dist);
                total_silhouette += silhouette;
                count += 1;
            }
        }
        
        if count > 0 {
            total_silhouette / count as f64
        } else {
            0.0
        }
    }
}