#pragma once
#include "Matrix.hpp"

/**
 * @brief library that uses Matrix.hpp and performs SVD on a matrix
 */
namespace SVD {
/**
 * @brief Compute the Singular Value Decomposition (SVD) of this matrix.
 *
 * Decomposes A into U × Σ × Vᵀ where:
 *   - U is an m×k orthogonal matrix (left singular vectors)
 *   - Σ is a k×k diagonal matrix with non-negative singular values
 *     (stored as a k×1 column vector)
 *   - Vᵀ is a k×n orthogonal matrix (right singular vectors, transposed)
 *   - k = min(m, n)
 *
 * The decomposition satisfies: A ≈ U × diag(Σ) × Vᵀ
 * Singular values are returned in descending order.
 *
 * @param U Output: left singular vectors (m×k orthogonal matrix)
 * @param sigma Output: singular values as k×1 column vector, sorted descending
 * @param Vt Output: right singular vectors transposed (k×n matrix)
 *
 * @note This implementation uses the Golub-Kahan-Reinsch algorithm:
 *       1. Householder bidiagonalization of A
 *       2. Implicit QR iteration on the bidiagonal matrix
 *       3. Accumulation of U and V factors throughout
 */
template <uint8_t rows, uint8_t columns>
void SVD(Matrix<rows, columns> &matrixToDecompose, Matrix<rows, columns> &U,
         Matrix<columns, 1> &sigma, Matrix<columns, columns> &Vt);

// ========================================================================
// SVD Building Block Functions (for unit testing)
// These operate on internal 5×5 working arrays for maximum flexibility.
// ========================================================================

/**
 * @brief Compute a Householder reflector vector.
 *
 * Given input vector x, computes normalized v and scalar alpha such that:
 *   (I - 2·v·vᵀ) · x = [alpha, 0, 0, ...]ᵀ
 *
 * @param x Input vector (up to 5 elements)
 * @param len Number of valid elements in x
 * @param v Output: normalized Householder vector (v[0] is the first element)
 * @param alpha Output: the resulting first element after reflection
 * @return The norm of the input vector x
 */
static float ComputeHouseholder(const float *x, uint8_t len, float *v,
                                float &alpha);

/**
 * @brief Apply a Householder reflection from the left.
 *
 * Transforms W = (I - 2·v·vᵀ) · W where v operates on rows [startRow..endRow].
 *
 * @param W Input/output: matrix to transform (5×5 working array)
 * @param v Householder vector (length = endRow - startRow + 1)
 * @param startRow First row index
 * @param endRow Last row index
 */
static void ApplyHouseholderLeft(Matrix<5, 5> &W, const float *v,
                                 uint8_t startRow, uint8_t endRow);

/**
 * @brief Apply a Householder reflection from the right.
 *
 * Transforms W = W · (I - 2·v·vᵀ) where v operates on columns
 * [startCol..endCol].
 *
 * @param W Input/output: matrix to transform (5×5 working array)
 * @param v Householder vector (length = endCol - startCol + 1)
 * @param startCol First column index
 * @param endCol Last column index
 */
static void ApplyHouseholderRight(Matrix<5, 5> &W, const float *v,
                                  uint8_t startCol, uint8_t endCol);

/**
 * @brief Reduce a matrix to upper bidiagonal form using Householder reflections.
 *
 * Applies a sequence of Householder reflections to reduce the input matrix
 * W (m×q, where q ≥ p) to upper bidiagonal form B (p×q), accumulating
 * the left and right transformation matrices in QL and QR respectively.
 *
 * Algorithm (Golub-Kahan bidiagonalization):
 *   For k = 0 to p-1:
 *     1. Left HH on column k, rows k..m-1: zero out subdiagonal below B[k+1][k]
 *     2. Right HH on row k, cols k+2..q-1: zero out superdiagonal above B[k][k+1]
 *
 * The accumulated transformations satisfy:
 *   QLᵀ · W_original · QR = B (upper bidiagonal)
 *
 * @param W Input/output: matrix to bidiagonalize (5×5, must be at least p×q)
 * @param m Number of rows in the working matrix
 * @param q Number of columns in the working matrix (q ≥ p)
 * @param p Rank = min(m, original_columns) — number of bidiagonalization steps
 * @param QL Input/output: left Householder accumulation (initialized to identity,
 *         output: QLᵀ such that QLᵀ·W = B)
 * @param QR Input/output: right Householder accumulation (initialized to identity,
 *         output: QR such that W·QR = B after left apply)
 */
static void Bidiagonalize(Matrix<5, 5> &W,
                          uint8_t m, uint8_t q, uint8_t p,
                          Matrix<5, 5> &QL,
                          Matrix<5, 5> &QR);

/**
 * @brief Deflate a bidiagonal matrix by zeroing negligible superdiagonals.
 *
 * Scans the p×p upper-bidiagonal matrix stored in W and zeros out any
 * superdiagonal element W[i][i+1] whose magnitude is negligible relative to
 * the local diagonal scale (|W[i][i]| + |W[i+1][i+1]|). Deflating splits
 * the matrix into independent unreduced blocks that can each be solved
 * separately.
 *
 * @param W Input/output: bidiagonal matrix (5×5 working array, first p×p used)
 * @param p Size of the bidiagonal matrix (min(rows, columns))
 * @param tol Relative deflation tolerance (e.g. 1e-8f)
 */
static void DeflateBidiagonal(Matrix<5, 5> &W, uint8_t p, float tol);

/**
 * @brief Check whether a bidiagonal matrix has fully reduced to diagonal.
 *
 * Returns true when every superdiagonal element of the p×p bidiagonal
 * matrix in W is (numerically) zero, i.e. the diagonal entries are the
 * (unsorted) singular values and no unreduced blocks remain.
 *
 * @param W Input: bidiagonal matrix (5×5 working array, first p×p used)
 * @param p Size of the bidiagonal matrix (min(rows, columns))
 * @param tol Numerical zero threshold multiplier
 * @return true when all superdiagonal elements are ~0
 */
static bool BidiagonalIsDiagonal(const Matrix<5, 5> &W, uint8_t p, float tol);

/**
 * @brief Compute the full SVD of a 2×2 upper-bidiagonal block (pure).
 *
 * Decomposes B = [[a, b], [0, d]] as:
 *   B = Ublock · diag(sigma[0], sigma[1]) · Vblockᵀ
 *
 * Guarantees:
 *   - sigma[0] ≥ sigma[1] ≥ 0 (singular values, from eigenvalues of BᵀB)
 *   - Ublock and Vblock are orthogonal (columns are the left/right
 *     singular vectors respectively; Vblock = scipy's Vᵀᵀ)
 *   - Ublock · diag(sigma) · Vblockᵀ == B (within float tolerance)
 *
 * Math: eigenvectors of BᵀB = [[a², ab], [ab, b²+d²]] give the right
 * singular vectors (v1 = normalize(ab, σ1²−a²) with a safe fallback when
 * that vector is ~0; v2 = (−v1y, v1x)); left singular vectors are
 * uᵢ = B·vᵢ/σᵢ with a rank-deficiency guard: when σᵢ ≈ 0 (i.e. ~1e-30),
 * that U column is filled with the signed orthogonal complement of the
 * other U column instead of dividing by ~0.
 *
 * @param a B[0][0] (first diagonal element)
 * @param b B[0][1] (superdiagonal element)
 * @param d B[1][1] (second diagonal element)
 * @param Ublock Output: 2×2 left singular vectors (columns)
 * @param Vblock Output: 2×2 right singular vectors (columns)
 * @param sigma Output: singular values, sigma[0] ≥ sigma[1] ≥ 0
 */
static void SolveBidiagonalBlock2x2(float a, float b, float d,
                                    float Ublock[2][2], float Vblock[2][2],
                                    float sigma[2]);

/**
 * @brief Cyclic Jacobi eigenvalue algorithm for a symmetric matrix (pure).
 *
 * Reduces symmetric n×n matrix T to (near-)diagonal form IN PLACE using
 * cyclic Jacobi rotations, accumulating the eigenvectors in V.
 *
 * On return:
 *   - T's diagonal entries are the eigenvalues (off-diagonals ~0)
 *   - evals[i] = T[i][i], UNSORTED
 *   - columns of V are the corresponding eigenvectors (T·V = V·Λ)
 *
 * @param T Input/output: symmetric matrix (5×5 storage, first n×n used,
 *        destroyed in place)
 * @param n Matrix size (≤ 5)
 * @param evals Output: eigenvalues, unsorted (evals[i] = T[i][i])
 * @param V Output: eigenvector matrix, columns are eigenvectors
 */
static void JacobiEigenSymmetric(float T[5][5], uint8_t n, float evals[5],
                                 float V[5][5]);

/**
 * @brief Fold a block SVD's factors into the QL/QR accumulators.
 *
 * Given the block SVD of a bidiagonal block, B = Ublock·Σ·Vblockᵀ, the
 * accumulated Householder matrices must absorb the block factors:
 *   QL[:, blockStart..blockStart+blockSize−1] ← QL[:, ...] · Ublock
 *     (rows 0..rowsQL−1)
 *   QR[:, blockStart..blockStart+blockSize−1] ← QR[:, ...] · Vblock
 *     (rows 0..rowsQR−1)
 *
 * rowsQL / rowsQR are the meaningful row extents of the accumulators
 * (e.g. for a wide matrix W = Aᵀ, QL carries n = rows(W) meaningful
 * rows while QR is read back over its first m rows).
 *
 * @param blockStart First column/row index of the block in W
 * @param blockSize Size of the block (2, or > 2 for the Jacobi path)
 * @param Ublock Left singular-vector factor of the block (blockSize×blockSize in 5×5 storage)
 * @param Vblock Right singular-vector factor of the block (blockSize×blockSize in 5×5 storage)
 * @param rowsQL Number of meaningful rows of QL
 * @param rowsQR Number of meaningful rows of QR
 * @param QL Input/output: left transformation accumulator
 * @param QR Input/output: right transformation accumulator
 */
static void ApplyBlockFactorsToAccumulators(uint8_t blockStart,
                                            uint8_t blockSize,
                                            const float Ublock[5][5],
                                            const float Vblock[5][5],
                                            uint8_t rowsQL, uint8_t rowsQR,
                                            Matrix<5, 5> &QL,
                                            Matrix<5, 5> &QR);

/**
 * @brief Solve a bidiagonal block larger than 2×2 via Jacobi eigen of BᵀB.
 *
 * Computes the full SVD of the unreduced upper-bidiagonal block
 * W[blockStart..blockStart+blockSize−1] via eigendecomposition of the
 * tridiagonal T = BᵀB:
 *   1. Snapshot the ORIGINAL block diagonal/superdiagonal from W
 *   2. JacobiEigenSymmetric on T → eigenvalues (unsorted) + V
 *   3. Sort eigenvalues descending, reordering V
 *   4. Ublock = B_orig · V · Σ⁻¹ (from the snapshot, so W is not
 *      overwritten before Ublock is computed)
 *   5. Fold Ublock/Vblock into QL/QR via ApplyBlockFactorsToAccumulators
 *   6. Only then write sqrt(eigenvalues) into W's diagonal and zero the
 *      block's superdiagonals
 *
 * @param W Input/output: bidiagonal matrix (5×5 working array); the block's
 *        diagonal holds the singular values and its superdiagonals are
 *        zeroed on return
 * @param blockStart First column/row index of the block
 * @param blockSize Size of the block (> 2, ≤ 5)
 * @param rowsQL Number of meaningful rows of QL
 * @param rowsQR Number of meaningful rows of QR
 * @param QL Input/output: left transformation accumulator
 * @param QR Input/output: right transformation accumulator
 * @param tol (unused: Jacobi convergence tolerance is internal)
 */
static void SolveBidiagonalBlockJacobi(Matrix<5, 5> &W, uint8_t blockStart,
                                       uint8_t blockSize, uint8_t rowsQL,
                                       uint8_t rowsQR, Matrix<5, 5> &QL,
                                       Matrix<5, 5> &QR, float tol);

/**
 * @brief Extract singular values from bidiagonal matrix diagonal and sort.
 *
 * Extracts absolute values of diagonal elements of W as singular values,
 * then sorts them in descending order while reordering columns of QL
 * and QR to maintain consistency.
 *
 * @param W Input: bidiagonal matrix (5×5 working array)
 * @param sigma Output: sorted singular values (5×1 column vector, only first p used)
 * @param p Number of singular values (min(rows, columns))
 * @param QL Input/output: left transformation matrix (modified during sort)
 * @param QR Input/output: right transformation matrix (modified during sort)
 */
static void ExtractAndSortSingularValues(Matrix<5, 5> &W,
                                          Matrix<5, 1> &sigma,
                                          uint8_t p,
                                          Matrix<5, 5> &QL,
                                          Matrix<5, 5> &QR);

/**
 * @brief Assemble final U and Vt matrices from QL/QR.
 *
 * Computes the final left singular vectors (U) and right singular vectors
 * transposed (Vt) from the accumulated Householder transformations.
 *
 * For non-transpose case: U = QL[:,0:p], Vt = QR[:,0:p]ᵀ
 * For transpose case:    U = QR[:,0:p]ᵀ, Vt = QL[:,0:p]ᵀ
 *
 * @param m Number of rows in original matrix
 * @param n Number of columns in original matrix
 * @param p Rank = min(m, n)
 * @param transposeNeeded True if we computed SVD(Aᵀ) instead of SVD(A)
 * @param QL Left Householder accumulation (5×5)
 * @param QR Right Householder accumulation (5×5)
 * @param U Output: left singular vectors (m×n matrix, only first p columns used)
 * @param Vt Output: right singular vectors transposed (n×n matrix, only first p rows used)
 */
static void AssembleUAndVt(uint8_t m, uint8_t n, uint8_t p,
                           bool transposeNeeded,
                           const Matrix<5, 5> &QL,
                           const Matrix<5, 5> &QR,
                           Matrix<5, 5> &U,
                           Matrix<5, 5> &Vt);

/**
 * @brief Compute a Givens rotation that zeros out y.
 *
 * Computes c, s such that:
 *   [c  s] [x] = [r]
 *   [-s c] [y]   [0]
 * where r = sqrt(x² + y²).
 *
 * @param x First element
 * @param y Second element (to be zeroed)
 * @param c Output: cosine of rotation angle
 * @param s Output: sine of rotation angle
 */
static void ComputeGivens(float x, float y, float &c, float &s);

/**
 * @brief Apply a Givens rotation from the left to rows i and j.
 *
 * Applies [c s; -s c] to rows i, j of W (columns startCol..endCol).
 *
 * @param W Input/output: matrix to transform
 * @param i First row index
 * @param j Second row index
 * @param c Cosine of rotation angle
 * @param s Sine of rotation angle
 * @param startCol First column to transform
 * @param endCol Last column to transform
 */
static void ApplyGivensLeft(Matrix<5, 5> &W, uint8_t i, uint8_t j, float c,
                            float s, uint8_t startCol, uint8_t endCol);

/**
 * @brief Apply a Givens rotation from the right to columns i and j.
 *
 * Applies [c -s; s c]ᵀ to columns i, j of W (rows startRow..endRow).
 *
 * @param W Input/output: matrix to transform
 * @param i First column index
 * @param j Second column index
 * @param c Cosine of rotation angle
 * @param s Sine of rotation angle
 * @param startRow First row to transform
 * @param endRow Last row to transform
 */
static void ApplyGivensRight(Matrix<5, 5> &W, uint8_t i, uint8_t j, float c,
                             float s, uint8_t startRow, uint8_t endRow);
} // namespace SVD

#ifndef SVD_H_
#include "SVD.cpp"
#endif // SVD_H_