fdars-core 0.28.0

//! Advanced functional clustering algorithms.
//!
//! This module provides four advanced functional clustering paradigms beyond the
//! basic k-means and fuzzy c-means found in [`clustering`](crate::clustering):
//!
//! - **DBSCAN** ([`dbscan_fd`]): Density-based clustering over precomputed functional
//!   L2 distances. Discovers clusters of arbitrary shape and flags noise curves as
//!   [`None`] in the assignment vector. No need to specify k.
//!
//! - **kCFC** ([`kcfc_cluster`]): K-means-style assignment loop where each cluster's
//!   centroid is replaced by a per-cluster FPCA model. A curve is assigned to the
//!   cluster whose FPCA basis produces the smallest reconstruction error.
//!
//! - **funFEM** and **Align-and-Cluster**: Discriminative-subspace GMM clustering and
//!   elastic joint clustering — implemented in plan 33-03.
//!
//! All algorithms are strictly additive relative to the existing `clustering` module.
//! No existing public signature is modified.

use crate::distance::l2_distance_matrix;
use crate::error::FdarError;
use crate::helpers::simpsons_weights;
use crate::matrix::FdMatrix;
use crate::regression::{fdata_to_pc_1d, FpcaResult};
use rand::prelude::*;

// ────────────────────────────────────────────────────────────────────────────
// DBSCAN over functional L2 distances
// ────────────────────────────────────────────────────────────────────────────

/// Configuration for DBSCAN density clustering over functional data.
///
/// DBSCAN discovers clusters of arbitrary shape by expanding dense
/// regions of curves in functional L2 distance space. Curves that
/// do not belong to any dense region are flagged as noise ([`None`]
/// in [`DbscanResult::cluster`]).
///
/// ## Distance units
///
/// `eps` is in the same units as `l2_distance_matrix` — that is,
/// functional L2 distance with Simpson's-rule integration weights.
/// For a constant-1 curve on `argvals` spanning \[0, 1\], the L2 norm
/// is ≈ 1.0. As a practical starting point, set `eps` to a fraction
/// (e.g. 0.1–0.5) of the dataset's median pairwise L2 distance:
///
/// ```
/// # use fdars_core::distance::l2_distance_matrix;
/// # use fdars_core::matrix::FdMatrix;
/// # let data = FdMatrix::zeros(10, 20);
/// # let argvals: Vec<f64> = (0..20).map(|i| i as f64 / 19.0).collect();
/// let dist = l2_distance_matrix(&data, &argvals);
/// let n = data.nrows();
/// let mut upper: Vec<f64> = Vec::new();
/// for i in 0..n {
///     for j in (i + 1)..n {
///         upper.push(dist[(i, j)]);
///     }
/// }
/// upper.sort_by(|a, b| a.partial_cmp(b).unwrap());
/// let median_dist = upper[upper.len() / 2];
/// let eps = 0.3 * median_dist; // start here, tune as needed
/// let _ = eps;
/// ```
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct DbscanConfig {
    /// Neighbourhood radius in functional L2 distance units (default: 0.5).
    ///
    /// Must be strictly positive (`eps > 0`). A tiny positive value makes
    /// every curve a noise point; a very large value merges everything into
    /// one cluster.
    pub eps: f64,
    /// Minimum number of curves (including the point itself) required in the
    /// `eps`-neighbourhood for a point to be considered a core point (default: 3).
    ///
    /// Must be ≥ 1.
    pub min_points: usize,
}

impl Default for DbscanConfig {
    fn default() -> Self {
        Self {
            eps: 0.5,
            min_points: 3,
        }
    }
}

/// Result of DBSCAN density clustering over functional data.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct DbscanResult {
    /// Cluster assignment for each curve.
    ///
    /// `None` = noise (not part of any dense cluster).
    /// `Some(c)` = cluster index `c` (0-based, contiguous).
    pub cluster: Vec<Option<usize>>,
    /// Number of discovered clusters (excludes noise points).
    pub n_clusters: usize,
    /// Number of noise points (curves assigned `None`).
    pub n_noise: usize,
    /// Precomputed n × n pairwise L2 distance matrix used internally.
    pub distances: FdMatrix,
}

/// Run DBSCAN over functional L2 distances.
///
/// Discovers arbitrarily-shaped clusters in functional data by expanding
/// dense regions. Curves not reachable from any core point are labelled
/// noise (`None`).
///
/// # Arguments
///
/// * `data` — Functional data matrix (n × m, column-major).
/// * `argvals` — Evaluation grid (length m).
/// * `config` — Algorithm parameters; see [`DbscanConfig`].
///
/// # Errors
///
/// Returns [`FdarError::InvalidDimension`] if `n == 0`, `m == 0`, or
/// `argvals.len() != m`.  Returns [`FdarError::InvalidParameter`] if
/// `config.eps <= 0` or `config.min_points == 0`.
///
/// # Examples
///
/// ```
/// use fdars_core::clustering_advanced::{dbscan_fd, DbscanConfig};
/// use fdars_core::matrix::FdMatrix;
/// use std::f64::consts::PI;
///
/// // Two tight clusters of 5 sin curves each, well separated
/// let m = 30;
/// let n = 10;
/// let t: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64).collect();
/// let mut col_major = vec![0.0_f64; n * m];
/// for i in 0..5 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).sin();
///     }
/// }
/// for i in 5..10 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).sin() + 5.0;
///     }
/// }
/// let data = FdMatrix::from_column_major(col_major, n, m).unwrap();
///
/// let mut cfg = DbscanConfig::default();
/// cfg.eps = 1.0;
/// cfg.min_points = 2;
/// let result = dbscan_fd(&data, &t, &cfg).unwrap();
/// assert_eq!(result.n_clusters, 2);
/// assert_eq!(result.n_noise, 0);
/// ```
#[must_use = "expensive computation whose result should not be discarded"]
pub fn dbscan_fd(
    data: &FdMatrix,
    argvals: &[f64],
    config: &DbscanConfig,
) -> Result<DbscanResult, FdarError> {
    let (n, m) = data.shape();

    // Validation
    if n == 0 || m == 0 {
        return Err(FdarError::InvalidDimension {
            parameter: "data",
            expected: "at least 1 row and 1 column".to_string(),
            actual: format!("{n} rows, {m} columns"),
        });
    }
    if argvals.len() != m {
        return Err(FdarError::InvalidDimension {
            parameter: "argvals",
            expected: format!("{m}"),
            actual: format!("{}", argvals.len()),
        });
    }
    if config.eps <= 0.0 {
        return Err(FdarError::InvalidParameter {
            parameter: "eps",
            message: format!("eps must be > 0, got {}", config.eps),
        });
    }
    if config.min_points == 0 {
        return Err(FdarError::InvalidParameter {
            parameter: "min_points",
            message: "min_points must be >= 1".to_string(),
        });
    }

    let dist = l2_distance_matrix(data, argvals);

    // Standard DBSCAN
    // labels: None = unvisited/noise, Some(c) = cluster c
    let mut labels: Vec<Option<usize>> = vec![None; n];
    let mut visited: Vec<bool> = vec![false; n];
    let mut cluster_id: usize = 0;

    for i in 0..n {
        if visited[i] {
            continue;
        }
        visited[i] = true;

        // Compute eps-neighbourhood of i (excluding i itself)
        let neighbors: Vec<usize> = (0..n)
            .filter(|&j| j != i && dist[(i, j)] <= config.eps)
            .collect();

        // Core-point rule: i + its neighbors must reach min_points
        if neighbors.len() + 1 < config.min_points {
            // Leave as noise for now (may be absorbed as border point later)
            continue;
        }

        // i is a core point — start a new cluster
        labels[i] = Some(cluster_id);

        // BFS expansion
        let mut queue = neighbors.clone();
        let mut qi = 0;
        while qi < queue.len() {
            let j = queue[qi];
            qi += 1;

            if !visited[j] {
                visited[j] = true;
                let j_neighbors: Vec<usize> = (0..n)
                    .filter(|&k| k != j && dist[(j, k)] <= config.eps)
                    .collect();
                if j_neighbors.len() + 1 >= config.min_points {
                    // j is also a core point — add its unqueued neighbours
                    for nb in j_neighbors {
                        if !queue.contains(&nb) {
                            queue.push(nb);
                        }
                    }
                }
            }

            // Absorb j into cluster if not yet assigned
            if labels[j].is_none() {
                labels[j] = Some(cluster_id);
            }
        }

        cluster_id += 1;
    }

    let n_clusters = cluster_id;
    let n_noise = labels.iter().filter(|l| l.is_none()).count();

    Ok(DbscanResult {
        cluster: labels,
        n_clusters,
        n_noise,
        distances: dist,
    })
}

// ────────────────────────────────────────────────────────────────────────────
// kCFC: per-cluster FPCA reassignment loop
// ────────────────────────────────────────────────────────────────────────────

/// Configuration for kCFC (k-means-like clustering via Functional Components).
///
/// Each cluster is represented by a per-cluster FPCA model rather than a
/// centroid. A curve is assigned to the cluster whose FPCA basis gives the
/// smallest L2 reconstruction error.
///
/// **Reference:** Chiou & Li (2007), "Functional clustering and identifying
/// substructures of longitudinal data." R baseline: `fdapace::kCFC`.
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct KcfcConfig {
    /// Number of clusters (default: 2). Must be ≥ 1 and ≤ n.
    pub k: usize,
    /// Number of per-cluster FPC components to use for reconstruction (default: 3).
    ///
    /// Clamped internally to `min(n_k, m)` where `n_k` is the cluster size.
    pub ncomp: usize,
    /// Maximum number of reassignment iterations (default: 50).
    pub max_iter: usize,
    /// Random seed for k-means++ initialization (default: 42).
    pub seed: u64,
}

impl Default for KcfcConfig {
    fn default() -> Self {
        Self {
            k: 2,
            ncomp: 3,
            max_iter: 50,
            seed: 42,
        }
    }
}

/// Result of kCFC per-cluster FPCA clustering.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct KcfcResult {
    /// Cluster assignment for each curve (0-based, contiguous, length n).
    pub cluster: Vec<usize>,
    /// Per-cluster FPCA models.
    ///
    /// `fpca_models[k]` is `None` if cluster k was empty at convergence.
    pub fpca_models: Vec<Option<FpcaResult>>,
    /// Reconstruction error matrix (n × k).
    ///
    /// `reconstruction_errors[(i, k)]` is the L2 squared reconstruction error
    /// of curve i against cluster k's FPCA model.
    pub reconstruction_errors: FdMatrix,
    /// Number of iterations performed.
    pub iterations: usize,
    /// Whether the algorithm converged (no label changes in the last iteration).
    pub converged: bool,
}

/// Cluster functional data using per-cluster FPCA reconstruction errors (kCFC).
///
/// Initialises with k-means++ hard labels, then iterates: fit a per-cluster
/// FPCA model for each cluster, compute the reconstruction error of every
/// curve against every cluster's model, and reassign each curve to the cluster
/// with the smallest error. Repeats until convergence or `config.max_iter`.
///
/// # Arguments
///
/// * `data` — Functional data matrix (n × m, column-major).
/// * `argvals` — Evaluation grid (length m).
/// * `config` — Algorithm parameters; see [`KcfcConfig`].
///
/// # Errors
///
/// Returns [`FdarError::InvalidDimension`] if `n == 0`, `m == 0`, or
/// `argvals.len() != m`.  Returns [`FdarError::InvalidParameter`] if
/// `config.k == 0` or `config.k > n`.
///
/// # Examples
///
/// ```
/// use fdars_core::clustering_advanced::{kcfc_cluster, KcfcConfig};
/// use fdars_core::matrix::FdMatrix;
/// use std::f64::consts::PI;
///
/// let m = 30;
/// let n = 20;  // 10 per cluster
/// let t: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64).collect();
/// let mut col_major = vec![0.0_f64; n * m];
/// for i in 0..10 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).sin();
///     }
/// }
/// for i in 10..20 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).cos() + 5.0;
///     }
/// }
/// let data = FdMatrix::from_column_major(col_major, n, m).unwrap();
///
/// let mut cfg = KcfcConfig::default();
/// cfg.k = 2;
/// cfg.ncomp = 2;
/// let result = kcfc_cluster(&data, &t, &cfg).unwrap();
/// assert_eq!(result.cluster.len(), n);
/// ```
#[must_use = "expensive computation whose result should not be discarded"]
pub fn kcfc_cluster(
    data: &FdMatrix,
    argvals: &[f64],
    config: &KcfcConfig,
) -> Result<KcfcResult, FdarError> {
    let (n, m) = data.shape();

    // Validation
    if n == 0 || m == 0 {
        return Err(FdarError::InvalidDimension {
            parameter: "data",
            expected: "at least 1 row and 1 column".to_string(),
            actual: format!("{n} rows, {m} columns"),
        });
    }
    if argvals.len() != m {
        return Err(FdarError::InvalidDimension {
            parameter: "argvals",
            expected: format!("{m}"),
            actual: format!("{}", argvals.len()),
        });
    }
    if config.k == 0 {
        return Err(FdarError::InvalidParameter {
            parameter: "k",
            message: "k must be >= 1".to_string(),
        });
    }
    if config.k > n {
        return Err(FdarError::InvalidParameter {
            parameter: "k",
            message: format!("k={} exceeds number of curves n={}", config.k, n),
        });
    }
    if config.ncomp == 0 {
        return Err(FdarError::InvalidParameter {
            parameter: "ncomp",
            message: "ncomp must be >= 1".to_string(),
        });
    }

    let k = config.k;
    let weights = simpsons_weights(argvals);

    // ── k-means++ initialisation over raw row curves ──────────────────────
    let row_major = data.to_row_major(); // n * m, row-major buffer
    let mut rng = StdRng::seed_from_u64(config.seed);

    // Select first center uniformly
    let mut center_indices: Vec<usize> = Vec::with_capacity(k);
    center_indices.push(rng.gen_range(0..n));

    // Maintain min-distance-squared to nearest chosen center
    let mut min_dist_sq: Vec<f64> = (0..n)
        .map(|i| {
            let c0 = center_indices[0];
            let d = l2_dist_rowmajor(&row_major, i, c0, m, &weights);
            d * d
        })
        .collect();

    while center_indices.len() < k {
        // Sample proportional to D^2
        let total: f64 = min_dist_sq.iter().sum();
        let chosen = if total < 1e-15 {
            rng.gen_range(0..n)
        } else {
            let r = rng.gen::<f64>() * total;
            let mut cumsum = 0.0;
            let mut sel = n - 1;
            for (i, &d) in min_dist_sq.iter().enumerate() {
                cumsum += d;
                if cumsum >= r {
                    sel = i;
                    break;
                }
            }
            sel
        };
        center_indices.push(chosen);

        // Update min_dist_sq with distance to new center
        for i in 0..n {
            let d = l2_dist_rowmajor(&row_major, i, chosen, m, &weights);
            let d2 = d * d;
            if d2 < min_dist_sq[i] {
                min_dist_sq[i] = d2;
            }
        }
    }

    // Initial assignment: assign each curve to nearest center
    let mut cluster: Vec<usize> = (0..n)
        .map(|i| {
            center_indices
                .iter()
                .enumerate()
                .min_by(|(_, &c1), (_, &c2)| {
                    let d1 = l2_dist_rowmajor(&row_major, i, c1, m, &weights);
                    let d2 = l2_dist_rowmajor(&row_major, i, c2, m, &weights);
                    d1.partial_cmp(&d2).unwrap_or(std::cmp::Ordering::Equal)
                })
                .map(|(ki, _)| ki)
                .unwrap_or(0)
        })
        .collect();

    // ── Reassignment loop ─────────────────────────────────────────────────
    let mut fpca_models: Vec<Option<FpcaResult>> = vec![None; k];
    let mut reconstruction_errors = FdMatrix::zeros(n, k);
    let mut converged = false;
    let mut iterations = 0;

    for _iter in 0..config.max_iter {
        iterations += 1;

        // ── Fit per-cluster FPCA models ───────────────────────────────────
        for ki in 0..k {
            let member_indices: Vec<usize> = (0..n).filter(|&i| cluster[i] == ki).collect();

            if member_indices.is_empty() {
                // Keep previous model (or None on first iteration)
                continue;
            }

            // Gather member rows into a (n_k x m) FdMatrix
            let n_k = member_indices.len();
            let mut col_major_k = vec![0.0_f64; n_k * m];
            for (row_in_k, &orig_i) in member_indices.iter().enumerate() {
                for j in 0..m {
                    col_major_k[row_in_k + j * n_k] = data[(orig_i, j)];
                }
            }
            let data_k = FdMatrix::from_column_major(col_major_k, n_k, m)?;

            // Fit FPCA (ncomp clamped internally to min(n_k, m))
            match fdata_to_pc_1d(&data_k, config.ncomp, argvals) {
                Ok(fpca) => {
                    fpca_models[ki] = Some(fpca);
                }
                Err(_) => {
                    // Degenerate cluster; keep prior model
                }
            }
        }

        // ── Compute reconstruction errors for all curves vs all clusters ──
        for i in 0..n {
            let curve_row = data.row(i);
            let curve_mat = FdMatrix::from_slice(&curve_row, 1, m)?;

            for ki in 0..k {
                let err = match &fpca_models[ki] {
                    None => f64::INFINITY,
                    Some(fpca) => {
                        let ncomp_eff = fpca.rotation.ncols();
                        match fpca.project(&curve_mat) {
                            Ok(scores) => {
                                match fpca.reconstruct(&scores, ncomp_eff) {
                                    Ok(recon) => {
                                        // L2^2 reconstruction error with Simpson weights
                                        let mut err_sq = 0.0;
                                        for j in 0..m {
                                            let diff = curve_row[j] - recon[(0, j)];
                                            err_sq += diff * diff * weights[j];
                                        }
                                        err_sq
                                    }
                                    Err(_) => f64::INFINITY,
                                }
                            }
                            Err(_) => f64::INFINITY,
                        }
                    }
                };
                reconstruction_errors[(i, ki)] = err;
            }
        }

        // ── Reassign each curve to the cluster with minimum error ─────────
        let mut changed = false;
        for i in 0..n {
            let best_k = (0..k)
                .min_by(|&a, &b| {
                    reconstruction_errors[(i, a)]
                        .partial_cmp(&reconstruction_errors[(i, b)])
                        .unwrap_or(std::cmp::Ordering::Equal)
                })
                .unwrap_or(0);
            if cluster[i] != best_k {
                cluster[i] = best_k;
                changed = true;
            }
        }

        if !changed {
            converged = true;
            break;
        }
    }

    Ok(KcfcResult {
        cluster,
        fpca_models,
        reconstruction_errors,
        iterations,
        converged,
    })
}

// ────────────────────────────────────────────────────────────────────────────
// funFEM: Discriminative-subspace (Fisher-EM) functional clustering
// ────────────────────────────────────────────────────────────────────────────

/// Configuration for funFEM discriminative-subspace clustering.
///
/// funFEM applies the Fisher-EM algorithm to functional data: it first extracts
/// global FPC scores via [`fdata_to_pc_1d`](crate::regression::fdata_to_pc_1d),
/// then alternates between finding a discriminative subspace (maximising
/// between-class vs within-class scatter) and running a GMM E/M step in that
/// subspace.
///
/// ## Divergence from R `funFEM`
///
/// This is a simplified Fisher-EM implementation. The discriminative subspace is
/// found by computing W\_soft^{-1} B\_soft via Cholesky inversion followed by
/// an SVD (instead of a proper generalized-eigenvalue solver), because no
/// generalized-eigenvalue crate is used. The multi-pass outer loop re-estimates
/// the subspace at each iteration. This diverges from the iterative schedule in
/// the original paper (Bouveyron & Brunet, 2014) but gives practical cluster
/// recovery on well-separated functional data.
///
/// ## Example
///
/// ```
/// use fdars_core::clustering_advanced::{funfem_cluster, FunFemConfig};
/// use fdars_core::matrix::FdMatrix;
/// use std::f64::consts::PI;
///
/// let m = 30;
/// let n = 12;
/// let t: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64).collect();
/// let mut col_major = vec![0.0_f64; n * m];
/// for i in 0..6 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).sin();
///     }
/// }
/// for i in 6..12 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).sin() + 5.0;
///     }
/// }
/// let data = FdMatrix::from_column_major(col_major, n, m).unwrap();
///
/// let mut cfg = FunFemConfig::default();
/// cfg.k = 2;
/// cfg.ncomp = 4;
/// let result = funfem_cluster(&data, &t, &cfg).unwrap();
/// assert_eq!(result.cluster.len(), n);
/// ```
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct FunFemConfig {
    /// Number of clusters (default: 2). Must be ≥ 1 and ≤ n.
    pub k: usize,
    /// Number of global FPC components for the score space (default: 10).
    /// Clamped internally to `min(n, m)`.
    pub ncomp: usize,
    /// Discriminative subspace dimension (default: 0 = auto = min(k-1, ncomp_eff)).
    /// Clamped to `ncomp_eff` if larger.
    pub p_disc: usize,
    /// Maximum outer Fisher-EM iterations (default: 50).
    pub max_iter: usize,
    /// Log-likelihood convergence tolerance (default: 1e-6).
    pub tol: f64,
    /// Random seed for k-means++ initialization (default: 42).
    pub seed: u64,
}

impl Default for FunFemConfig {
    fn default() -> Self {
        Self {
            k: 2,
            ncomp: 10,
            p_disc: 0,
            max_iter: 50,
            tol: 1e-6,
            seed: 42,
        }
    }
}

/// Result of funFEM discriminative-subspace clustering.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct FunFemResult {
    /// Cluster assignment for each curve (0-based, contiguous, length n).
    pub cluster: Vec<usize>,
    /// Soft membership matrix (n × k).
    pub membership: FdMatrix,
    /// Discriminative directions (ncomp_eff × p_disc_eff).
    pub disc_subspace: FdMatrix,
    /// Final log-likelihood.
    pub log_likelihood: f64,
    /// Number of outer iterations performed.
    pub iterations: usize,
    /// Whether the algorithm converged.
    pub converged: bool,
}

/// Cluster functional data using the Fisher-EM discriminative-subspace GMM (funFEM).
///
/// Extracts global FPC scores, then iterates between estimating a discriminative
/// subspace and running a GMM E/M step within it. Returns cluster assignments
/// and the discriminative directions used.
///
/// # Arguments
///
/// * `data` — Functional data matrix (n × m, column-major).
/// * `argvals` — Evaluation grid (length m).
/// * `config` — Algorithm parameters; see [`FunFemConfig`].
///
/// # Errors
///
/// Returns [`FdarError::InvalidDimension`] if `n == 0`, `m == 0`, or
/// `argvals.len() != m`. Returns [`FdarError::InvalidParameter`] if
/// `config.k == 0`, `config.k > n`, or `config.ncomp == 0`.
#[must_use = "expensive computation whose result should not be discarded"]
pub fn funfem_cluster(
    data: &FdMatrix,
    argvals: &[f64],
    config: &FunFemConfig,
) -> Result<FunFemResult, FdarError> {
    let (n, m) = data.shape();

    // Validation
    if n == 0 || m == 0 {
        return Err(FdarError::InvalidDimension {
            parameter: "data",
            expected: "at least 1 row and 1 column".to_string(),
            actual: format!("{n} rows, {m} columns"),
        });
    }
    if argvals.len() != m {
        return Err(FdarError::InvalidDimension {
            parameter: "argvals",
            expected: format!("{m}"),
            actual: format!("{}", argvals.len()),
        });
    }
    if config.k == 0 {
        return Err(FdarError::InvalidParameter {
            parameter: "k",
            message: "k must be >= 1".to_string(),
        });
    }
    if config.k > n {
        return Err(FdarError::InvalidParameter {
            parameter: "k",
            message: format!("k={} exceeds number of curves n={}", config.k, n),
        });
    }
    if config.ncomp == 0 {
        return Err(FdarError::InvalidParameter {
            parameter: "ncomp",
            message: "ncomp must be >= 1".to_string(),
        });
    }

    let k = config.k;

    // Step 1: Global FPC scores via fdata_to_pc_1d (clamps ncomp to min(n,m))
    let fpca = fdata_to_pc_1d(data, config.ncomp, argvals)?;
    // scores: n x ncomp_eff (row-major semantics via FdMatrix indexing)
    let scores = &fpca.scores; // FdMatrix n x ncomp_eff
    let ncomp_eff = scores.ncols();

    // Effective discriminative dimension
    let p_disc_eff = if config.p_disc == 0 {
        (k - 1).max(1).min(ncomp_eff)
    } else {
        config.p_disc.min(ncomp_eff)
    };

    // Step 2: K-means++ initialization on the score rows
    let weights_uniform = vec![1.0; ncomp_eff]; // uniform weight for Euclidean distance in score space
    let row_major_scores = scores.to_row_major(); // n * ncomp_eff
    let mut rng = StdRng::seed_from_u64(config.seed);

    let mut center_indices: Vec<usize> = Vec::with_capacity(k);
    center_indices.push(rng.gen_range(0..n));
    let mut min_dist_sq: Vec<f64> = (0..n)
        .map(|i| {
            let c0 = center_indices[0];
            let d = l2_dist_rowmajor(&row_major_scores, i, c0, ncomp_eff, &weights_uniform);
            d * d
        })
        .collect();
    while center_indices.len() < k {
        let total: f64 = min_dist_sq.iter().sum();
        let chosen = if total < 1e-15 {
            rng.gen_range(0..n)
        } else {
            let r = rng.gen::<f64>() * total;
            let mut cumsum = 0.0;
            let mut sel = n - 1;
            for (i, &d) in min_dist_sq.iter().enumerate() {
                cumsum += d;
                if cumsum >= r {
                    sel = i;
                    break;
                }
            }
            sel
        };
        center_indices.push(chosen);
        for i in 0..n {
            let d = l2_dist_rowmajor(&row_major_scores, i, chosen, ncomp_eff, &weights_uniform);
            let d2 = d * d;
            if d2 < min_dist_sq[i] {
                min_dist_sq[i] = d2;
            }
        }
    }

    // Initial hard labels from nearest center
    let mut cluster: Vec<usize> = (0..n)
        .map(|i| {
            center_indices
                .iter()
                .enumerate()
                .min_by(|(_, &c1), (_, &c2)| {
                    let d1 =
                        l2_dist_rowmajor(&row_major_scores, i, c1, ncomp_eff, &weights_uniform);
                    let d2 =
                        l2_dist_rowmajor(&row_major_scores, i, c2, ncomp_eff, &weights_uniform);
                    d1.partial_cmp(&d2).unwrap_or(std::cmp::Ordering::Equal)
                })
                .map(|(ki, _)| ki)
                .unwrap_or(0)
        })
        .collect();

    // Initial uniform cluster weights and diagonal covariances in score space
    let mut pi: Vec<f64> = vec![1.0 / k as f64; k];
    // Per-cluster means in score space (ncomp_eff-dim)
    let mut mu_k: Vec<Vec<f64>> = vec![vec![0.0; ncomp_eff]; k];
    // Per-cluster diagonal variance in score space
    let mut sigma_k: Vec<Vec<f64>> = vec![vec![1.0; ncomp_eff]; k];

    // Initialize mu_k from initial assignment
    update_gmm_params_from_hard(
        &row_major_scores,
        &cluster,
        k,
        ncomp_eff,
        &mut pi,
        &mut mu_k,
        &mut sigma_k,
    );

    // Discriminative subspace directions (ncomp_eff x p_disc_eff) — identity init
    let mut disc_dirs: Vec<f64> = {
        let mut v = vec![0.0_f64; ncomp_eff * p_disc_eff];
        for d in 0..p_disc_eff {
            if d < ncomp_eff {
                v[d + d * ncomp_eff] = 1.0; // column d = e_d
            }
        }
        v
    };

    let mut prev_ll = f64::NEG_INFINITY;
    // Initialize responsibilities from hard cluster assignment
    let mut resp = vec![0.0_f64; n * k];
    for i in 0..n {
        let ki = cluster[i].min(k - 1);
        resp[i * k + ki] = 1.0;
    }
    let mut converged = false;
    let mut iterations = 0;

    for _iter in 0..config.max_iter {
        iterations += 1;

        // ── Project scores onto current discriminative subspace ───────────
        // proj_scores[i, d] = sum_j scores[(i,j)] * disc_dirs[j + d*ncomp_eff]
        // disc_dirs is ncomp_eff x p_disc_eff column-major
        let mut proj_scores = vec![0.0_f64; n * p_disc_eff]; // n x p_disc_eff row-major
        for i in 0..n {
            for d in 0..p_disc_eff {
                let mut val = 0.0;
                for j in 0..ncomp_eff {
                    val += scores[(i, j)] * disc_dirs[j + d * ncomp_eff];
                }
                proj_scores[i * p_disc_eff + d] = val;
            }
        }

        // ── GMM in projected subspace: per-cluster mean and diag variance ─
        // Compute cluster means in projected space
        let mut mu_disc: Vec<Vec<f64>> = vec![vec![0.0; p_disc_eff]; k];
        let mut n_k_soft: Vec<f64> = vec![0.0; k];
        for i in 0..n {
            for ki in 0..k {
                let r = resp[i * k + ki];
                n_k_soft[ki] += r;
                for d in 0..p_disc_eff {
                    mu_disc[ki][d] += r * proj_scores[i * p_disc_eff + d];
                }
            }
        }
        for ki in 0..k {
            if n_k_soft[ki] > 1e-10 {
                for d in 0..p_disc_eff {
                    mu_disc[ki][d] /= n_k_soft[ki];
                }
            }
        }

        // Per-cluster diagonal variance in projected space
        let mut var_disc: Vec<Vec<f64>> = vec![vec![1.0; p_disc_eff]; k];
        for ki in 0..k {
            if n_k_soft[ki] > 1e-10 {
                for d in 0..p_disc_eff {
                    let mut v = 0.0;
                    for i in 0..n {
                        let diff = proj_scores[i * p_disc_eff + d] - mu_disc[ki][d];
                        v += resp[i * k + ki] * diff * diff;
                    }
                    var_disc[ki][d] = (v / n_k_soft[ki]).max(1e-8);
                }
            }
        }

        // ── E-step: compute log-responsibilities in projected space ────────
        let mut log_resp = vec![0.0_f64; n * k];
        let mut ll = 0.0;
        for i in 0..n {
            let mut log_components = vec![0.0_f64; k];
            for ki in 0..k {
                let log_pi = if pi[ki] > 1e-300 { pi[ki].ln() } else { -700.0 };
                let mut log_lik = log_pi;
                for d in 0..p_disc_eff {
                    let diff = proj_scores[i * p_disc_eff + d] - mu_disc[ki][d];
                    let var = var_disc[ki][d];
                    log_lik -= 0.5 * (var.ln() + diff * diff / var);
                }
                log_lik -= 0.5 * (p_disc_eff as f64) * std::f64::consts::TAU.ln();
                log_components[ki] = log_lik;
            }
            // Log-sum-exp normalization
            let log_sum = log_sum_exp(&log_components);
            ll += log_sum;
            for ki in 0..k {
                log_resp[i * k + ki] = log_components[ki] - log_sum;
            }
        }

        // Exponentiate
        resp.fill(0.0);
        for i in 0..n {
            for ki in 0..k {
                resp[i * k + ki] = log_resp[i * k + ki].exp().max(1e-300);
            }
        }

        // ── Update pi from responsibilities ────────────────────────────────
        let mut n_k_new: Vec<f64> = vec![0.0; k];
        for i in 0..n {
            for ki in 0..k {
                n_k_new[ki] += resp[i * k + ki];
            }
        }
        let n_total: f64 = n_k_new.iter().sum();
        for ki in 0..k {
            pi[ki] = (n_k_new[ki] / n_total).max(1e-300);
        }

        // Hard assignments from responsibilities (for subspace update)
        cluster = (0..n)
            .map(|i| {
                (0..k)
                    .max_by(|&a, &b| {
                        resp[i * k + a]
                            .partial_cmp(&resp[i * k + b])
                            .unwrap_or(std::cmp::Ordering::Equal)
                    })
                    .unwrap_or(0)
            })
            .collect();

        // Update score-space params for next subspace computation
        update_gmm_params_from_soft(
            &row_major_scores,
            &resp,
            k,
            ncomp_eff,
            n,
            &mut pi,
            &mut mu_k,
            &mut sigma_k,
        );

        // ── Fisher-EM discriminative subspace update ───────────────────────
        // Compute between-scatter B_soft and within-scatter W_soft in score space
        let global_mean: Vec<f64> = (0..ncomp_eff)
            .map(|j| {
                (0..n)
                    .map(|i| row_major_scores[i * ncomp_eff + j])
                    .sum::<f64>()
                    / n as f64
            })
            .collect();

        let mut b_soft = vec![0.0_f64; ncomp_eff * ncomp_eff]; // ncomp_eff x ncomp_eff row-major
        let mut w_soft = vec![0.0_f64; ncomp_eff * ncomp_eff];

        // Between-scatter: sum_k n_k (mu_k - mu_global)(mu_k - mu_global)^T
        for ki in 0..k {
            let nk = n_k_new[ki].max(1.0);
            for j in 0..ncomp_eff {
                for l in 0..ncomp_eff {
                    b_soft[j * ncomp_eff + l] +=
                        nk * (mu_k[ki][j] - global_mean[j]) * (mu_k[ki][l] - global_mean[l]);
                }
            }
        }

        // Within-scatter: sum_i sum_k r_{ik} (x_i - mu_k)(x_i - mu_k)^T
        for i in 0..n {
            for ki in 0..k {
                let r = resp[i * k + ki];
                for j in 0..ncomp_eff {
                    let dj = row_major_scores[i * ncomp_eff + j] - mu_k[ki][j];
                    for l in 0..ncomp_eff {
                        let dl = row_major_scores[i * ncomp_eff + l] - mu_k[ki][l];
                        w_soft[j * ncomp_eff + l] += r * dj * dl;
                    }
                }
            }
        }

        // Add data-scaled regularization floor to W_soft diagonal (Pitfall 3)
        let trace_w: f64 = (0..ncomp_eff)
            .map(|j| w_soft[j * ncomp_eff + j])
            .sum::<f64>();
        let reg_floor = (trace_w / ncomp_eff as f64 * 1e-4).max(1e-8);
        for j in 0..ncomp_eff {
            w_soft[j * ncomp_eff + j] += reg_floor;
        }

        // Compute W^{-1} B via Cholesky + SVD: new discriminative directions
        // W_soft Cholesky, then solve each column of B
        match crate::linalg::cholesky_factor(&w_soft, ncomp_eff) {
            Ok(l_w) => {
                // Form W^{-1} B by solving W * X = B column by column
                let mut winv_b = vec![0.0_f64; ncomp_eff * ncomp_eff];
                for col in 0..ncomp_eff {
                    let b_col: Vec<f64> = (0..ncomp_eff)
                        .map(|r| b_soft[r * ncomp_eff + col])
                        .collect();
                    let x = crate::linalg::cholesky_forward_back(&l_w, &b_col, ncomp_eff);
                    for r in 0..ncomp_eff {
                        winv_b[r * ncomp_eff + col] = x[r];
                    }
                }

                // SVD of W^{-1}B to get top p_disc_eff directions
                use nalgebra::{DMatrix, SVD};
                let mat = DMatrix::from_row_slice(ncomp_eff, ncomp_eff, &winv_b);
                let svd = SVD::new(mat, true, false);
                if let Some(u) = svd.u {
                    // u columns = left singular vectors (ncomp_eff x ncomp_eff)
                    // Store top p_disc_eff columns as new disc_dirs (ncomp_eff x p_disc_eff column-major)
                    let mut new_dirs = vec![0.0_f64; ncomp_eff * p_disc_eff];
                    for d in 0..p_disc_eff {
                        for r in 0..ncomp_eff {
                            new_dirs[r + d * ncomp_eff] = u[(r, d)];
                        }
                    }
                    disc_dirs = new_dirs;
                }
                // If SVD u unavailable, keep previous disc_dirs
            }
            Err(_) => {
                // Cholesky failed — keep previous discriminative directions (identity fallback)
            }
        }

        // Convergence check
        let delta = (ll - prev_ll).abs();
        prev_ll = ll;
        if _iter > 0 && delta < config.tol {
            converged = true;
            break;
        }
    }

    // Build membership FdMatrix (n x k)
    let mut membership_data = vec![0.0_f64; n * k];
    for i in 0..n {
        for ki in 0..k {
            // column-major: membership[(i, ki)] = membership_data[i + ki*n]
            membership_data[i + ki * n] = resp[i * k + ki];
        }
    }
    let membership = FdMatrix::from_column_major(membership_data, n, k)?;

    // Build disc_subspace FdMatrix (ncomp_eff x p_disc_eff)
    // disc_dirs is already ncomp_eff x p_disc_eff column-major
    let disc_subspace = FdMatrix::from_column_major(disc_dirs, ncomp_eff, p_disc_eff)?;

    Ok(FunFemResult {
        cluster,
        membership,
        disc_subspace,
        log_likelihood: prev_ll,
        iterations,
        converged,
    })
}

/// Update GMM parameters from hard assignments.
fn update_gmm_params_from_hard(
    scores_rm: &[f64],
    cluster: &[usize],
    k: usize,
    d: usize,
    pi: &mut [f64],
    mu_k: &mut [Vec<f64>],
    sigma_k: &mut [Vec<f64>],
) {
    let n = cluster.len();
    let mut counts = vec![0usize; k];
    for &c in cluster {
        if c < k {
            counts[c] += 1;
        }
    }
    for ki in 0..k {
        pi[ki] = (counts[ki] as f64 / n as f64).max(1e-300);
        mu_k[ki] = vec![0.0; d];
        // Initialize to 0 so that the accumulation below yields Σ(x-μ)²
        // (not 1 + Σ(x-μ)², which was the previous incorrect value).
        sigma_k[ki] = vec![0.0; d];
        for i in 0..n {
            if cluster[i] == ki {
                for j in 0..d {
                    mu_k[ki][j] += scores_rm[i * d + j];
                }
            }
        }
        if counts[ki] > 0 {
            for j in 0..d {
                mu_k[ki][j] /= counts[ki] as f64;
            }
        }
        for i in 0..n {
            if cluster[i] == ki {
                for j in 0..d {
                    let diff = scores_rm[i * d + j] - mu_k[ki][j];
                    sigma_k[ki][j] += diff * diff;
                }
            }
        }
        if counts[ki] > 1 {
            for j in 0..d {
                sigma_k[ki][j] = (sigma_k[ki][j] / counts[ki] as f64).max(1e-8);
            }
        } else {
            for j in 0..d {
                sigma_k[ki][j] = 1.0;
            }
        }
    }
}

/// Update GMM parameters from soft responsibilities.
fn update_gmm_params_from_soft(
    scores_rm: &[f64],
    resp: &[f64],
    k: usize,
    d: usize,
    n: usize,
    pi: &mut [f64],
    mu_k: &mut [Vec<f64>],
    sigma_k: &mut [Vec<f64>],
) {
    let mut n_k = vec![0.0_f64; k];
    for i in 0..n {
        for ki in 0..k {
            n_k[ki] += resp[i * k + ki];
        }
    }
    let total: f64 = n_k.iter().sum();
    for ki in 0..k {
        pi[ki] = (n_k[ki] / total).max(1e-300);
        mu_k[ki] = vec![0.0; d];
        sigma_k[ki] = vec![1.0; d];
        if n_k[ki] > 1e-10 {
            for i in 0..n {
                let r = resp[i * k + ki];
                for j in 0..d {
                    mu_k[ki][j] += r * scores_rm[i * d + j];
                }
            }
            for j in 0..d {
                mu_k[ki][j] /= n_k[ki];
            }
            let mut var_j = vec![0.0_f64; d];
            for i in 0..n {
                let r = resp[i * k + ki];
                for j in 0..d {
                    let diff = scores_rm[i * d + j] - mu_k[ki][j];
                    var_j[j] += r * diff * diff;
                }
            }
            for j in 0..d {
                sigma_k[ki][j] = (var_j[j] / n_k[ki]).max(1e-8);
            }
        }
    }
}

/// Log-sum-exp of a slice (numerically stable).
fn log_sum_exp(v: &[f64]) -> f64 {
    let max_v = v.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    if max_v == f64::NEG_INFINITY {
        return f64::NEG_INFINITY;
    }
    max_v + v.iter().map(|&x| (x - max_v).exp()).sum::<f64>().ln()
}

/// Compute L2 distance between two rows in a flat row-major buffer.
///
/// `buf[i * m .. (i+1) * m]` is row `i`.
fn l2_dist_rowmajor(buf: &[f64], i: usize, j: usize, m: usize, weights: &[f64]) -> f64 {
    let mut sq = 0.0;
    for t in 0..m {
        let d = buf[i * m + t] - buf[j * m + t];
        sq += d * d * weights[t];
    }
    sq.sqrt()
}

// ────────────────────────────────────────────────────────────────────────────
// Align-and-cluster: elastic k-means via Karcher-mean templates
// ────────────────────────────────────────────────────────────────────────────

/// Configuration for joint align-and-cluster (elastic k-means).
///
/// Alternates between updating per-cluster template curves (via Karcher mean
/// in the elastic metric) and reassigning each curve to the nearest template
/// by elastic distance. This is the only clusterer in this module that is
/// shape-invariant — it finds clusters based on amplitude shape, not phase.
///
/// **Reference:** Inspired by Sangalli et al. (2010) joint clustering and alignment
/// and the `fdasrvf` elastic k-means variant (Srivastava et al.).
///
/// ## Example
///
/// ```
/// use fdars_core::clustering_advanced::{align_cluster_fd, AlignClusterConfig};
/// use fdars_core::matrix::FdMatrix;
/// use std::f64::consts::PI;
///
/// let m = 25;
/// let n = 10;
/// let t: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64).collect();
/// let mut col_major = vec![0.0_f64; n * m];
/// for i in 0..5 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).sin();
///     }
/// }
/// for i in 5..10 {
///     for (j, &tj) in t.iter().enumerate() {
///         col_major[i + j * n] = (2.0 * PI * tj).sin() + 5.0;
///     }
/// }
/// let data = FdMatrix::from_column_major(col_major, n, m).unwrap();
///
/// let mut cfg = AlignClusterConfig::default();
/// cfg.k = 2;
/// cfg.max_iter = 10;
/// cfg.karcher_max_iter = 8;
/// let result = align_cluster_fd(&data, &t, &cfg).unwrap();
/// assert_eq!(result.cluster.len(), n);
/// assert_eq!(result.templates.len(), 2);
/// ```
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub struct AlignClusterConfig {
    /// Number of clusters (default: 2). Must be ≥ 1 and ≤ n.
    pub k: usize,
    /// Maximum number of outer iterations (default: 20).
    pub max_iter: usize,
    /// Random seed for initial template selection (default: 42).
    pub seed: u64,
    /// If `true`, use `amplitude_distance` (shape-invariant); otherwise use
    /// `elastic_distance` (full Fisher-Rao distance including phase). Default: `true`.
    pub use_amplitude_only: bool,
    /// Penalty weight for `elastic_distance` when `use_amplitude_only = false` (default: 0.0).
    pub elastic_lambda: f64,
    /// Maximum inner iterations for Karcher mean (default: 15).
    pub karcher_max_iter: usize,
    /// Convergence tolerance for Karcher mean (default: 1e-4).
    pub karcher_tol: f64,
}

impl Default for AlignClusterConfig {
    fn default() -> Self {
        Self {
            k: 2,
            max_iter: 20,
            seed: 42,
            use_amplitude_only: true,
            elastic_lambda: 0.0,
            karcher_max_iter: 15,
            karcher_tol: 1e-4,
        }
    }
}

/// Result of joint align-and-cluster (elastic k-means).
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct AlignClusterResult {
    /// Cluster assignment for each curve (0-based, contiguous, length n).
    pub cluster: Vec<usize>,
    /// Per-cluster template curves (k entries, each of length m).
    pub templates: Vec<Vec<f64>>,
    /// Distance matrix (n × k) — `distances[(i, ki)]` is the elastic distance
    /// from curve i to template ki.
    pub distances: FdMatrix,
    /// Number of outer iterations performed.
    pub iterations: usize,
    /// Whether the algorithm converged (no label changes in last iteration).
    pub converged: bool,
}

/// Cluster functional data using elastic k-means with Karcher-mean templates.
///
/// Initialises k templates by random distinct-curve selection, then alternates
/// between reassigning each curve to the nearest template (by elastic distance)
/// and updating each template to the Karcher mean of its cluster members.
/// Empty clusters are reinitialized to a random non-member curve.
///
/// # Arguments
///
/// * `data` — Functional data matrix (n × m, column-major).
/// * `argvals` — Evaluation grid (length m).
/// * `config` — Algorithm parameters; see [`AlignClusterConfig`].
///
/// # Errors
///
/// Returns [`FdarError::InvalidDimension`] if `n == 0`, `m == 0`, or
/// `argvals.len() != m`. Returns [`FdarError::InvalidParameter`] if
/// `config.k == 0` or `config.k > n`.
#[must_use = "expensive computation whose result should not be discarded"]
pub fn align_cluster_fd(
    data: &FdMatrix,
    argvals: &[f64],
    config: &AlignClusterConfig,
) -> Result<AlignClusterResult, FdarError> {
    use crate::alignment::{amplitude_distance, elastic_distance, karcher_mean};

    let (n, m) = data.shape();

    // Validation
    if n == 0 || m == 0 {
        return Err(FdarError::InvalidDimension {
            parameter: "data",
            expected: "at least 1 row and 1 column".to_string(),
            actual: format!("{n} rows, {m} columns"),
        });
    }
    if argvals.len() != m {
        return Err(FdarError::InvalidDimension {
            parameter: "argvals",
            expected: format!("{m}"),
            actual: format!("{}", argvals.len()),
        });
    }
    if config.k == 0 {
        return Err(FdarError::InvalidParameter {
            parameter: "k",
            message: "k must be >= 1".to_string(),
        });
    }
    if config.k > n {
        return Err(FdarError::InvalidParameter {
            parameter: "k",
            message: format!("k={} exceeds number of curves n={}", config.k, n),
        });
    }

    let k = config.k;

    // Initialize k templates by seeded random distinct-curve selection.
    // Use a shuffled index list and pick every n/k curves for spread,
    // then shuffle within the stratum using the seeded RNG to ensure
    // reproducibility while covering the data range.
    let mut rng = StdRng::seed_from_u64(config.seed);
    let mut shuffled: Vec<usize> = (0..n).collect();
    // Fisher-Yates shuffle
    for i in (1..n).rev() {
        let j = rng.gen_range(0..=i);
        shuffled.swap(i, j);
    }
    // Pick k evenly-strided indices from the shuffled list for spread
    let step = n / k;
    let template_indices: Vec<usize> = (0..k).map(|ki| shuffled[(ki * step).min(n - 1)]).collect();

    // Initialize templates as the selected curve values
    let mut templates: Vec<Vec<f64>> = template_indices.iter().map(|&ci| data.row(ci)).collect();

    let mut cluster: Vec<usize> = vec![0; n];
    let mut distances = FdMatrix::zeros(n, k);
    let mut converged = false;
    let mut iterations = 0;

    for _iter in 0..config.max_iter {
        iterations += 1;

        // ── Reassignment: compute distances to each template ──────────────
        for i in 0..n {
            let curve_i = data.row(i);
            for ki in 0..k {
                let dist = if config.use_amplitude_only {
                    amplitude_distance(&curve_i, &templates[ki], argvals, config.elastic_lambda)
                } else {
                    elastic_distance(&curve_i, &templates[ki], argvals, config.elastic_lambda)
                };
                distances[(i, ki)] = dist;
            }
        }

        // Assign each curve to nearest template
        let mut changed = false;
        for i in 0..n {
            let best_k = (0..k)
                .min_by(|&a, &b| {
                    distances[(i, a)]
                        .partial_cmp(&distances[(i, b)])
                        .unwrap_or(std::cmp::Ordering::Equal)
                })
                .unwrap_or(0);
            if cluster[i] != best_k {
                cluster[i] = best_k;
                changed = true;
            }
        }

        // ── Template update: Karcher mean of cluster members ──────────────
        // Track whether any template was changed (including by empty-cluster
        // reinit) so convergence is not declared on the same iteration that a
        // reinit occurred — the new template must be tested in the next pass.
        let mut template_changed = false;
        for ki in 0..k {
            let member_indices: Vec<usize> = (0..n).filter(|&i| cluster[i] == ki).collect();

            if member_indices.is_empty() {
                // Empty-cluster fallback (Pitfall 6): reinit to a random non-member curve
                // Find all non-members
                let non_members: Vec<usize> = (0..n).filter(|&i| cluster[i] != ki).collect();
                if !non_members.is_empty() {
                    let rand_idx = rng.gen_range(0..non_members.len());
                    templates[ki] = data.row(non_members[rand_idx]);
                    template_changed = true;
                }
                // (If somehow all n are in one cluster, keep the old template)
                continue;
            }

            // Gather member rows into (n_k x m) FdMatrix
            let n_k = member_indices.len();
            let mut col_major_k = vec![0.0_f64; n_k * m];
            for (row_in_k, &orig_i) in member_indices.iter().enumerate() {
                for j in 0..m {
                    col_major_k[row_in_k + j * n_k] = data[(orig_i, j)];
                }
            }
            let data_k = FdMatrix::from_column_major(col_major_k, n_k, m)?;

            // Karcher mean — returns KarcherMeanResult, template is `.mean` field
            let km = karcher_mean(
                &data_k,
                argvals,
                config.karcher_max_iter,
                config.karcher_tol,
                config.elastic_lambda,
            );
            templates[ki] = km.mean;
        }

        // Converge only when no curve moved AND no template was reinit-ed due to
        // an empty cluster.  Including template_changed prevents the algorithm
        // from declaring convergence on the same iteration as an empty-cluster
        // reinit, before the new template has been used in a subsequent
        // reassignment pass.
        if !changed && !template_changed {
            converged = true;
            break;
        }
    }

    Ok(AlignClusterResult {
        cluster,
        templates,
        distances,
        iterations,
        converged,
    })
}

// ────────────────────────────────────────────────────────────────────────────
// Tests
// ────────────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::{adjusted_rand_index, uniform_grid};
    use std::f64::consts::PI;

    // ── Synthetic data generators ────────────────────────────────────────

    /// Two tight clusters, n_per curves each. Cluster 0: sin wave.
    /// Cluster 1: sin wave shifted up by 5.
    fn two_tight_clusters(n_per: usize, m: usize) -> (FdMatrix, Vec<f64>, Vec<usize>) {
        let t = uniform_grid(m);
        let n = 2 * n_per;
        let mut col_major = vec![0.0_f64; n * m];
        for i in 0..n_per {
            for (j, &tj) in t.iter().enumerate() {
                col_major[i + j * n] = (2.0 * PI * tj).sin();
            }
        }
        for i in 0..n_per {
            for (j, &tj) in t.iter().enumerate() {
                col_major[(i + n_per) + j * n] = (2.0 * PI * tj).sin() + 5.0;
            }
        }
        let labels: Vec<usize> = (0..n).map(|i| if i < n_per { 0 } else { 1 }).collect();
        (
            FdMatrix::from_column_major(col_major, n, m).unwrap(),
            t,
            labels,
        )
    }

    /// 2 tight clusters (n_per each) + 2 far constant-offset outlier curves.
    fn clusters_with_noise(n_per: usize, m: usize) -> (FdMatrix, Vec<f64>) {
        let t = uniform_grid(m);
        let n = 2 * n_per + 2;
        let mut col_major = vec![0.0_f64; n * m];
        // Cluster 0: sin wave
        for i in 0..n_per {
            for (j, &tj) in t.iter().enumerate() {
                col_major[i + j * n] = (2.0 * PI * tj).sin();
            }
        }
        // Cluster 1: sin wave + 5
        for i in 0..n_per {
            for (j, &tj) in t.iter().enumerate() {
                col_major[(i + n_per) + j * n] = (2.0 * PI * tj).sin() + 5.0;
            }
        }
        // Outlier 0: constant 100
        let o0 = 2 * n_per;
        for j in 0..m {
            col_major[o0 + j * n] = 100.0;
        }
        // Outlier 1: constant -100
        let o1 = 2 * n_per + 1;
        for j in 0..m {
            col_major[o1 + j * n] = -100.0;
        }
        (FdMatrix::from_column_major(col_major, n, m).unwrap(), t)
    }

    /// Two well-separated clusters for kCFC testing.
    fn two_separated_clusters(n_per: usize, m: usize) -> (FdMatrix, Vec<f64>, Vec<usize>) {
        let t = uniform_grid(m);
        let n = 2 * n_per;
        let mut col_major = vec![0.0_f64; n * m];
        // Cluster 0: sin waves
        for i in 0..n_per {
            for (j, &tj) in t.iter().enumerate() {
                col_major[i + j * n] = (2.0 * PI * tj).sin() + 0.05 * (i as f64 / n_per as f64);
            }
        }
        // Cluster 1: cos waves shifted up by 8 (very different shape)
        for i in 0..n_per {
            for (j, &tj) in t.iter().enumerate() {
                col_major[(i + n_per) + j * n] =
                    (2.0 * PI * tj).cos() + 8.0 + 0.05 * (i as f64 / n_per as f64);
            }
        }
        let labels: Vec<usize> = (0..n).map(|i| if i < n_per { 0 } else { 1 }).collect();
        (
            FdMatrix::from_column_major(col_major, n, m).unwrap(),
            t,
            labels,
        )
    }

    // ── DBSCAN tests ─────────────────────────────────────────────────────

    #[test]
    fn test_dbscan_core_points() {
        let m = 30;
        let n_per = 5;
        let (data, t, _labels) = two_tight_clusters(n_per, m);
        let result = dbscan_fd(
            &data,
            &t,
            &DbscanConfig {
                eps: 1.0,
                min_points: 2,
                ..Default::default()
            },
        )
        .unwrap();
        assert_eq!(result.n_clusters, 2, "expected 2 clusters");
        assert_eq!(result.n_noise, 0, "expected 0 noise points");
        assert_eq!(result.cluster.len(), 2 * n_per);
    }

    #[test]
    fn test_dbscan_noise_flagging() {
        let m = 30;
        let n_per = 5;
        let (data, t) = clusters_with_noise(n_per, m);
        let result = dbscan_fd(
            &data,
            &t,
            &DbscanConfig {
                eps: 1.5,
                min_points: 2,
                ..Default::default()
            },
        )
        .unwrap();
        // The 2 outlier curves should be noise
        assert_eq!(
            result.n_noise, 2,
            "expected exactly 2 noise points, got {}",
            result.n_noise
        );
        assert_eq!(result.n_clusters, 2, "expected 2 clusters");
        // Verify outlier indices (last 2) are None
        let n = data.nrows();
        assert!(result.cluster[n - 2].is_none(), "outlier 0 should be noise");
        assert!(result.cluster[n - 1].is_none(), "outlier 1 should be noise");
    }

    #[test]
    fn test_dbscan_zero_eps_returns_err() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        assert!(
            dbscan_fd(
                &data,
                &t,
                &DbscanConfig {
                    eps: 0.0,
                    ..Default::default()
                }
            )
            .is_err(),
            "eps=0 should return Err"
        );
    }

    #[test]
    fn test_dbscan_negative_eps_returns_err() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        assert!(
            dbscan_fd(
                &data,
                &t,
                &DbscanConfig {
                    eps: -1.0,
                    ..Default::default()
                }
            )
            .is_err(),
            "eps=-1 should return Err"
        );
    }

    #[test]
    fn test_dbscan_invalid_min_points_zero() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        assert!(
            dbscan_fd(
                &data,
                &t,
                &DbscanConfig {
                    min_points: 0,
                    ..Default::default()
                }
            )
            .is_err(),
            "min_points=0 should return Err"
        );
    }

    #[test]
    fn test_dbscan_empty_data() {
        let data = FdMatrix::zeros(0, 0);
        let t: Vec<f64> = vec![];
        assert!(
            dbscan_fd(&data, &t, &DbscanConfig::default()).is_err(),
            "empty data should return Err"
        );
    }

    #[test]
    fn test_dbscan_mismatched_argvals() {
        let m = 20;
        let (data, _t, _) = two_tight_clusters(5, m);
        let wrong_t = uniform_grid(m + 1);
        assert!(
            dbscan_fd(&data, &wrong_t, &DbscanConfig::default()).is_err(),
            "mismatched argvals should return Err"
        );
    }

    #[test]
    fn test_dbscan_distances_shape() {
        let m = 20;
        let n_per = 4;
        let (data, t, _) = two_tight_clusters(n_per, m);
        let result = dbscan_fd(
            &data,
            &t,
            &DbscanConfig {
                eps: 1.0,
                min_points: 2,
                ..Default::default()
            },
        )
        .unwrap();
        let n = 2 * n_per;
        assert_eq!(
            result.distances.shape(),
            (n, n),
            "distance matrix must be n x n"
        );
    }

    // ── kCFC tests ───────────────────────────────────────────────────────

    #[test]
    fn test_kcfc_recovery() {
        let m = 40;
        let n_per = 10;
        let (data, t, ground_truth) = two_separated_clusters(n_per, m);
        let result = kcfc_cluster(
            &data,
            &t,
            &KcfcConfig {
                k: 2,
                ncomp: 3,
                max_iter: 50,
                seed: 42,
                ..Default::default()
            },
        )
        .unwrap();
        let ari = adjusted_rand_index(&result.cluster, &ground_truth);
        assert!(
            ari >= 0.90,
            "kCFC ARI={ari:.3} should be >= 0.90 on well-separated data"
        );
    }

    #[test]
    fn test_kcfc_errors_ordering() {
        // Curves in cluster 0 should have smaller error against cluster 0's FPCA
        // than against cluster 1's FPCA, and vice versa.
        let m = 40;
        let n_per = 10;
        let (data, t, ground_truth) = two_separated_clusters(n_per, m);
        let result = kcfc_cluster(
            &data,
            &t,
            &KcfcConfig {
                k: 2,
                ncomp: 3,
                max_iter: 50,
                seed: 42,
                ..Default::default()
            },
        )
        .unwrap();

        // Determine which result cluster corresponds to ground truth 0 and 1.
        // Use an explicit lookup instead of arithmetic (1 - gt0_cluster) so the
        // test is not k==2-specific and does not risk wrapping on usize.
        let gt0_cluster = result.cluster[0]; // first curve is in ground truth 0
        let gt1_cluster = result.cluster[n_per]; // first curve in ground truth 1

        let mut correct_ordering = 0;
        let mut total = 0;
        for i in 0..data.nrows() {
            let expected_cluster = if ground_truth[i] == 0 {
                gt0_cluster
            } else {
                gt1_cluster
            };
            // Find the "other" cluster by scanning — avoids k==2-only arithmetic.
            let other_cluster = (0..2).find(|&c| c != expected_cluster).unwrap_or(0);
            let err_own = result.reconstruction_errors[(i, expected_cluster)];
            let err_other = result.reconstruction_errors[(i, other_cluster)];
            if err_own.is_finite() && err_other.is_finite() {
                if err_own < err_other {
                    correct_ordering += 1;
                }
                total += 1;
            }
        }
        // At least 80% should have correct error ordering
        assert!(
            correct_ordering * 10 >= total * 8,
            "only {correct_ordering}/{total} curves had smaller error for their true cluster"
        );
    }

    #[test]
    fn test_kcfc_deterministic() {
        let m = 30;
        let n_per = 8;
        let (data, t, _) = two_separated_clusters(n_per, m);
        let cfg = KcfcConfig {
            k: 2,
            ncomp: 2,
            max_iter: 30,
            seed: 7,
            ..Default::default()
        };
        let r1 = kcfc_cluster(&data, &t, &cfg).unwrap();
        let r2 = kcfc_cluster(&data, &t, &cfg).unwrap();
        assert_eq!(
            r1.cluster, r2.cluster,
            "identical seed must produce identical assignments"
        );
    }

    #[test]
    fn test_kcfc_invalid_k_zero() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        assert!(
            kcfc_cluster(
                &data,
                &t,
                &KcfcConfig {
                    k: 0,
                    ..Default::default()
                }
            )
            .is_err(),
            "k=0 should return Err"
        );
    }

    #[test]
    fn test_kcfc_invalid_k_gt_n() {
        let m = 20;
        let n = 4;
        let (data, t, _) = two_tight_clusters(n / 2, m);
        assert!(
            kcfc_cluster(
                &data,
                &t,
                &KcfcConfig {
                    k: n + 1,
                    ..Default::default()
                }
            )
            .is_err(),
            "k>n should return Err"
        );
    }

    #[test]
    fn test_kcfc_empty_data() {
        let data = FdMatrix::zeros(0, 0);
        let t: Vec<f64> = vec![];
        assert!(
            kcfc_cluster(&data, &t, &KcfcConfig::default()).is_err(),
            "empty data should return Err"
        );
    }

    #[test]
    fn test_kcfc_mismatched_argvals() {
        let m = 20;
        let (data, _t, _) = two_tight_clusters(5, m);
        let wrong_t = uniform_grid(m + 3);
        assert!(
            kcfc_cluster(&data, &wrong_t, &KcfcConfig::default()).is_err(),
            "mismatched argvals should return Err"
        );
    }

    #[test]
    fn test_kcfc_result_shapes() {
        let m = 20;
        let n_per = 5;
        let (data, t, _) = two_separated_clusters(n_per, m);
        let n = 2 * n_per;
        let result = kcfc_cluster(
            &data,
            &t,
            &KcfcConfig {
                k: 2,
                ncomp: 2,
                ..Default::default()
            },
        )
        .unwrap();
        assert_eq!(result.cluster.len(), n);
        assert_eq!(result.fpca_models.len(), 2);
        assert_eq!(result.reconstruction_errors.shape(), (n, 2));
    }

    #[test]
    fn test_kcfc_ncomp_zero_returns_err() {
        // ncomp == 0 must be caught at entry rather than silently assigning all
        // curves to cluster 0 (which was the pre-fix behavior when fdata_to_pc_1d
        // returned Err and was swallowed by the degenerate-cluster arm).
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        let result = kcfc_cluster(
            &data,
            &t,
            &KcfcConfig {
                k: 2,
                ncomp: 0,
                ..Default::default()
            },
        );
        assert!(result.is_err(), "ncomp=0 must return Err, got Ok");
        if let Err(FdarError::InvalidParameter { parameter, .. }) = result {
            assert_eq!(parameter, "ncomp");
        } else {
            panic!("expected InvalidParameter {{ parameter: \"ncomp\" }}");
        }
    }

    // ── funFEM tests ─────────────────────────────────────────────────────────

    #[test]
    fn test_funfem_recovery() {
        let m = 40;
        let n_per = 10;
        let (data, t, ground_truth) = two_separated_clusters(n_per, m);
        let result = funfem_cluster(
            &data,
            &t,
            &FunFemConfig {
                k: 2,
                ncomp: 5,
                p_disc: 1,
                max_iter: 30,
                tol: 1e-5,
                seed: 42,
            },
        )
        .unwrap();
        let ari = adjusted_rand_index(&result.cluster, &ground_truth);
        assert!(
            ari >= 0.90,
            "funFEM ARI={ari:.3} should be >= 0.90 on well-separated data"
        );
    }

    #[test]
    fn test_funfem_deterministic() {
        let m = 30;
        let n_per = 8;
        let (data, t, _) = two_separated_clusters(n_per, m);
        let cfg = FunFemConfig {
            k: 2,
            ncomp: 4,
            p_disc: 1,
            max_iter: 20,
            tol: 1e-5,
            seed: 99,
        };
        let r1 = funfem_cluster(&data, &t, &cfg).unwrap();
        let r2 = funfem_cluster(&data, &t, &cfg).unwrap();
        assert_eq!(
            r1.cluster, r2.cluster,
            "same seed must produce identical assignments"
        );
    }

    #[test]
    fn test_funfem_invalid_k_zero() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        assert!(
            funfem_cluster(
                &data,
                &t,
                &FunFemConfig {
                    k: 0,
                    ..FunFemConfig::default()
                }
            )
            .is_err(),
            "k=0 must return Err"
        );
    }

    #[test]
    fn test_funfem_invalid_k_gt_n() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(3, m);
        assert!(
            funfem_cluster(
                &data,
                &t,
                &FunFemConfig {
                    k: 10,
                    ..FunFemConfig::default()
                }
            )
            .is_err(),
            "k>n must return Err"
        );
    }

    #[test]
    fn test_funfem_invalid_ncomp_zero() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        assert!(
            funfem_cluster(
                &data,
                &t,
                &FunFemConfig {
                    ncomp: 0,
                    ..FunFemConfig::default()
                }
            )
            .is_err(),
            "ncomp=0 must return Err"
        );
    }

    #[test]
    fn test_funfem_invalid_empty_data() {
        let data = FdMatrix::zeros(0, 0);
        let t: Vec<f64> = vec![];
        assert!(
            funfem_cluster(&data, &t, &FunFemConfig::default()).is_err(),
            "empty data must return Err"
        );
    }

    #[test]
    fn test_funfem_invalid_argvals_mismatch() {
        let m = 20;
        let (data, _t, _) = two_tight_clusters(5, m);
        let wrong_t = uniform_grid(m + 2);
        assert!(
            funfem_cluster(&data, &wrong_t, &FunFemConfig::default()).is_err(),
            "argvals mismatch must return Err"
        );
    }

    #[test]
    fn test_funfem_output_shapes() {
        let m = 30;
        let n_per = 6;
        let (data, t, _) = two_separated_clusters(n_per, m);
        let n = 2 * n_per;
        let result = funfem_cluster(
            &data,
            &t,
            &FunFemConfig {
                k: 2,
                ncomp: 4,
                p_disc: 1,
                max_iter: 10,
                tol: 1e-4,
                seed: 1,
            },
        )
        .unwrap();
        assert_eq!(result.cluster.len(), n);
        assert_eq!(result.membership.shape(), (n, 2));
    }

    // ── Align-and-cluster tests ───────────────────────────────────────────

    /// Two clusters differing in fundamental shape:
    /// Group 0 = flat curves (constant value near 0)
    /// Group 1 = sin curves (oscillating shape, amplitude offset +5)
    ///
    /// Each group has small within-group phase/warp variability via different
    /// time-warps (g(t) = t^alpha with alpha varying slightly across members).
    /// The shapes are so distinct that elastic k-means trivially recovers them,
    /// demonstrating shape-invariant clustering on data with within-group warps.
    fn time_warped_clusters(n_per: usize, m: usize) -> (FdMatrix, Vec<f64>, Vec<usize>) {
        let t = uniform_grid(m);
        let n = 2 * n_per;
        let mut col_major = vec![0.0_f64; n * m];
        // Cluster 0: sin(2πt) evaluated at time-warped grid t^alpha (alpha near 1)
        // Small phase warp within cluster — but all are sin-shaped
        for i in 0..n_per {
            let alpha = 1.0 + 0.1 * (i as f64 / n_per as f64); // alpha in [1.0, 1.1]
            for (j, &tj) in t.iter().enumerate() {
                let warped = tj.powf(alpha);
                col_major[i + j * n] = (2.0 * PI * warped).sin();
            }
        }
        // Cluster 1: flat curves (constant value 8) — completely different shape
        for i in 0..n_per {
            for j in 0..m {
                col_major[(i + n_per) + j * n] = 8.0 + 0.05 * i as f64;
            }
        }
        let labels: Vec<usize> = (0..n).map(|i| if i < n_per { 0 } else { 1 }).collect();
        (
            FdMatrix::from_column_major(col_major, n, m).unwrap(),
            t,
            labels,
        )
    }

    #[test]
    fn test_align_cluster_shape_shift() {
        // Use small n_per and m as per the anti-stall note (elastic DP is slow)
        let m = 30;
        let n_per = 6;
        let (data, t, ground_truth) = time_warped_clusters(n_per, m);
        let result = align_cluster_fd(
            &data,
            &t,
            &AlignClusterConfig {
                k: 2,
                max_iter: 15,
                seed: 42,
                use_amplitude_only: true,
                elastic_lambda: 0.0,
                karcher_max_iter: 10,
                karcher_tol: 1e-3,
            },
        )
        .unwrap();
        let ari = adjusted_rand_index(&result.cluster, &ground_truth);
        assert!(
            ari >= 0.90,
            "align_cluster ARI={ari:.3} on shape-distinct data should be >= 0.90"
        );
    }

    #[test]
    fn test_align_cluster_recovery() {
        let m = 30;
        let n_per = 6;
        let (data, t, ground_truth) = two_separated_clusters(n_per, m);
        let result = align_cluster_fd(
            &data,
            &t,
            &AlignClusterConfig {
                k: 2,
                max_iter: 15,
                seed: 7,
                use_amplitude_only: true,
                elastic_lambda: 0.0,
                karcher_max_iter: 10,
                karcher_tol: 1e-3,
            },
        )
        .unwrap();
        let ari = adjusted_rand_index(&result.cluster, &ground_truth);
        assert!(
            ari >= 0.90,
            "align_cluster ARI={ari:.3} on amplitude-separated data should be >= 0.90"
        );
    }

    #[test]
    fn test_align_cluster_invalid_k_zero() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(5, m);
        assert!(
            align_cluster_fd(
                &data,
                &t,
                &AlignClusterConfig {
                    k: 0,
                    ..AlignClusterConfig::default()
                }
            )
            .is_err(),
            "k=0 must return Err"
        );
    }

    #[test]
    fn test_align_cluster_invalid_k_gt_n() {
        let m = 20;
        let (data, t, _) = two_tight_clusters(3, m);
        assert!(
            align_cluster_fd(
                &data,
                &t,
                &AlignClusterConfig {
                    k: 20,
                    ..AlignClusterConfig::default()
                }
            )
            .is_err(),
            "k>n must return Err"
        );
    }

    #[test]
    fn test_align_cluster_invalid_empty_data() {
        let data = FdMatrix::zeros(0, 0);
        let t: Vec<f64> = vec![];
        assert!(
            align_cluster_fd(&data, &t, &AlignClusterConfig::default()).is_err(),
            "empty data must return Err"
        );
    }

    #[test]
    fn test_align_cluster_invalid_argvals_mismatch() {
        let m = 20;
        let (data, _t, _) = two_tight_clusters(5, m);
        let wrong_t = uniform_grid(m + 5);
        assert!(
            align_cluster_fd(&data, &wrong_t, &AlignClusterConfig::default()).is_err(),
            "argvals mismatch must return Err"
        );
    }

    #[test]
    fn test_align_cluster_output_shapes() {
        let m = 20;
        let n_per = 4;
        let (data, t, _) = two_separated_clusters(n_per, m);
        let n = 2 * n_per;
        let result = align_cluster_fd(
            &data,
            &t,
            &AlignClusterConfig {
                k: 2,
                max_iter: 5,
                seed: 1,
                karcher_max_iter: 5,
                karcher_tol: 1e-2,
                ..AlignClusterConfig::default()
            },
        )
        .unwrap();
        assert_eq!(result.cluster.len(), n);
        assert_eq!(result.templates.len(), 2);
        assert!(result.templates.iter().all(|t| t.len() == m));
        assert_eq!(result.distances.shape(), (n, 2));
    }
}