torsh-functional 0.1.2

//! Matrix decomposition operations
//!
//! This module provides fundamental matrix decompositions including LU, QR, Cholesky,
//! SVD, and eigenvalue decomposition. These decompositions are building blocks for
//! many numerical algorithms and have extensive applications in machine learning,
//! scientific computing, and engineering.

use torsh_core::{Result as TorshResult, TorshError};
use torsh_tensor::{
    creation::{eye, ones},
    Tensor,
};

use super::core::tensor_to_array2;

/// LU decomposition with partial pivoting
///
/// Computes the LU decomposition of a square matrix A:
/// ```text
/// P A = L U
/// ```text
///
/// ## Mathematical Background
///
/// LU decomposition factors a matrix into:
/// - **P**: Permutation matrix (partial pivoting for numerical stability)
/// - **L**: Lower triangular matrix with unit diagonal
/// - **U**: Upper triangular matrix
///
/// ## Applications
/// - **Linear systems**: Solve Ax = b via forward/backward substitution
/// - **Matrix inversion**: A⁻¹ = U⁻¹ L⁻¹ P^T
/// - **Determinant**: det(A) = det(P) × ∏ᵢ Uᵢᵢ
///
/// ## Parameters
/// * `tensor` - Square matrix to decompose
///
/// ## Returns
/// * Tuple (P, L, U) where P A = L U
///
/// ## Note
/// This is currently a placeholder implementation. A production version would
/// implement the actual LU decomposition algorithm with partial pivoting.
pub fn lu(tensor: &Tensor) -> TorshResult<(Tensor, Tensor, Tensor)> {
    // Validate input
    if tensor.shape().ndim() != 2 {
        return Err(TorshError::invalid_argument_with_context(
            "LU decomposition requires 2D tensor",
            "lu",
        ));
    }

    let n = tensor.shape().dims()[0];
    let m = tensor.shape().dims()[1];

    if n != m {
        return Err(TorshError::invalid_argument_with_context(
            "LU decomposition requires square matrix",
            "lu",
        ));
    }

    // This is a placeholder implementation
    // Real implementation would perform actual LU decomposition
    let p = eye(n)?;
    let l = eye(n)?;
    let u = tensor.clone();

    Ok((p, l, u))
}

/// QR Decomposition
///
/// Computes the QR decomposition of a matrix A:
/// ```text
/// A = QR
/// ```text
///
/// ## Mathematical Definition
///
/// For an m×n matrix A, the QR decomposition is a factorization:
/// - **Q**: m×m orthogonal matrix (when reduced=false) or m×min(m,n) (when reduced=true)
/// - **R**: m×n upper triangular matrix (when reduced=false) or min(m,n)×n (when reduced=true)
///
/// ## Mathematical Properties
///
/// 1. **Orthogonality**: Q^T Q = I (Q has orthonormal columns)
/// 2. **Upper triangular**: R[i,j] = 0 for i > j
/// 3. **Uniqueness**: If A has full column rank and R has positive diagonal elements,
///    then the QR decomposition is unique
/// 4. **Determinant**: det(A) = det(R) (since det(Q) = ±1)
///
/// ## Gram-Schmidt Process
///
/// The QR decomposition can be computed via Gram-Schmidt orthogonalization:
/// ```text
/// q₁ = a₁ / ||a₁||
/// q₂ = (a₂ - (q₁·a₂)q₁) / ||(a₂ - (q₁·a₂)q₁)||
/// qₖ = (aₖ - Σᵢ₌₁ᵏ⁻¹(qᵢ·aₖ)qᵢ) / ||aₖ - Σᵢ₌₁ᵏ⁻¹(qᵢ·aₖ)qᵢ||
/// ```text
///
/// ## Parameters
/// - `tensor`: Input matrix A (m×n)
/// - `reduced`: If true, returns "thin" QR with Q (m×min(m,n)) and R (min(m,n)×n).
///             If false, returns "full" QR with Q (m×m) and R (m×n).
///
/// ## Returns
/// Tuple (Q, R) where:
/// - Q: Orthogonal matrix with orthonormal columns
/// - R: Upper triangular matrix
///
/// ## Applications
/// - **Least squares**: Solve Ax = b via Rx = Q^T b
/// - **Eigenvalue algorithms**: QR iteration for eigenvalue computation
/// - **Orthogonalization**: Extract orthonormal basis from column space
/// - **Matrix inversion**: A⁻¹ = R⁻¹ Q^T (when A is square and invertible)
pub fn qr(tensor: &Tensor, reduced: bool) -> TorshResult<(Tensor, Tensor)> {
    if tensor.shape().ndim() != 2 {
        return Err(TorshError::invalid_argument_with_context(
            "QR decomposition requires 2D tensor",
            "qr",
        ));
    }

    // Use torsh-linalg's QR implementation
    // Note: torsh-linalg::decomposition::qr always returns reduced form
    // We need to adapt it if full_matrices is requested
    let (q, r) = torsh_linalg::decomposition::qr(tensor)?;

    if reduced {
        // Already in reduced form
        Ok((q, r))
    } else {
        // Need to expand Q to full m×m matrix
        let shape = tensor.shape();
        let dims = shape.dims();
        let m = dims[0];
        let n = dims[1];
        let k = m.min(n);

        if k == m {
            // Already full size
            Ok((q, r))
        } else {
            // Expand Q from m×k to m×m by padding with identity
            let mut q_data = vec![0.0f32; m * m];

            // Copy existing Q
            let q_vec = q.to_vec()?;
            for i in 0..m {
                for j in 0..k {
                    q_data[i * m + j] = q_vec[i * k + j];
                }
            }

            // Add identity for remaining columns
            for i in k..m {
                q_data[i * m + i] = 1.0;
            }

            let q_full = Tensor::from_data(q_data, vec![m, m], tensor.device())?;
            Ok((q_full, r))
        }
    }
}

/// Cholesky Decomposition
///
/// Computes the Cholesky decomposition of a positive definite matrix A:
/// ```text
/// A = L L^T  (lower form)
/// A = U^T U  (upper form)
/// ```text
///
/// ## Mathematical Definition
///
/// For a symmetric positive definite matrix A, the Cholesky decomposition uniquely factors A as:
/// - **Lower form**: A = L L^T where L is lower triangular with positive diagonal
/// - **Upper form**: A = U^T U where U is upper triangular with positive diagonal
///
/// ## Mathematical Properties
///
/// 1. **Uniqueness**: The Cholesky factor is unique for positive definite matrices
/// 2. **Efficiency**: Requires ~n³/3 operations (half of LU decomposition)
/// 3. **Numerical stability**: More stable than LU for positive definite systems
/// 4. **Determinant**: det(A) = (∏ᵢ Lᵢᵢ)² = ∏ᵢ Lᵢᵢ²
///
/// ## Cholesky Algorithm
///
/// For lower triangular L where A = L L^T:
/// ```text
/// for i = 0 to n-1:
///     Lᵢᵢ = √(Aᵢᵢ - Σⱼ₌₀ⁱ⁻¹ Lᵢⱼ²)
///     for k = i+1 to n-1:
///         Lₖᵢ = (Aₖᵢ - Σⱼ₌₀ⁱ⁻¹ Lₖⱼ Lᵢⱼ) / Lᵢᵢ
/// ```text
///
/// ## Parameters
/// - `tensor`: Input matrix A (n×n, must be symmetric positive definite)
/// - `upper`: If true, returns upper triangular U such that A = U^T U.
///           If false, returns lower triangular L such that A = L L^T.
///
/// ## Returns
/// Cholesky factor:
/// - Lower triangular L (if upper=false)
/// - Upper triangular U (if upper=true)
///
/// ## Applications
/// - **Linear systems**: Solve Ax = b via Ly = b, L^T x = y
/// - **Matrix inversion**: A⁻¹ = (L^T)⁻¹ L⁻¹
/// - **Determinant**: det(A) = ∏ᵢ Lᵢᵢ²
/// - **Gaussian sampling**: Generate x ~ N(μ, A) via x = μ + L z where z ~ N(0,I)
/// - **Optimization**: Newton's method with positive definite Hessians
pub fn cholesky(tensor: &Tensor, upper: bool) -> TorshResult<Tensor> {
    if tensor.shape().ndim() != 2 {
        return Err(TorshError::invalid_argument_with_context(
            "Cholesky decomposition requires 2D tensor",
            "cholesky",
        ));
    }

    let binding = tensor.shape();
    let dims = binding.dims();
    if dims[0] != dims[1] {
        return Err(TorshError::invalid_argument_with_context(
            "Cholesky decomposition requires square matrix",
            "cholesky",
        ));
    }

    // Use torsh-linalg's Cholesky implementation
    torsh_linalg::decomposition::cholesky(tensor, upper)
}

/// Singular Value Decomposition (SVD)
///
/// Computes the singular value decomposition of a matrix A:
/// ```text
/// A = U Σ V^T
/// ```text
///
/// ## Mathematical Definition
///
/// For an m×n matrix A, the SVD is a factorization:
/// - **U**: m×m orthogonal matrix (left singular vectors)
/// - **Σ**: m×n diagonal matrix (singular values σ₁ ≥ σ₂ ≥ ... ≥ σₘᵢₙ ≥ 0)
/// - **V^T**: n×n orthogonal matrix (right singular vectors, transposed)
///
/// ## Mathematical Properties
///
/// 1. **Orthogonality**: U^T U = I, V^T V = I
/// 2. **Rank**: rank(A) = number of non-zero singular values
/// 3. **Norm preservation**: ||A||₂ = σ₁ (largest singular value)
/// 4. **Eckart-Young theorem**: Best rank-k approximation A_k = U_k Σ_k V_k^T
///
/// ## Parameters
/// - `tensor`: Input matrix A (m×n)
/// - `full_matrices`: If true, return full U (m×m) and V^T (n×n).
///                   If false, return reduced form U (m×min(m,n)) and V^T (min(m,n)×n)
///
/// ## Returns
/// Tuple (U, Σ, V^T) where:
/// - U: Left singular vectors
/// - Σ: Singular values (1D tensor)
/// - V^T: Right singular vectors (transposed)
///
/// ## Example Applications
/// - **Matrix pseudoinverse**: A⁺ = V Σ⁺ U^T
/// - **Principal Component Analysis**: Components from U or V
/// - **Low-rank approximation**: A ≈ Σᵢ₌₁ᵏ σᵢ uᵢ vᵢ^T
/// - **Condition number**: κ(A) = σ₁/σₘᵢₙ
pub fn svd(tensor: &Tensor, full_matrices: bool) -> TorshResult<(Tensor, Tensor, Tensor)> {
    if tensor.shape().ndim() != 2 {
        return Err(TorshError::invalid_argument_with_context(
            "SVD requires 2D tensor",
            "svd",
        ));
    }

    // Use torsh-linalg's SVD implementation
    torsh_linalg::decomposition::svd(tensor, full_matrices)
}

/// Eigenvalue decomposition for square matrices
///
/// Computes the eigenvalue decomposition of a square matrix A:
/// ```text
/// A V = V Λ  or  A = V Λ V⁻¹
/// ```text
///
/// ## Mathematical Background
///
/// For a square matrix A ∈ ℝⁿˣⁿ, the eigendecomposition finds:
/// - **Eigenvalues** λᵢ: Scalars such that A vᵢ = λᵢ vᵢ
/// - **Eigenvectors** vᵢ: Non-zero vectors satisfying the eigenvalue equation
///
/// ## Properties
/// - **Characteristic polynomial**: det(A - λI) = 0
/// - **Trace**: tr(A) = Σᵢ λᵢ
/// - **Determinant**: det(A) = ∏ᵢ λᵢ
/// - **Spectral radius**: ρ(A) = max |λᵢ|
///
/// ## Parameters
/// * `tensor` - Square matrix to decompose
///
/// ## Returns
/// * Tuple (eigenvalues, eigenvectors) where eigenvalues are 1D and eigenvectors are 2D
///
/// ## Applications
/// - **Principal Component Analysis**: Covariance matrix eigendecomposition
/// - **Stability analysis**: Eigenvalues determine system stability
/// - **Matrix powers**: A^k = V Λ^k V⁻¹
/// - **Matrix functions**: f(A) = V f(Λ) V⁻¹
pub fn eig(tensor: &Tensor) -> TorshResult<(Tensor, Tensor)> {
    if tensor.shape().ndim() != 2 {
        return Err(TorshError::invalid_argument_with_context(
            "Eigenvalue decomposition requires 2D tensor",
            "eig",
        ));
    }

    let shape = tensor.shape();
    let dims = shape.dims();
    if dims[0] != dims[1] {
        return Err(TorshError::invalid_argument_with_context(
            "Eigenvalue decomposition requires square matrix",
            "eig",
        ));
    }

    // Convert tensor to ndarray for scirs2-linalg processing
    let _array = tensor_to_array2(tensor)?;

    // Simple fallback implementation for eigenvalue decomposition
    // Note: Advanced linalg dependencies not available, using basic fallback

    // For now, return identity matrix as a placeholder
    // This should be replaced with proper eigenvalue computation when dependencies are available
    let shape_binding = tensor.shape();
    let dims = shape_binding.dims();
    let eigenvalues_tensor = ones(&[dims[0]])?;
    let eigenvectors_tensor = eye(dims[0])?;

    // Return basic implementation
    Ok((eigenvalues_tensor, eigenvectors_tensor))
}

/// Low-rank SVD approximation using randomized algorithms
///
/// Computes an approximate SVD for large matrices using randomized algorithms,
/// which can be much faster than full SVD for low-rank approximations.
///
/// ## Mathematical Background
///
/// Randomized SVD uses random projections to efficiently compute low-rank
/// approximations of large matrices. The algorithm works by:
/// 1. Finding a good subspace Q that captures the range of A
/// 2. Computing B = Q^T A (smaller projected matrix)
/// 3. Computing SVD of B and mapping back to original space
///
/// ## Parameters
/// * `tensor` - Input matrix to decompose
/// * `rank` - Target rank for approximation (default: min(m,n,10))
/// * `niter` - Number of power iterations for accuracy (default: 2)
///
/// ## Returns
/// * Tuple (U, S, V) representing the low-rank SVD approximation
pub fn svd_lowrank(
    tensor: &Tensor,
    rank: Option<usize>,
    niter: Option<usize>,
) -> TorshResult<(Tensor, Tensor, Tensor)> {
    // Validate input
    if tensor.shape().ndim() != 2 {
        return Err(TorshError::invalid_argument_with_context(
            "SVD requires 2D tensor",
            "svd_lowrank",
        ));
    }

    let (m, n) = (tensor.shape().dims()[0], tensor.shape().dims()[1]);
    let k = rank.unwrap_or(m.min(n).min(10));
    let _niter = niter.unwrap_or(2);

    // This is a placeholder implementation
    // Real implementation would use randomized SVD algorithm
    use torsh_tensor::creation::randn;
    let u = randn(&[m, k])?;
    let s = ones(&[k])?;
    let v = randn(&[k, n])?;

    Ok((u, s, v))
}

/// Low-rank PCA approximation
///
/// Computes principal components using randomized SVD for efficient
/// dimensionality reduction of large datasets.
///
/// ## Mathematical Background
///
/// PCA finds the directions of maximum variance in data:
/// 1. Center the data: X̃ = X - μ
/// 2. Compute covariance: C = X̃^T X̃ / (n-1)
/// 3. Find eigenvectors: principal components are eigenvectors of C
///
/// ## Parameters
/// * `tensor` - Input data matrix (samples × features)
/// * `rank` - Number of principal components (default: min dimensions)
/// * `center` - Whether to center the data (default: true)
///
/// ## Returns
/// * Tuple (U, S, V^T) where V^T contains the principal components
pub fn pca_lowrank(
    tensor: &Tensor,
    rank: Option<usize>,
    center: bool,
) -> TorshResult<(Tensor, Tensor, Tensor)> {
    // Validate input
    if tensor.shape().ndim() != 2 {
        return Err(TorshError::invalid_argument_with_context(
            "PCA requires 2D tensor",
            "pca_lowrank",
        ));
    }

    let (m, n) = (tensor.shape().dims()[0], tensor.shape().dims()[1]);
    let k = rank.unwrap_or(m.min(n));

    // For now, simple placeholder implementation
    let mut data_tensor = tensor.clone();
    if center {
        // Center the data (subtract mean)
        // This is a simplified version - real implementation would compute actual mean
        data_tensor = tensor.clone();
    }

    // Use SVD for PCA
    let (u, s, v) = svd_lowrank(&data_tensor, Some(k), None)?;

    // For PCA, we typically return V^T as the principal components
    Ok((u, s, v.transpose(-2, -1)?))
}