numrs2 0.3.3

#![cfg(feature = "lapack")]

use crate::array::Array;
use crate::error::{NumRs2Error, Result};
use crate::new_modules::matrix_decomp::svd;
use num_traits::{Float, Zero};
use std::fmt::Debug;

/// Type alias for complex least squares return type
/// Returns (solution, residuals, rank, singular_values)
type LstsqResult<T> = Result<(
    Array<T>,
    Array<T>, // Residuals are same type as matrix elements
    usize,
    Array<T>, // Singular values are same type as matrix elements
)>;

/// Compute the condition number of a matrix using the SVD method
///
/// The condition number is the ratio of the largest to smallest singular value
/// and provides a measure of how sensitive matrix operations are to numerical errors.
/// A high condition number indicates an ill-conditioned matrix.
///
/// This implementation includes various numerical stability enhancements:
/// 1. Proper handling of very small singular values
/// 2. Scaling to avoid overflow in calculations
/// 3. Classification of different condition number ranges
pub fn condition_number<T>(a: &Array<T>) -> Result<T>
where
    T: Float + Clone + Debug,
{
    // SVD is the most numerically stable method for computing condition number
    let (_, s, _) = svd(a)?;

    // Convert to vector for easier manipulation
    let s_vec = s.to_vec();

    // Find the largest and smallest singular values
    if s_vec.is_empty() {
        return Err(NumRs2Error::ComputationError(
            "Cannot compute condition number of empty matrix".to_string(),
        ));
    }

    // Get the largest singular value
    let max_sv = s_vec
        .iter()
        .cloned()
        .fold(T::zero(), |a, b| if a > b { a } else { b });

    // Get the smallest non-zero singular value
    // We use a threshold based on machine epsilon to determine effective zeros
    let eps = T::epsilon();
    let threshold = max_sv
        * eps
        * T::from(std::cmp::max(a.shape()[0], a.shape()[1])).unwrap_or_else(|| T::one());

    // Filter out singular values effectively zero
    let non_zero_sv = s_vec
        .iter()
        .cloned()
        .filter(|&sv| sv > threshold)
        .collect::<Vec<_>>();

    if non_zero_sv.is_empty() || max_sv == T::zero() {
        // Matrix is numerically singular (all singular values effectively zero)
        return Ok(T::infinity());
    }

    // Also check if there are singular values that are almost zero
    let min_sv_all = s_vec
        .iter()
        .cloned()
        .fold(max_sv, |a, b| if a < b { a } else { b });

    // If the ratio between largest and smallest is very large, return infinity
    if max_sv / min_sv_all > T::from(1e16).unwrap_or_else(|| T::max_value()) {
        return Ok(T::infinity());
    }

    // Get the smallest non-zero singular value
    let min_sv = non_zero_sv
        .iter()
        .cloned()
        .fold(max_sv, |a, b| if a < b { a } else { b });

    // Compute the condition number as the ratio of largest to smallest singular values
    let cond = max_sv / min_sv;

    // Check for overflow and handle appropriately
    if cond.is_infinite() || cond.is_nan() {
        // If we get overflow or NaN, return a high but finite condition number
        return Ok(T::max_value());
    }

    Ok(cond)
}

/// Calculate the reciprocal condition number, which is more numerically stable
/// for very ill-conditioned matrices where the ratio might overflow.
///
/// Returns a value between 0 and 1, where values close to 0 indicate ill-conditioning,
/// and values close to 1 indicate good conditioning.
pub fn rcond<T>(a: &Array<T>) -> Result<T>
where
    T: Float + Clone + Debug,
{
    let cond = condition_number(a)?;

    // Compute the reciprocal, handling potential underflow
    if cond.is_infinite() {
        Ok(T::zero())
    } else {
        Ok(T::one() / cond)
    }
}

/// Compute the sign and (natural) logarithm of the determinant of a matrix.
///
/// This is a more numerically stable way to compute the determinant for large matrices
/// where the determinant might overflow or underflow. Returns a tuple (sign, logdet)
/// where sign is -1, 0, or 1, and logdet is the natural logarithm of the absolute
/// value of the determinant.
///
/// # Arguments
/// * `a` - Input square matrix
///
/// # Returns
/// A tuple (sign, logdet) where:
/// - sign: The sign of the determinant (-1, 0, or 1)
/// - logdet: The natural logarithm of the absolute value of the determinant
///
/// # Examples
/// ```
/// use numrs2::prelude::*;
/// use numrs2::new_modules::matrix_decomp::condition::slogdet;
///
/// let a = Array::from_vec(vec![2.0, 0.0, 0.0, 3.0]).reshape(&[2, 2]);
/// let (sign, logdet) = slogdet(&a).expect("slogdet should succeed for square matrix");
/// // For this diagonal matrix: det = 2*3 = 6, so sign = 1, logdet = ln(6) ≈ 1.79
/// ```
pub fn slogdet<T>(a: &Array<T>) -> Result<(i8, T)>
where
    T: Float
        + Clone
        + Debug
        + std::ops::AddAssign
        + std::ops::MulAssign
        + std::ops::DivAssign
        + std::ops::SubAssign
        + std::fmt::Display,
{
    // Check if the matrix is square
    let shape = a.shape();
    if shape.len() != 2 || shape[0] != shape[1] {
        return Err(NumRs2Error::DimensionMismatch(
            "slogdet requires a square matrix".to_string(),
        ));
    }

    let n = shape[0];

    // For very small matrices, compute determinant directly
    if n <= 3 {
        let det = a.det()?;

        if det == T::zero() {
            return Ok((0, T::neg_infinity()));
        }

        let sign = if det > T::zero() { 1 } else { -1 };

        let logdet = num_traits::Float::ln(num_traits::Float::abs(det));
        return Ok((sign, logdet));
    }

    // For larger matrices, use LU decomposition which is more numerically stable
    let lu_result = crate::new_modules::matrix_decomp::lu::lu(a);

    match lu_result {
        Ok((_l, u, p)) => {
            // For LU decomposition PA = LU:
            // det(A) = det(P^-1) * det(L) * det(U)
            // det(L) = 1 (since L has 1's on diagonal)
            // det(U) = product of diagonal elements
            // det(P^-1) = det(P) = (-1)^(number of swaps)

            // Count the number of swaps in the permutation
            let p_vec = p.to_vec();
            let mut visited = vec![false; n];
            let mut num_swaps = 0;

            for i in 0..n {
                if !visited[i] {
                    let mut j = i;
                    let mut cycle_length = 0;

                    while !visited[j] {
                        visited[j] = true;
                        j = p_vec[j];
                        cycle_length += 1;
                    }

                    // A cycle of length k requires k-1 swaps
                    if cycle_length > 1 {
                        num_swaps += cycle_length - 1;
                    }
                }
            }

            // Calculate the sign from permutation
            let mut sign = if num_swaps % 2 == 0 { 1i8 } else { -1i8 };

            // Calculate log absolute value of determinant from U's diagonal
            let mut logdet = T::zero();

            for i in 0..n {
                let u_diag = u.get(&[i, i])?;

                if u_diag == T::zero() {
                    // Matrix is singular
                    return Ok((0, T::neg_infinity()));
                }

                // Update sign if diagonal element is negative
                if u_diag < T::zero() {
                    sign = -sign;
                }

                // Add log of absolute value
                logdet += num_traits::Float::ln(num_traits::Float::abs(u_diag));
            }

            Ok((sign, logdet))
        }
        Err(_) => {
            // If LU decomposition fails, fall back to SVD
            let (_, s, _) = svd(a)?;
            let s_vec = s.to_vec();

            // Check if any singular value is effectively zero
            let eps = T::epsilon();
            let threshold = eps
                * T::from(n).unwrap_or_else(|| T::one())
                * s_vec
                    .iter()
                    .cloned()
                    .fold(T::zero(), |a, b| if a > b { a } else { b });

            let mut logdet = T::zero();
            let mut zero_count = 0;

            for &sv in &s_vec {
                if sv <= threshold {
                    zero_count += 1;
                } else {
                    logdet += num_traits::Float::ln(sv);
                }
            }

            if zero_count > 0 {
                // Matrix is singular or numerically singular
                Ok((0, T::neg_infinity()))
            } else {
                // For SVD, determinant is always positive since we're using |det|
                // But we need to account for the actual sign of the determinant
                // This requires more complex analysis, so we'll use a simple heuristic
                let sign = 1; // Simplified for now - both branches were identical
                Ok((sign, logdet))
            }
        }
    }
}

/// Solve a linear least-squares problem using SVD decomposition.
///
/// Computes the least-squares solution to the linear system Ax = b.
/// If the system is over-determined (more equations than unknowns), this
/// finds the solution that minimizes ||Ax - b||₂.
/// If the system is under-determined, this finds the minimum-norm solution.
///
/// # Arguments
/// * `a` - Coefficient matrix (m × n)
/// * `b` - Right-hand side vector or matrix (m × k)
/// * `rcond` - Cutoff for small singular values (relative to largest singular value).
///   Singular values smaller than rcond * largest_sv are set to zero.
///   If None, uses machine precision.
///
/// # Returns
/// A tuple (x, residuals, rank, singular_values) where:
/// - x: Least-squares solution(s)
/// - residuals: Sum of squared residuals (empty if m <= n or rank deficient)
/// - rank: Effective rank of matrix a
/// - singular_values: Singular values of a in descending order
///
/// # Examples
/// ```
/// use numrs2::prelude::*;
/// use numrs2::new_modules::matrix_decomp::condition::lstsq;
///
/// // Solve simple 2x2 system
/// let a = Array::from_vec(vec![1.0, 1.0, 1.0, 2.0]).reshape(&[2, 2]);
/// let b = Array::from_vec(vec![3.0, 4.0]);
/// let (x, residuals, rank, sv) = lstsq(&a, &b, None).expect("lstsq should succeed");
/// ```
pub fn lstsq<T>(a: &Array<T>, b: &Array<T>, rcond: Option<T>) -> LstsqResult<T>
where
    T: Float + Clone + Debug,
{
    // Check input dimensions
    let a_shape = a.shape();
    let b_shape = b.shape();

    if a_shape.len() != 2 {
        return Err(NumRs2Error::DimensionMismatch(
            "lstsq requires a 2D coefficient matrix".to_string(),
        ));
    }

    let m = a_shape[0]; // number of equations
    let n = a_shape[1]; // number of unknowns

    // Check that b has the correct number of rows
    if b_shape[0] != m {
        return Err(NumRs2Error::ShapeMismatch {
            expected: vec![m],
            actual: b_shape.clone(),
        });
    }

    // Determine if b is a vector or matrix
    let k = if b_shape.len() == 1 { 1 } else { b_shape[1] };
    let b_is_vector = b_shape.len() == 1;

    // Compute SVD of A
    let (u, s, vt) = svd(a)?;
    let s_vec = s.to_vec();

    // Determine the effective rank using rcond
    let max_sv = s_vec
        .iter()
        .cloned()
        .fold(T::zero(), |a, b| if a > b { a } else { b });

    let cutoff = match rcond {
        Some(rc) => rc * max_sv,
        None => {
            let eps = T::epsilon();
            eps * T::from(std::cmp::max(m, n)).unwrap_or_else(|| T::one()) * max_sv
        }
    };

    let rank = s_vec.iter().filter(|&&sv| sv > cutoff).count();

    // Compute the pseudo-inverse solution: x = V * S^+ * U^T * b
    // where S^+ is the Moore-Penrose pseudo-inverse of the diagonal matrix S

    // Step 1: Compute U^T * b
    let ut_b = u.transpose().matmul(b)?;

    // Step 2: Apply S^+ (multiply by 1/s_i for non-zero singular values)
    let mut s_pinv_ut_b = ut_b.clone();
    let min_dim = std::cmp::min(m, n);

    for i in 0..min_dim {
        for j in 0..k {
            let ut_b_val = ut_b.get(&[i, j])?;
            if i < s_vec.len() && s_vec[i] > cutoff {
                let sv_as_t = <T as num_traits::NumCast>::from(s_vec[i])
                    .expect("singular value should convert to float type");
                let new_val = ut_b_val / sv_as_t;
                s_pinv_ut_b.set(&[i, j], new_val)?;
            } else {
                // Zero out contributions from small singular values
                s_pinv_ut_b.set(&[i, j], T::zero())?;
            }
        }
    }

    // Step 3: Compute x = V * (S^+ * U^T * b)
    let v = vt.transpose();
    let x = v.matmul(&s_pinv_ut_b)?;

    // Reshape x if b was a vector
    let x_final = if b_is_vector && x.shape().len() == 2 && x.shape()[1] == 1 {
        let x_vec: Vec<T> = (0..x.shape()[0])
            .map(|i| x.get(&[i, 0]).expect("index should be valid"))
            .collect();
        Array::from_vec(x_vec)
    } else {
        x
    };

    // Compute residuals if the system is over-determined and full rank
    let residuals = if m > n && rank == n {
        // Compute ||Ax - b||²
        let ax = a.matmul(&x_final)?;
        let diff = if b_is_vector {
            // Both ax and b should be vectors for element-wise operations
            let ax_vec: Vec<T> = (0..ax.shape()[0])
                .map(|i| ax.get(&[i]).expect("index should be valid"))
                .collect();
            let b_vec = b.to_vec();
            let diff_vec: Vec<T> = ax_vec
                .iter()
                .zip(b_vec.iter())
                .map(|(&ax_val, &b_val)| ax_val - b_val)
                .collect();
            Array::from_vec(diff_vec)
        } else {
            let mut diff_result = Array::zeros(&b.shape());
            for i in 0..b.shape()[0] {
                for j in 0..k {
                    let ax_val = ax.get(&[i, j])?;
                    let b_val = b.get(&[i, j])?;
                    diff_result.set(&[i, j], ax_val - b_val)?;
                }
            }
            diff_result
        };

        // Compute sum of squares for each column
        let mut residuals_vec = Vec::with_capacity(k);
        for j in 0..k {
            let mut sum_sq = T::zero();
            for i in 0..m {
                let val = if b_is_vector && k == 1 {
                    diff.get(&[i])?
                } else {
                    diff.get(&[i, j])?
                };
                sum_sq = sum_sq + val * val;
            }
            residuals_vec.push(sum_sq);
        }

        // Both branches were identical, so no condition needed
        Array::from_vec(residuals_vec)
    } else {
        // No residuals for under-determined or rank-deficient systems
        Array::from_vec(vec![])
    };

    Ok((x_final, residuals, rank, s))
}