// Copyright (c) 2026 Soumyadip Sarkar.
// All rights reserved.
//
// This source code is licensed under the Apache-style license found in the
// LICENSE file in the root directory of this source tree.

use crate::errors::{ClustorError, ClustorResult};
use crate::metrics::{Metric, cosine_distance};
use crate::utils::{compute_row_norms, validate_data_shape};
use std::collections::HashMap;

fn validate(
    data: &[f64],
    n_samples: usize,
    n_features: usize,
    labels: &[i64],
) -> ClustorResult<()> {
    if n_samples == 0 || n_features == 0 {
        return Err(ClustorError::InvalidArg("X must be non-empty".into()));
    }
    validate_data_shape(
        data.len(),
        n_samples,
        n_features,
        "X length mismatch",
        "Input too large",
    )?;
    if labels.len() != n_samples {
        return Err(ClustorError::InvalidArg("labels length mismatch".into()));
    }
    for (i, &lab) in labels.iter().enumerate() {
        if lab < -1 {
            return Err(ClustorError::InvalidArg(
                "labels must be >= -1, where -1 denotes noise".into(),
            ));
        }
        if lab >= 0 && !row_is_finite(data, i, n_features) {
            return Err(ClustorError::InvalidArg(
                "X must contain only finite values for non-noise samples".into(),
            ));
        }
    }
    Ok(())
}

#[inline]
fn row_is_finite(data: &[f64], i: usize, n_features: usize) -> bool {
    data[i * n_features..(i + 1) * n_features]
        .iter()
        .all(|v| v.is_finite())
}

#[inline]
fn dist(
    a: &[f64],
    an: Option<f64>,
    b: &[f64],
    bn: Option<f64>,
    metric: Metric,
    euclidean_inv_scale: f64,
) -> f64 {
    match metric {
        Metric::Euclidean => euclidean_distance_scaled(a, b, euclidean_inv_scale),
        Metric::Cosine => cosine_distance(a, an.unwrap(), b, bn.unwrap()),
    }
}

#[inline]
fn euclidean_distance_scaled(a: &[f64], b: &[f64], inv_scale: f64) -> f64 {
    let mut s = 0.0;
    for i in 0..a.len() {
        let d = a[i] * inv_scale - b[i] * inv_scale;
        s += d * d;
    }
    s.sqrt()
}

#[inline]
fn non_noise_inv_scale_from_labels(
    data: &[f64],
    n_samples: usize,
    n_features: usize,
    labels: &[i64],
) -> f64 {
    let mut max_abs = 0.0;
    for (i, &lab) in labels.iter().take(n_samples).enumerate() {
        if lab < 0 {
            continue;
        }
        let row = &data[i * n_features..(i + 1) * n_features];
        for &v in row {
            let a = v.abs();
            if a > max_abs {
                max_abs = a;
            }
        }
    }
    if max_abs > 0.0 { 1.0 / max_abs } else { 1.0 }
}

#[inline]
fn build_clusters(labels: &[i64]) -> Vec<Vec<usize>> {
    let mut non_noise_count = 0usize;
    let mut min_label = i64::MAX;
    let mut max_label = i64::MIN;

    for &lab in labels {
        if lab < 0 {
            continue;
        }
        non_noise_count += 1;
        if lab < min_label {
            min_label = lab;
        }
        if lab > max_label {
            max_label = lab;
        }
    }

    if non_noise_count == 0 {
        return Vec::new();
    }

    // Fast path for dense, near-contiguous label spaces (even when shifted away from zero).
    // Guarded by checked arithmetic, a density threshold, and fallible reservations so sparse
    // or very large label spaces safely fall back to hashing.
    let dense_span_opt = max_label
        .checked_sub(min_label)
        .and_then(|d| d.checked_add(1))
        .and_then(|span| usize::try_from(span).ok());

    if let Some(dense_span) = dense_span_opt
        && dense_span <= non_noise_count.saturating_mul(4)
    {
        let offset = min_label;
        let mut dense_map = Vec::new();
        if dense_map.try_reserve_exact(dense_span).is_ok() {
            dense_map.resize(dense_span, usize::MAX);
            let mut clusters: Vec<Vec<usize>> = Vec::new();
            let _ = clusters.try_reserve(non_noise_count);

            for (i, &lab) in labels.iter().enumerate() {
                if lab < 0 {
                    continue;
                }
                let dense_idx = (lab - offset) as usize;
                let idx = if dense_map[dense_idx] == usize::MAX {
                    let next = clusters.len();
                    dense_map[dense_idx] = next;
                    clusters.push(Vec::new());
                    next
                } else {
                    dense_map[dense_idx]
                };
                clusters[idx].push(i);
            }

            return clusters;
        }
    }

    let mut label_to_cluster: HashMap<i64, usize> = HashMap::new();
    let _ = label_to_cluster.try_reserve(non_noise_count);
    let mut clusters: Vec<Vec<usize>> = Vec::new();
    let _ = clusters.try_reserve(non_noise_count);

    for (i, &lab) in labels.iter().enumerate() {
        if lab < 0 {
            continue;
        }
        let idx = *label_to_cluster.entry(lab).or_insert_with(|| {
            let idx = clusters.len();
            clusters.push(Vec::new());
            idx
        });
        clusters[idx].push(i);
    }

    clusters
}

/// Silhouette score (mean over samples). Noise label -1 is ignored.
/// O(n^2) implementation intended for moderate n.
pub fn silhouette_score(
    data: &[f64],
    n_samples: usize,
    n_features: usize,
    labels: &[i64],
    metric: Metric,
) -> ClustorResult<f64> {
    validate(data, n_samples, n_features, labels)?;

    let clusters = build_clusters(labels);
    let cluster_sizes: Vec<usize> = clusters.iter().map(|idxs| idxs.len()).collect();
    let k = cluster_sizes.len();
    let m = cluster_sizes.iter().sum::<usize>();
    let mut non_noise_indices: Vec<usize> = Vec::with_capacity(m);
    let mut non_noise_cluster_indices: Vec<usize> = Vec::with_capacity(m);
    let mut euclidean_max_abs = 0.0;

    for (cluster_idx, idxs) in clusters.iter().enumerate() {
        for &i in idxs {
            non_noise_indices.push(i);
            non_noise_cluster_indices.push(cluster_idx);
            if metric == Metric::Euclidean {
                let row = &data[i * n_features..(i + 1) * n_features];
                for &v in row {
                    let a = v.abs();
                    if a > euclidean_max_abs {
                        euclidean_max_abs = a;
                    }
                }
            }
        }
    }

    debug_assert_eq!(non_noise_indices.len(), m);
    debug_assert_eq!(non_noise_cluster_indices.len(), m);
    if k <= 1 {
        return Err(ClustorError::InvalidArg(
            "Need at least 2 clusters (excluding noise) for silhouette_score".into(),
        ));
    }
    let norms = if metric == Metric::Cosine {
        Some(compute_row_norms(data, n_samples, n_features))
    } else {
        None
    };
    let euclidean_inv_scale = if metric == Metric::Euclidean && euclidean_max_abs > 0.0 {
        1.0 / euclidean_max_abs
    } else {
        1.0
    };

    let sum_len = m
        .checked_mul(k)
        .ok_or_else(|| ClustorError::InvalidArg("Input too large".into()))?;
    let mut sums = vec![0.0; sum_len];

    for pos_i in 0..m {
        let i = non_noise_indices[pos_i];
        let ci = non_noise_cluster_indices[pos_i];
        let xi = &data[i * n_features..(i + 1) * n_features];
        let ni = norms.as_ref().map(|v| v[i]);
        for pos_j in (pos_i + 1)..m {
            let j = non_noise_indices[pos_j];
            let cj = non_noise_cluster_indices[pos_j];
            let xj = &data[j * n_features..(j + 1) * n_features];
            let nj = norms.as_ref().map(|v| v[j]);
            let d = dist(xi, ni, xj, nj, metric, euclidean_inv_scale);
            sums[pos_i * k + cj] += d;
            sums[pos_j * k + ci] += d;
        }
    }

    let mut total = 0.0;

    for (pos, _i) in non_noise_indices.iter().enumerate() {
        let ci = non_noise_cluster_indices[pos];
        let size = cluster_sizes[ci];
        if size <= 1 {
            continue;
        }

        let a = sums[pos * k + ci] / ((size - 1) as f64);
        let mut b = f64::INFINITY;
        for c in 0..k {
            if c == ci {
                continue;
            }
            let mean = sums[pos * k + c] / (cluster_sizes[c] as f64);
            if mean < b {
                b = mean;
            }
        }
        let denom = a.max(b).max(1e-12);
        total += (b - a) / denom;
    }

    Ok(total / (m as f64))
}

/// Calinski-Harabasz index (higher is better). Noise (-1) ignored.
pub fn calinski_harabasz_score(
    data: &[f64],
    n_samples: usize,
    n_features: usize,
    labels: &[i64],
) -> ClustorResult<f64> {
    validate(data, n_samples, n_features, labels)?;
    let clusters = build_clusters(labels);
    let k = clusters.len();
    if k <= 1 {
        return Err(ClustorError::InvalidArg(
            "Need at least 2 clusters for calinski_harabasz_score".into(),
        ));
    }

    let n = clusters.iter().map(|v| v.len()).sum::<usize>();
    if n <= k {
        return Err(ClustorError::InvalidArg("Not enough samples".into()));
    }

    let inv_scale = non_noise_inv_scale_from_labels(data, n_samples, n_features, labels);

    let mut overall = vec![0.0; n_features];
    for i in 0..n_samples {
        if labels[i] < 0 {
            continue;
        }
        let x = &data[i * n_features..(i + 1) * n_features];
        for (overall_j, x_j) in overall.iter_mut().zip(x.iter()) {
            *overall_j += x_j;
        }
    }
    for overall_j in overall.iter_mut() {
        *overall_j /= n as f64;
    }

    let mut b = 0.0;
    let mut w = 0.0;
    for idxs in clusters.iter() {
        let mut mu = vec![0.0; n_features];
        for &i in idxs.iter() {
            let x = &data[i * n_features..(i + 1) * n_features];
            for (mu_j, x_j) in mu.iter_mut().zip(x.iter()) {
                *mu_j += x_j;
            }
        }
        for mu_j in mu.iter_mut() {
            *mu_j /= idxs.len() as f64;
        }
        let mut dist2 = 0.0;
        for j in 0..n_features {
            let d = (mu[j] - overall[j]) * inv_scale;
            dist2 += d * d;
        }
        b += (idxs.len() as f64) * dist2;
        for &i in idxs.iter() {
            let x = &data[i * n_features..(i + 1) * n_features];
            let mut s = 0.0;
            for j in 0..n_features {
                let d = (x[j] - mu[j]) * inv_scale;
                s += d * d;
            }
            w += s;
        }
    }
    if w == 0.0 {
        if b == 0.0 {
            return Err(ClustorError::InvalidArg(
                "Calinski-Harabasz score undefined for zero dispersion".into(),
            ));
        }
        return Ok(f64::INFINITY);
    }
    let kf = k as f64;
    let nf = n as f64;
    Ok((b / (kf - 1.0)) / (w / (nf - kf)))
}

/// Davies-Bouldin index (lower is better). Noise (-1) ignored.
pub fn davies_bouldin_score(
    data: &[f64],
    n_samples: usize,
    n_features: usize,
    labels: &[i64],
) -> ClustorResult<f64> {
    validate(data, n_samples, n_features, labels)?;
    let clusters = build_clusters(labels);
    let k = clusters.len();
    if k <= 1 {
        return Err(ClustorError::InvalidArg(
            "Need at least 2 clusters for davies_bouldin_score".into(),
        ));
    }

    let inv_scale = non_noise_inv_scale_from_labels(data, n_samples, n_features, labels);

    let mut centroids: Vec<Vec<f64>> = Vec::with_capacity(k);
    let mut scatters: Vec<f64> = Vec::with_capacity(k);

    for idxs in clusters.iter() {
        let mut mu = vec![0.0; n_features];
        for &i in idxs.iter() {
            let x = &data[i * n_features..(i + 1) * n_features];
            for (mu_j, x_j) in mu.iter_mut().zip(x.iter()) {
                *mu_j += x_j;
            }
        }
        for mu_j in mu.iter_mut() {
            *mu_j /= idxs.len() as f64;
        }
        let mut s = 0.0;
        for &i in idxs.iter() {
            let x = &data[i * n_features..(i + 1) * n_features];
            let mut d2 = 0.0;
            for j in 0..n_features {
                let d = (x[j] - mu[j]) * inv_scale;
                d2 += d * d;
            }
            s += d2.sqrt();
        }
        s /= idxs.len() as f64;
        centroids.push(mu);
        scatters.push(s);
    }

    let mut sum_r = 0.0;
    for i in 0..k {
        let mut max_r = f64::NEG_INFINITY;
        for j in 0..k {
            if i == j {
                continue;
            }
            let mut d2 = 0.0;
            for (&ci, &cj) in centroids[i].iter().zip(centroids[j].iter()) {
                let d = (ci - cj) * inv_scale;
                d2 += d * d;
            }
            let d = d2.sqrt().max(1e-12);
            let r = (scatters[i] + scatters[j]) / d;
            if r > max_r {
                max_r = r;
            }
        }
        sum_r += max_r;
    }
    Ok(sum_r / (k as f64))
}

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

    fn data_with_non_finite_noise() -> (Vec<f64>, Vec<i64>) {
        (
            vec![
                f64::NAN,
                f64::INFINITY,
                0.0,
                0.0,
                0.0,
                1.0,
                5.0,
                0.0,
                5.0,
                1.0,
            ],
            vec![-1, 0, 0, 1, 1],
        )
    }

    fn assert_invalid_arg(res: ClustorResult<f64>, expected_substr: &str) {
        match res {
            Err(ClustorError::InvalidArg(msg)) => assert!(msg.contains(expected_substr)),
            other => panic!("expected InvalidArg, got {other:?}"),
        }
    }

    #[test]
    fn build_clusters_handles_dense_and_sparse_labels() {
        let dense = vec![-1, 2, 0, 2, 1, -1, 0];
        let dense_clusters = build_clusters(&dense);
        assert_eq!(dense_clusters.len(), 3);
        assert_eq!(dense_clusters[0], vec![1, 3]);
        assert_eq!(dense_clusters[1], vec![2, 6]);
        assert_eq!(dense_clusters[2], vec![4]);

        let shifted_dense = vec![-1, 1_000_000, 1_000_002, 1_000_001, 1_000_002, 1_000_000];
        let shifted_dense_clusters = build_clusters(&shifted_dense);
        assert_eq!(shifted_dense_clusters.len(), 3);
        assert_eq!(shifted_dense_clusters[0], vec![1, 5]);
        assert_eq!(shifted_dense_clusters[1], vec![2, 4]);
        assert_eq!(shifted_dense_clusters[2], vec![3]);

        let sparse = vec![-1, 1_000_000_000, 7, 1_000_000_000, 7, -1, 42];
        let sparse_clusters = build_clusters(&sparse);
        assert_eq!(sparse_clusters.len(), 3);
        assert_eq!(sparse_clusters[0], vec![1, 3]);
        assert_eq!(sparse_clusters[1], vec![2, 4]);
        assert_eq!(sparse_clusters[2], vec![6]);
    }

    #[test]
    fn build_clusters_returns_empty_for_all_noise() {
        let labels = vec![-1, -1, -1];
        let clusters = build_clusters(&labels);
        assert!(clusters.is_empty());
    }

    #[test]
    fn silhouette_score_supports_cosine_metric() {
        let data = vec![
            1.0, 0.0, // c0
            0.9, 0.1, // c0
            0.0, 1.0, // c1
            0.1, 0.9, // c1
        ];
        let labels = vec![0, 0, 1, 1];

        let score = silhouette_score(&data, 4, 2, &labels, Metric::Cosine).unwrap();
        assert!(score.is_finite());
        assert!((-1.0..=1.0).contains(&score));
    }

    #[test]
    fn calinski_harabasz_errors_when_total_dispersion_is_zero() {
        let data = vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
        let labels = vec![0, 0, 1];

        assert_invalid_arg(
            calinski_harabasz_score(&data, 3, 2, &labels),
            "undefined for zero dispersion",
        );
    }

    #[test]
    fn calinski_harabasz_returns_infinity_when_within_cluster_dispersion_is_zero() {
        let data = vec![0.0, 0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0];
        let labels = vec![0, 0, 1, 1];

        let score = calinski_harabasz_score(&data, 4, 2, &labels).unwrap();
        assert!(score.is_infinite());
        assert!(score.is_sign_positive());
    }

    #[test]
    fn davies_bouldin_handles_identical_centroids() {
        let data = vec![0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0];
        let labels = vec![0, 0, 1, 1];

        let score = davies_bouldin_score(&data, 4, 2, &labels).unwrap();
        assert!(score.is_finite());
        assert!(score >= 0.0);
    }

    #[test]
    fn validation_metrics_ignore_non_finite_noise_rows() {
        let (data, labels) = data_with_non_finite_noise();

        let sil = silhouette_score(&data, 5, 2, &labels, Metric::Euclidean).unwrap();
        let ch = calinski_harabasz_score(&data, 5, 2, &labels).unwrap();
        let db = davies_bouldin_score(&data, 5, 2, &labels).unwrap();

        assert!(sil.is_finite());
        assert!((-1.0..=1.0).contains(&sil));
        assert!(ch.is_finite());
        assert!(db.is_finite());
    }

    #[test]
    fn validation_metrics_reject_non_finite_non_noise_rows() {
        let data = vec![0.0, 0.0, f64::NAN, 1.0, 2.0, f64::INFINITY];
        let labels = vec![0, 0, 1];

        assert_invalid_arg(
            silhouette_score(&data, 3, 2, &labels, Metric::Euclidean),
            "finite",
        );
        assert_invalid_arg(calinski_harabasz_score(&data, 3, 2, &labels), "finite");
        assert_invalid_arg(davies_bouldin_score(&data, 3, 2, &labels), "finite");
    }

    #[test]
    fn validation_rejects_invalid_negative_labels() {
        let data = vec![0.0, 0.0, 1.0, 1.0];
        let labels = vec![-2, 0];

        assert_invalid_arg(
            silhouette_score(&data, 2, 2, &labels, Metric::Euclidean),
            "labels must be >= -1",
        );
        assert_invalid_arg(
            calinski_harabasz_score(&data, 2, 2, &labels),
            "labels must be >= -1",
        );
        assert_invalid_arg(
            davies_bouldin_score(&data, 2, 2, &labels),
            "labels must be >= -1",
        );
    }

    #[test]
    fn validation_rejects_overflowing_shape_product() {
        assert_invalid_arg(
            silhouette_score(&[], usize::MAX, 2, &[], Metric::Euclidean),
            "Input too large",
        );
        assert_invalid_arg(
            calinski_harabasz_score(&[], usize::MAX, 2, &[]),
            "Input too large",
        );
        assert_invalid_arg(
            davies_bouldin_score(&[], usize::MAX, 2, &[]),
            "Input too large",
        );
    }

    #[test]
    fn validation_metrics_reject_all_noise_labels() {
        let data = vec![0.0, 0.0, 1.0, 1.0, 2.0, 2.0];
        let labels = vec![-1, -1, -1];

        assert_invalid_arg(
            silhouette_score(&data, 3, 2, &labels, Metric::Euclidean),
            "at least 2 clusters",
        );
        assert_invalid_arg(
            calinski_harabasz_score(&data, 3, 2, &labels),
            "at least 2 clusters",
        );
        assert_invalid_arg(
            davies_bouldin_score(&data, 3, 2, &labels),
            "at least 2 clusters",
        );
    }

    #[test]
    fn validation_metrics_reject_all_noise_labels_with_non_finite_rows() {
        let data = vec![f64::NAN, f64::INFINITY, 1.0, 1.0, -f64::INFINITY, 2.0];
        let labels = vec![-1, -1, -1];

        assert_invalid_arg(
            silhouette_score(&data, 3, 2, &labels, Metric::Euclidean),
            "at least 2 clusters",
        );
        assert_invalid_arg(
            calinski_harabasz_score(&data, 3, 2, &labels),
            "at least 2 clusters",
        );
        assert_invalid_arg(
            davies_bouldin_score(&data, 3, 2, &labels),
            "at least 2 clusters",
        );
    }

    #[test]
    fn validation_metrics_handle_large_finite_values_without_overflow() {
        let data = vec![0.0, 0.0, 1e308, 1e308, -1e308, -1e308];
        let labels = vec![0, 1, 1];

        let sil = silhouette_score(&data, 3, 2, &labels, Metric::Euclidean).unwrap();
        let db = davies_bouldin_score(&data, 3, 2, &labels).unwrap();
        let ch = calinski_harabasz_score(&data, 3, 2, &labels).unwrap();

        assert!(sil.is_finite());
        assert!((-1.0..=1.0).contains(&sil));
        assert!(db.is_finite());
        assert!(ch.is_finite());
    }

    #[test]
    fn validation_metric_scores_are_invariant_to_label_values() {
        let data = vec![0.0, 0.0, 0.0, 1.0, 10.0, 10.0, 10.0, 11.0, 50.0, 50.0];
        let labels_a = vec![5, 5, 9, 9, -1];
        let labels_b = vec![42, 42, 7, 7, -1];

        let sil_a = silhouette_score(&data, 5, 2, &labels_a, Metric::Euclidean).unwrap();
        let sil_b = silhouette_score(&data, 5, 2, &labels_b, Metric::Euclidean).unwrap();
        let ch_a = calinski_harabasz_score(&data, 5, 2, &labels_a).unwrap();
        let ch_b = calinski_harabasz_score(&data, 5, 2, &labels_b).unwrap();
        let db_a = davies_bouldin_score(&data, 5, 2, &labels_a).unwrap();
        let db_b = davies_bouldin_score(&data, 5, 2, &labels_b).unwrap();

        assert!((sil_a - sil_b).abs() <= 1e-12);
        assert!((ch_a - ch_b).abs() <= 1e-12);
        assert!((db_a - db_b).abs() <= 1e-12);
    }
}
