//! Unsupervised clustering bootstrap (cold-start for the tag pipeline).
//!
//! Groups a library's persisted feature vectors so the user can name clusters and seed
//! the first labels before any rules or manual tagging exist. k-means runs over
//! z-standardized features (MFCC and spectral scales differ by orders of magnitude) with
//! deterministic farthest-first seeding, no RNG dependency, reproducible results.
//!
//! Clustering itself writes nothing; `apply_cluster_tag` is what a named cluster does,
//! recording `source = 'cluster'` provenance (sticky, the rules engine never reconciles
//! these away). See [`crate::rules`] for the provenance model.

use crate::db::Database;
use crate::error::Result;
use tracing::instrument;

use super::classify::{FEATURE_VERSION, NUM_FEATURES};

/// One cluster: a representative (medoid) sample plus its members.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Cluster {
    pub id: usize,
    /// Member nearest the centroid, a playable representative for naming.
    pub medoid_hash: String,
    pub member_hashes: Vec<String>,
}

/// Result of clustering a library.
#[derive(Debug, Clone)]
pub struct ClusterResult {
    pub clusters: Vec<Cluster>,
    /// Feature-extractor version the vectors were produced with.
    pub feature_version: u32,
}

/// Suggested cluster count for `n` samples: ~sqrt(n/2), clamped to a sane UI range.
pub fn suggest_k(n: usize) -> usize {
    if n < 2 {
        return n;
    }
    let k = ((n as f64 / 2.0).sqrt()).round() as usize;
    k.clamp(2, 24).min(n)
}

/// Load all current-version feature vectors as `(hash, vector)` pairs.
fn load_vectors(db: &Database) -> Result<Vec<(String, Vec<f64>)>> {
    let mut stmt = db
        .conn()
        .prepare("SELECT hash, vector FROM sample_features WHERE feat_version = ?1")?;
    let rows = stmt.query_map([FEATURE_VERSION], |row| {
        Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
    })?;

    let mut out = Vec::new();
    for r in rows {
        let (hash, json) = r?;
        // A vector that carried a non-finite component was serialized to JSON
        // `null` (serde_json's non-finite behavior), which won't parse back into
        // `Vec<f64>`. Skip such a row rather than `?`-aborting the WHOLE clustering
        // run, one crafted sample must not break clustering for the entire
        // library. Also drop wrong-dimension / any-non-finite vectors, mirroring
        // the sibling guards in exemplar.rs / similarity.rs.
        let Ok(v) = serde_json::from_str::<Vec<f64>>(&json) else {
            continue;
        };
        if v.len() == NUM_FEATURES && v.iter().all(|x| x.is_finite()) {
            out.push((hash, v));
        }
    }
    Ok(out)
}

/// z-standardize each dimension across the dataset (in place). Zero-variance dims
/// collapse to 0.0. Returns nothing, operates on the supplied matrix.
fn standardize(vectors: &mut [Vec<f64>]) {
    if vectors.is_empty() {
        return;
    }
    let dims = vectors[0].len();
    let n = vectors.len() as f64;
    for d in 0..dims {
        let mean = vectors.iter().map(|v| v[d]).sum::<f64>() / n;
        let var = vectors.iter().map(|v| (v[d] - mean).powi(2)).sum::<f64>() / n;
        let std = var.sqrt();
        if std > f64::EPSILON {
            for v in vectors.iter_mut() {
                v[d] = (v[d] - mean) / std;
            }
        } else {
            for v in vectors.iter_mut() {
                v[d] = 0.0;
            }
        }
    }
}

fn sq_dist(a: &[f64], b: &[f64]) -> f64 {
    a.iter().zip(b).map(|(x, y)| (x - y).powi(2)).sum()
}

/// Deterministic farthest-first seeding (Gonzalez): start from the point nearest the
/// dataset mean, then repeatedly take the point maximizing its distance to the chosen
/// set. Ties break by lowest index.
fn init_centroids(points: &[Vec<f64>], k: usize) -> Vec<Vec<f64>> {
    let dims = points[0].len();
    let n = points.len() as f64;
    let mean: Vec<f64> = (0..dims)
        .map(|d| points.iter().map(|p| p[d]).sum::<f64>() / n)
        .collect();

    let first = (0..points.len())
        .min_by(|&a, &b| sq_dist(&points[a], &mean).total_cmp(&sq_dist(&points[b], &mean)))
        .expect("points is non-empty (guaranteed by caller)");

    let mut centroids = vec![points[first].clone()];
    let mut min_d: Vec<f64> = points.iter().map(|p| sq_dist(p, &points[first])).collect();

    while centroids.len() < k {
        let next = (0..points.len())
            .max_by(|&a, &b| min_d[a].total_cmp(&min_d[b]))
            .expect("points is non-empty (guaranteed by caller)");
        centroids.push(points[next].clone());
        for (i, p) in points.iter().enumerate() {
            min_d[i] = min_d[i].min(sq_dist(p, &points[next]));
        }
    }
    centroids
}

/// Cluster the library's feature vectors into `k` groups.
///
/// Returns clusters with member hashes and a medoid representative. `k` is clamped to
/// the number of vectors. Returns an empty result if nothing is analyzed yet.
#[instrument(skip_all)]
pub fn cluster_library(db: &Database, k: usize, max_iters: usize) -> Result<ClusterResult> {
    let loaded = load_vectors(db)?;
    if loaded.is_empty() || k == 0 {
        return Ok(ClusterResult {
            clusters: Vec::new(),
            feature_version: FEATURE_VERSION,
        });
    }
    let hashes: Vec<String> = loaded.iter().map(|(h, _)| h.clone()).collect();
    let mut points: Vec<Vec<f64>> = loaded.into_iter().map(|(_, v)| v).collect();
    standardize(&mut points);

    let k = k.min(points.len());
    let mut centroids = init_centroids(&points, k);
    let mut assign = vec![0usize; points.len()];

    for _ in 0..max_iters.max(1) {
        // Assignment step.
        let mut changed = false;
        for (i, p) in points.iter().enumerate() {
            let best = (0..k)
                .min_by(|&a, &b| sq_dist(p, &centroids[a]).total_cmp(&sq_dist(p, &centroids[b])))
                .expect("k >= 1 (empty/zero-k inputs return early)");
            if best != assign[i] {
                assign[i] = best;
                changed = true;
            }
        }
        if !changed {
            break;
        }
        // Update step. Empty clusters keep their previous centroid (no NaN drift).
        let dims = points[0].len();
        let mut sums = vec![vec![0.0f64; dims]; k];
        let mut counts = vec![0usize; k];
        for (i, p) in points.iter().enumerate() {
            let c = assign[i];
            counts[c] += 1;
            for d in 0..dims {
                sums[c][d] += p[d];
            }
        }
        for c in 0..k {
            if counts[c] > 0 {
                for d in 0..dims {
                    centroids[c][d] = sums[c][d] / counts[c] as f64;
                }
            }
        }
    }

    // Build clusters + medoids.
    let mut members: Vec<Vec<usize>> = vec![Vec::new(); k];
    for (i, &c) in assign.iter().enumerate() {
        members[c].push(i);
    }
    let mut clusters = Vec::new();
    for (c, idxs) in members.into_iter().enumerate() {
        if idxs.is_empty() {
            continue;
        }
        let medoid = *idxs
            .iter()
            .min_by(|&&a, &&b| {
                sq_dist(&points[a], &centroids[c]).total_cmp(&sq_dist(&points[b], &centroids[c]))
            })
            .expect("idxs is non-empty (empty clusters are skipped above)");
        clusters.push(Cluster {
            id: c,
            medoid_hash: hashes[medoid].clone(),
            member_hashes: idxs.into_iter().map(|i| hashes[i].clone()).collect(),
        });
    }

    Ok(ClusterResult {
        clusters,
        feature_version: FEATURE_VERSION,
    })
}

/// Name a cluster: apply `tag` to every member with `source = 'cluster'` provenance.
/// Returns the number of samples that gained the tag.
#[instrument(skip_all)]
pub fn apply_cluster_tag(db: &Database, member_hashes: &[String], tag: &str) -> Result<usize> {
    // One transaction for the whole cluster, not an autocommit per member.
    db.transaction_core(|tx| {
        let mut added = 0;
        for hash in member_hashes {
            if crate::rules::apply_tag_sourced(db, tx, hash, tag, "cluster")? {
                added += 1;
            }
        }
        Ok(added)
    })
}

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

    fn db_with_features(rows: &[(&str, [f64; NUM_FEATURES])]) -> Database {
        let db = Database::open_in_memory().unwrap();
        for (hash, vec) in rows {
            db.conn()
                .execute(
                    "INSERT INTO samples (hash, original_name, file_extension, file_size, import_date, last_modified) \
                     VALUES (?1, ?1, 'wav', 1, 0, 0)",
                    [hash],
                )
                .unwrap();
            let json = serde_json::to_string(&vec.to_vec()).unwrap();
            db.conn()
                .execute(
                    "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
                    rusqlite::params![hash, FEATURE_VERSION, json],
                )
                .unwrap();
        }
        db
    }

    fn vec_at(base: f64) -> [f64; NUM_FEATURES] {
        [base; NUM_FEATURES]
    }

    #[test]
    fn suggest_k_is_sane() {
        assert_eq!(suggest_k(0), 0);
        assert_eq!(suggest_k(1), 1);
        assert!((2..=24).contains(&suggest_k(100)));
        assert_eq!(suggest_k(3), 2); // sqrt(1.5) rounds to 1, clamped up to 2
    }

    #[test]
    fn empty_library_yields_no_clusters() {
        let db = Database::open_in_memory().unwrap();
        let res = cluster_library(&db, 4, 20).unwrap();
        assert!(res.clusters.is_empty());
    }

    #[test]
    fn separates_two_obvious_groups() {
        // Two tight groups far apart in feature space.
        let rows = [
            ("a1", vec_at(0.0)),
            ("a2", vec_at(0.1)),
            ("a3", vec_at(-0.1)),
            ("b1", vec_at(100.0)),
            ("b2", vec_at(100.1)),
            ("b3", vec_at(99.9)),
        ];
        let db = db_with_features(&rows);
        let res = cluster_library(&db, 2, 50).unwrap();
        assert_eq!(res.clusters.len(), 2);

        // Each cluster must be all-a or all-b.
        for cluster in &res.clusters {
            let all_a = cluster.member_hashes.iter().all(|h| h.starts_with('a'));
            let all_b = cluster.member_hashes.iter().all(|h| h.starts_with('b'));
            assert!(
                all_a || all_b,
                "cluster mixed groups: {:?}",
                cluster.member_hashes
            );
            assert_eq!(cluster.member_hashes.len(), 3);
        }
    }

    #[test]
    fn deterministic_across_runs() {
        let rows = [
            ("a1", vec_at(0.0)),
            ("b1", vec_at(50.0)),
            ("c1", vec_at(100.0)),
            ("a2", vec_at(1.0)),
            ("b2", vec_at(51.0)),
        ];
        let db = db_with_features(&rows);
        let r1 = cluster_library(&db, 3, 50).unwrap();
        let r2 = cluster_library(&db, 3, 50).unwrap();
        let sig = |r: &ClusterResult| {
            let mut v: Vec<Vec<String>> = r
                .clusters
                .iter()
                .map(|c| {
                    let mut m = c.member_hashes.clone();
                    m.sort();
                    m
                })
                .collect();
            v.sort();
            v
        };
        assert_eq!(sig(&r1), sig(&r2));
    }

    #[test]
    fn apply_cluster_tag_writes_sticky_provenance() {
        let rows = [("a1", vec_at(0.0)), ("a2", vec_at(0.1))];
        let db = db_with_features(&rows);
        let members = vec!["a1".to_string(), "a2".to_string()];
        assert_eq!(apply_cluster_tag(&db, &members, "kit.my-kicks").unwrap(), 2);

        assert!(
            crate::tags::get_sample_tags(&db, "a1")
                .unwrap()
                .contains(&"kit.my-kicks".to_string())
        );
        let prov = crate::rules::sample_tag_provenance(&db, "a1").unwrap();
        assert_eq!(prov[0].1, "cluster");

        // Removable as a batch (undo the pass).
        let removed = crate::rules::remove_tags_by_source(&db, "cluster").unwrap();
        assert_eq!(removed, 2);
        assert!(crate::tags::get_sample_tags(&db, "a1").unwrap().is_empty());
    }

    #[test]
    fn load_vectors_drops_non_finite_vectors() {
        let mut nan_vec = vec_at(0.5);
        nan_vec[0] = f64::NAN;
        let mut inf_vec = vec_at(0.5);
        inf_vec[3] = f64::INFINITY;
        let db = db_with_features(&[("good", vec_at(0.2)), ("nan", nan_vec), ("inf", inf_vec)]);

        let loaded = load_vectors(&db).unwrap();
        // Only the finite vector survives; NaN/Inf rows are dropped like the
        // sibling `exemplar::is_usable_vector` guard does.
        let hashes: Vec<&str> = loaded.iter().map(|(h, _)| h.as_str()).collect();
        assert_eq!(hashes, vec!["good"]);
    }

    #[test]
    fn clustering_ignores_non_finite_and_does_not_nan() {
        // A poisoned vector must not garbage the medoids: the run completes and
        // only clusters the finite members.
        let mut bad = vec_at(5.0);
        bad[1] = f64::INFINITY;
        let db = db_with_features(&[
            ("a1", vec_at(0.0)),
            ("a2", vec_at(0.1)),
            ("b1", vec_at(9.0)),
            ("bad", bad),
        ]);
        let res = cluster_library(&db, 2, 20).unwrap();
        let clustered: usize = res.clusters.iter().map(|c| c.member_hashes.len()).sum();
        assert_eq!(clustered, 3, "the Inf-poisoned vector should be excluded");
        for c in &res.clusters {
            assert!(!c.member_hashes.iter().any(|m| m == "bad"));
        }
    }
}