rosa 0.1.0

use crate::matrix::Matrix;

/// Distance metric for k-NN computation.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum DistanceMetric {
    Euclidean,
    Cosine,
}

/// Output mode for the recurrence matrix.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum RecurrenceMode {
    /// Binary 0/1 connectivity.
    Connectivity,
    /// `exp(-d / bandwidth)` affinity values.
    Affinity,
    /// Raw distances.
    Distance,
}

/// Parameters for `recurrence_matrix`.
pub struct RecurrenceParams {
    /// Number of nearest neighbors. `None` → `2 * ceil(sqrt(T - 2*width + 1))`, min 2.
    pub k: Option<usize>,
    /// Diagonal exclusion band: exclude `|i - j| < width`. Default 1.
    pub width: usize,
    /// Distance metric. Default `Euclidean`.
    pub metric: DistanceMetric,
    /// Output mode. Default `Connectivity`.
    pub mode: RecurrenceMode,
    /// Affinity bandwidth. `None` → median k-th neighbor distance.
    pub bandwidth: Option<f64>,
    /// If true, symmetrize via mutual nearest neighbors: `R = min(R, R^T)`.
    pub sym: bool,
    /// If true, populate the diagonal (self-links). Default false.
    pub self_links: bool,
}

impl Default for RecurrenceParams {
    fn default() -> Self {
        Self {
            k: None,
            width: 1,
            metric: DistanceMetric::Euclidean,
            mode: RecurrenceMode::Connectivity,
            bandwidth: None,
            sym: false,
            self_links: false,
        }
    }
}

/// Compute the recurrence (self-similarity) matrix.
///
/// Matches `librosa.segment.recurrence_matrix`:
/// 1. Find `k + 2*width` nearest neighbors (excluding self)
/// 2. Zero out diagonal band `|i - j| < width`
/// 3. Keep only top-k nearest neighbors per row
/// 4. For affinity mode, convert to `exp(-d / bandwidth)`
/// 5. Symmetrize if requested
/// 6. Transpose (librosa convention)
///
/// Input:  `data [n_features, T]` — each column is a feature vector
/// Output: `[T, T]` — recurrence/similarity/distance matrix
pub fn recurrence_matrix(data: &Matrix, params: &RecurrenceParams) -> Matrix {
    let t = data.cols();
    let width = params.width;

    // Default k: matches librosa's formula
    let k = params.k.unwrap_or_else(|| {
        if t > 2 * width + 1 {
            let k = 2 * (((t - 2 * width + 1) as f64).sqrt().ceil() as usize);
            k.max(2)
        } else {
            2
        }
    });
    let k = k.min(t.saturating_sub(1));

    // Number of neighbors to query (extra to account for band exclusion)
    let k_query = (k + 2 * width).min(t.saturating_sub(1));

    // Compute full pairwise distance matrix [T, T]
    let dist = compute_distance_matrix(data, params.metric);

    // Step 1: For each point, find k_query nearest neighbors (excluding self).
    // Store distances in rec (we need distances for band removal + top-k + affinity).
    let mut rec = vec![0.0f64; t * t];

    for i in 0..t {
        let mut candidates: Vec<(f64, usize)> = (0..t)
            .filter(|&j| j != i)
            .map(|j| (dist[i * t + j], j))
            .collect();

        candidates
            .sort_unstable_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

        let take = k_query.min(candidates.len());
        for &(d, j) in candidates.iter().take(take) {
            rec[i * t + j] = if params.mode == RecurrenceMode::Connectivity {
                1.0
            } else {
                d
            };
        }
    }

    // Step 2: Zero out diagonal band |i - j| < width
    for i in 0..t {
        for offset in 0..width {
            if i + offset < t {
                rec[i * t + (i + offset)] = 0.0;
            }
            if offset > 0 && i >= offset {
                rec[i * t + (i - offset)] = 0.0;
            }
        }
    }

    // Step 3: Retain only top-k links per row
    for i in 0..t {
        let mut links: Vec<usize> = (0..t).filter(|&j| rec[i * t + j] != 0.0).collect();

        if links.len() > k {
            if params.mode == RecurrenceMode::Connectivity {
                // For connectivity: argsort on constant values preserves column index order.
                // Links are already in ascending column index order (from the iterator).
                // Keep first k, remove the rest.
            } else {
                // Sort ascending by stored distance value
                links.sort_unstable_by(|&a, &b| {
                    rec[i * t + a]
                        .partial_cmp(&rec[i * t + b])
                        .unwrap_or(std::cmp::Ordering::Equal)
                });
            }
            for &j in links[k..].iter() {
                rec[i * t + j] = 0.0;
            }
        }
    }

    // Step 4: Handle self-links before symmetrization
    if !params.self_links {
        for i in 0..t {
            rec[i * t + i] = 0.0;
        }
    }

    // Step 5: Symmetrize if requested (before affinity conversion)
    if params.sym {
        let mut sym_rec = vec![0.0f64; t * t];
        for i in 0..t {
            for j in i..t {
                let rij = rec[i * t + j];
                let rji = rec[j * t + i];
                let m = rij.min(rji);
                sym_rec[i * t + j] = m;
                sym_rec[j * t + i] = m;
            }
        }
        rec = sym_rec;
    }

    // Step 6: For affinity mode, compute bandwidth and convert
    if matches!(params.mode, RecurrenceMode::Affinity) {
        // 'med_k_scalar': for each row, sort nonzero distances, take k-th,
        // then take the median across all rows.
        let bandwidth = params.bandwidth.unwrap_or_else(|| {
            let mut kth_dists: Vec<f64> = Vec::with_capacity(t);
            for i in 0..t {
                let mut row_dists: Vec<f64> = (0..t)
                    .filter(|&j| rec[i * t + j] > 0.0)
                    .map(|j| rec[i * t + j])
                    .collect();
                row_dists
                    .sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                let take = k.min(row_dists.len());
                if take > 0 {
                    kth_dists.push(row_dists[take - 1]);
                }
            }
            kth_dists
                .sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
            let n = kth_dists.len();
            if n == 0 {
                1.0
            } else if n.is_multiple_of(2) {
                (kth_dists[n / 2 - 1] + kth_dists[n / 2]) / 2.0
            } else {
                kth_dists[n / 2]
            }
        });
        let bandwidth = if bandwidth < 1e-15 { 1.0 } else { bandwidth };

        for v in rec.iter_mut() {
            if *v != 0.0 {
                *v = (-(*v) / bandwidth).exp();
            }
        }
    }

    // Step 7: Transpose (librosa convention: rec[i,j] means j is neighbor of i,
    // output is transposed so result[i,j] = 1 iff i is a neighbor of j)
    let mut transposed = vec![0.0f64; t * t];
    for i in 0..t {
        for j in 0..t {
            transposed[j * t + i] = rec[i * t + j];
        }
    }

    // Step 8: Handle connectivity type
    if params.mode == RecurrenceMode::Connectivity {
        for v in transposed.iter_mut() {
            *v = if *v != 0.0 { 1.0 } else { 0.0 };
        }
    }

    Matrix::from_vec(transposed, t, t)
}

fn compute_distance_matrix(data: &Matrix, metric: DistanceMetric) -> Vec<f64> {
    let t = data.cols();
    let d = data.rows();

    // Precompute norms for cosine
    let norms: Vec<f64> = if matches!(metric, DistanceMetric::Cosine) {
        (0..t)
            .map(|j| (0..d).map(|r| data.get(r, j).powi(2)).sum::<f64>().sqrt())
            .collect()
    } else {
        vec![]
    };

    let mut dist = vec![0.0f64; t * t];
    for i in 0..t {
        for j in (i + 1)..t {
            let d_val = match metric {
                DistanceMetric::Euclidean => {
                    let mut sum = 0.0f64;
                    for r in 0..d {
                        let diff = data.get(r, i) - data.get(r, j);
                        sum += diff * diff;
                    }
                    sum.sqrt()
                }
                DistanceMetric::Cosine => {
                    let ni = norms[i];
                    let nj = norms[j];
                    if ni < 1e-15 || nj < 1e-15 {
                        1.0
                    } else {
                        let dot: f64 = (0..d).map(|r| data.get(r, i) * data.get(r, j)).sum();
                        let cos = (dot / (ni * nj)).clamp(-1.0, 1.0);
                        1.0 - cos
                    }
                }
            };
            dist[i * t + j] = d_val;
            dist[j * t + i] = d_val;
        }
    }
    dist
}