sklears-manifold 0.1.2

//! Local Tangent Space Alignment (LTSA) implementation
//!
//! This module provides LTSA for non-linear dimensionality reduction through local tangent space alignment.

use scirs2_core::ndarray::{Array2, ArrayView2, Axis};
use scirs2_linalg::compat::{ArrayLinalgExt, UPLO};
use sklears_core::{
    error::{Result as SklResult, SklearsError},
    traits::{Estimator, Fit, Transform, Untrained},
    types::Float,
};

/// Local Tangent Space Alignment (LTSA)
///
/// LTSA constructs a neighborhood-preserving mapping by aligning
/// local tangent spaces learned on neighborhood patches. It uses
/// local PCA to compute tangent spaces and then aligns them globally.
///
/// # Parameters
///
/// * `n_neighbors` - Number of neighbors to consider for each point
/// * `n_components` - Number of coordinates for the manifold
/// * `reg` - Regularization constant for computation
/// * `eigen_solver` - The eigensolver to use
/// * `tol` - Tolerance for convergence
/// * `max_iter` - Maximum number of iterations
/// * `neighbors_algorithm` - Algorithm to use for nearest neighbors search
/// * `random_state` - Random state for reproducibility
/// * `n_jobs` - Number of parallel jobs
///
/// # Examples
///
/// ```rust,ignore
/// use sklears_manifold::LTSA;
/// use sklears_core::traits::{Transform, Fit};
/// use scirs2_core::ndarray::array;
///
/// let x = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0]];
///
/// let ltsa = LTSA::new()
///     .n_neighbors(4)
///     .n_components(2);
/// let fitted = ltsa.fit(&x.view(), &()).unwrap();
/// let embedded = fitted.transform(&x.view()).unwrap();
/// ```
#[derive(Debug, Clone)]
pub struct LTSA<S = Untrained> {
    state: S,
    n_neighbors: usize,
    n_components: usize,
    reg: f64,
    eigen_solver: String,
    tol: f64,
    max_iter: Option<usize>,
    neighbors_algorithm: String,
    random_state: Option<u64>,
    n_jobs: Option<i32>,
}

/// Trained state for LTSA
#[derive(Debug, Clone)]
pub struct LtsaTrained {
    /// The low-dimensional embedding of the training data
    pub embedding: Array2<f64>,
    /// The global alignment matrix
    pub alignment_matrix: Array2<f64>,
    /// Local tangent spaces for each point
    pub local_tangent_spaces: Vec<Array2<f64>>,
}

impl LTSA<Untrained> {
    /// Create a new LTSA instance
    pub fn new() -> Self {
        Self {
            state: Untrained,
            n_neighbors: 5,
            n_components: 2,
            reg: 1e-3,
            eigen_solver: "auto".to_string(),
            tol: 1e-6,
            max_iter: Some(100),
            neighbors_algorithm: "auto".to_string(),
            random_state: None,
            n_jobs: None,
        }
    }

    /// Set the number of neighbors
    pub fn n_neighbors(mut self, n_neighbors: usize) -> Self {
        self.n_neighbors = n_neighbors;
        self
    }

    /// Set the number of components
    pub fn n_components(mut self, n_components: usize) -> Self {
        self.n_components = n_components;
        self
    }

    /// Set the regularization constant
    pub fn reg(mut self, reg: f64) -> Self {
        self.reg = reg;
        self
    }

    /// Set the eigen solver
    pub fn eigen_solver(mut self, eigen_solver: &str) -> Self {
        self.eigen_solver = eigen_solver.to_string();
        self
    }

    /// Set the tolerance
    pub fn tol(mut self, tol: f64) -> Self {
        self.tol = tol;
        self
    }

    /// Set the maximum iterations
    pub fn max_iter(mut self, max_iter: Option<usize>) -> Self {
        self.max_iter = max_iter;
        self
    }

    /// Set the neighbors algorithm
    pub fn neighbors_algorithm(mut self, neighbors_algorithm: &str) -> Self {
        self.neighbors_algorithm = neighbors_algorithm.to_string();
        self
    }

    /// Set the random state
    pub fn random_state(mut self, random_state: Option<u64>) -> Self {
        self.random_state = random_state;
        self
    }

    /// Set the number of jobs
    pub fn n_jobs(mut self, n_jobs: Option<i32>) -> Self {
        self.n_jobs = n_jobs;
        self
    }
}

impl Default for LTSA<Untrained> {
    fn default() -> Self {
        Self::new()
    }
}

impl Estimator for LTSA<Untrained> {
    type Config = ();
    type Error = SklearsError;
    type Float = Float;

    fn config(&self) -> &Self::Config {
        &()
    }
}

impl Fit<ArrayView2<'_, Float>, ()> for LTSA<Untrained> {
    type Fitted = LTSA<LtsaTrained>;

    fn fit(self, x: &ArrayView2<'_, Float>, _y: &()) -> SklResult<Self::Fitted> {
        let x = x.mapv(|x| x);
        let (n_samples, n_features) = x.dim();

        if n_samples <= self.n_components {
            return Err(SklearsError::InvalidInput(
                "Number of samples must be greater than n_components".to_string(),
            ));
        }

        if self.n_neighbors >= n_samples {
            return Err(SklearsError::InvalidInput(
                "n_neighbors must be less than number of samples".to_string(),
            ));
        }

        if self.n_neighbors <= self.n_components {
            return Err(SklearsError::InvalidInput(
                "LTSA requires n_neighbors > n_components".to_string(),
            ));
        }

        // Step 1: Find k-nearest neighbors for each point
        let neighbor_indices = self.find_neighbors(&x)?;

        // Step 2: Compute local tangent spaces
        let local_tangent_spaces =
            self.compute_local_tangent_spaces(&x, &neighbor_indices, n_features)?;

        // Step 3: Align local coordinate systems
        let alignment_matrix =
            self.compute_alignment_matrix(&x, &neighbor_indices, &local_tangent_spaces)?;

        // Step 4: Find global embedding
        let embedding = self.compute_global_embedding(&alignment_matrix)?;

        Ok(LTSA {
            state: LtsaTrained {
                embedding,
                alignment_matrix,
                local_tangent_spaces,
            },
            n_neighbors: self.n_neighbors,
            n_components: self.n_components,
            reg: self.reg,
            eigen_solver: self.eigen_solver,
            tol: self.tol,
            max_iter: self.max_iter,
            neighbors_algorithm: self.neighbors_algorithm,
            random_state: self.random_state,
            n_jobs: self.n_jobs,
        })
    }
}

impl LTSA<Untrained> {
    fn find_neighbors(&self, x: &Array2<f64>) -> SklResult<Array2<usize>> {
        let n_samples = x.nrows();
        let mut neighbor_indices = Array2::zeros((n_samples, self.n_neighbors));

        for i in 0..n_samples {
            let mut distances: Vec<(f64, usize)> = Vec::new();
            for j in 0..n_samples {
                if i != j {
                    let diff = &x.row(i) - &x.row(j);
                    let dist = diff.mapv(|x| x * x).sum().sqrt();
                    distances.push((dist, j));
                }
            }

            // Sort by distance and take k nearest neighbors
            distances.sort_by(|a, b| a.0.partial_cmp(&b.0).expect("operation should succeed"));
            for (neighbor_idx, &(_, j)) in distances.iter().take(self.n_neighbors).enumerate() {
                neighbor_indices[[i, neighbor_idx]] = j;
            }
        }

        Ok(neighbor_indices)
    }

    fn compute_local_tangent_spaces(
        &self,
        x: &Array2<f64>,
        neighbor_indices: &Array2<usize>,
        n_features: usize,
    ) -> SklResult<Vec<Array2<f64>>> {
        let n_samples = x.nrows();
        let mut local_tangent_spaces = Vec::with_capacity(n_samples);

        for i in 0..n_samples {
            // Extract neighborhood
            let neighbors: Vec<usize> = (0..self.n_neighbors)
                .map(|j| neighbor_indices[[i, j]])
                .collect();

            // Create neighborhood matrix
            let mut neighborhood = Array2::zeros((self.n_neighbors, n_features));
            for (k, &neighbor_idx) in neighbors.iter().enumerate() {
                for d in 0..n_features {
                    neighborhood[[k, d]] = x[[neighbor_idx, d]];
                }
            }

            // Center the neighborhood
            let mean = neighborhood
                .mean_axis(Axis(0))
                .expect("operation should succeed");
            for k in 0..self.n_neighbors {
                for d in 0..n_features {
                    neighborhood[[k, d]] -= mean[d];
                }
            }

            // Compute local tangent space via SVD
            let (_, s, vt) = neighborhood
                .svd(true)
                .map_err(|e| SklearsError::InvalidInput(format!("SVD failed: {e}")))?;

            let vt_matrix = vt;
            // Take the first n_components principal components as tangent space
            let mut tangent_space = Array2::zeros((self.n_components, n_features));
            for comp in 0..self.n_components {
                if comp < s.len() && s[comp] > 1e-12 {
                    for d in 0..n_features {
                        tangent_space[[comp, d]] = vt_matrix[[comp, d]];
                    }
                }
            }
            local_tangent_spaces.push(tangent_space);
        }

        Ok(local_tangent_spaces)
    }

    fn compute_alignment_matrix(
        &self,
        x: &Array2<f64>,
        neighbor_indices: &Array2<usize>,
        local_tangent_spaces: &[Array2<f64>],
    ) -> SklResult<Array2<f64>> {
        let n_samples = x.nrows();
        let mut alignment_matrix = Array2::zeros((n_samples, n_samples));

        for i in 0..n_samples {
            let neighbors: Vec<usize> = (0..self.n_neighbors)
                .map(|j| neighbor_indices[[i, j]])
                .collect();

            // Compute local coordinate representation
            let tangent_space = &local_tangent_spaces[i];

            // Project neighborhood onto local tangent space
            let mut local_coords = Array2::zeros((self.n_neighbors, self.n_components));

            // Get the mean of the neighborhood
            let mut center = Array2::<f64>::zeros((1, x.ncols()));
            for &neighbor_idx in &neighbors {
                for d in 0..x.ncols() {
                    center[[0, d]] += x[[neighbor_idx, d]];
                }
            }
            for d in 0..x.ncols() {
                center[[0, d]] /= self.n_neighbors as f64;
            }

            for (k, &neighbor_idx) in neighbors.iter().enumerate() {
                for comp in 0..self.n_components {
                    let mut coord = 0.0;
                    for d in 0..x.ncols() {
                        coord += (x[[neighbor_idx, d]] - center[[0, d]]) * tangent_space[[comp, d]];
                    }
                    local_coords[[k, comp]] = coord;
                }
            }

            // Compute reconstruction weights that preserve local geometry
            let weights = self.compute_reconstruction_weights(&local_coords)?;

            // Add to alignment matrix
            for (a, &neighbor_a) in neighbors.iter().enumerate() {
                for (b, &neighbor_b) in neighbors.iter().enumerate() {
                    alignment_matrix[[neighbor_a, neighbor_b]] += weights[[a, b]];
                }
            }
        }

        Ok(alignment_matrix)
    }

    fn compute_reconstruction_weights(&self, local_coords: &Array2<f64>) -> SklResult<Array2<f64>> {
        let n_neighbors = local_coords.nrows();
        let mut weights = Array2::zeros((n_neighbors, n_neighbors));

        // Simplified weight computation
        // In practice, this would involve solving for optimal alignment weights
        for i in 0..n_neighbors {
            weights[[i, i]] = 1.0;
            for j in 0..n_neighbors {
                if i != j {
                    // Compute similarity based on local coordinate distance
                    let mut dist_sq = 0.0;
                    for comp in 0..self.n_components {
                        let diff = local_coords[[i, comp]] - local_coords[[j, comp]];
                        dist_sq += diff * diff;
                    }
                    weights[[i, j]] = (-dist_sq / (2.0 * self.reg)).exp();
                }
            }
        }

        // Normalize weights
        for i in 0..n_neighbors {
            let row_sum: f64 = weights.row(i).sum();
            if row_sum > 1e-12 {
                for j in 0..n_neighbors {
                    weights[[i, j]] /= row_sum;
                }
            }
        }

        Ok(weights)
    }

    fn compute_global_embedding(&self, alignment_matrix: &Array2<f64>) -> SklResult<Array2<f64>> {
        let n_samples = alignment_matrix.nrows();

        // Eigendecomposition of alignment matrix
        let (eigenvals, eigenvecs) = alignment_matrix
            .eigh(UPLO::Lower)
            .map_err(|e| SklearsError::InvalidInput(format!("Eigendecomposition failed: {e}")))?;

        // Sort eigenvalues and eigenvectors in ascending order
        let mut eigen_pairs: Vec<(f64, usize)> = eigenvals
            .iter()
            .enumerate()
            .map(|(i, &val)| (val, i))
            .collect();
        eigen_pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).expect("operation should succeed"));

        // Take eigenvectors corresponding to smallest non-zero eigenvalues
        // Skip the first eigenvector (corresponding to eigenvalue 0)
        let mut embedding = Array2::zeros((n_samples, self.n_components));
        for (comp_idx, &(eigenval, eigen_idx)) in eigen_pairs
            .iter()
            .skip(1)
            .take(self.n_components)
            .enumerate()
        {
            if eigenval > 1e-12 {
                for i in 0..n_samples {
                    embedding[[i, comp_idx]] = eigenvecs[[i, eigen_idx]];
                }
            }
        }

        Ok(embedding)
    }
}

impl Transform<ArrayView2<'_, Float>, Array2<Float>> for LTSA<LtsaTrained> {
    fn transform(&self, _x: &ArrayView2<'_, Float>) -> SklResult<Array2<Float>> {
        // LTSA doesn't support transforming new data in this implementation
        Err(SklearsError::InvalidOperation(
            "LTSA does not support transforming new data. Use fit_transform for training data."
                .to_string(),
        ))
    }
}

impl LTSA<LtsaTrained> {
    /// Get the embedding
    pub fn embedding(&self) -> &Array2<f64> {
        &self.state.embedding
    }

    /// Get the global alignment matrix
    pub fn alignment_matrix(&self) -> &Array2<f64> {
        &self.state.alignment_matrix
    }

    /// Get the local tangent spaces
    pub fn local_tangent_spaces(&self) -> &[Array2<f64>] {
        &self.state.local_tangent_spaces
    }
}