// This #ifndef section makes clangd happy so that it can properly do type hints
// in this file
#ifndef SVD_H_
#define SVD_H_
#include "SVD.hpp"
#endif

#ifdef SVD_H_ // since the .cpp file has to be included by the .hpp file this
              // will evaluate to true
#include "SVD.hpp"
#include <cstdint>

// ============================================================================
// SVD Building Block Implementations
// ============================================================================

float SVD::ComputeHouseholder(const float *x, uint8_t len, float *v,
                              float &alpha) {
  // Compute ||x||
  float norm = 0.0f;
  for (uint8_t i = 0; i < len; i++) {
    norm += x[i] * x[i];
  }
  norm = sqrtf(norm);

  if (norm < 1e-30f) {
    alpha = 0.0f;
    for (uint8_t i = 0; i < len; i++) {
      v[i] = 0.0f;
    }
    return 0.0f;
  }

  // Choose sign to avoid cancellation: alpha has opposite sign of x[0]
  alpha = (x[0] >= 0.0f) ? -norm : norm;

  // v = x - alpha * e1, then normalize
  float v0 = x[0] - alpha;

  // Compute ||v||² directly: v0² + x₁² + ... + xₙ₋₁²
  float vv = v0 * v0;
  for (uint8_t i = 1; i < len; i++) {
    vv += x[i] * x[i];
  }

  if (vv < 1e-30f) {
    // Already aligned with e1
    for (uint8_t i = 0; i < len; i++) {
      v[i] = (i == 0) ? 1.0f : 0.0f;
    }
    return norm;
  }

  float scale = 1.0f / sqrtf(vv);
  for (uint8_t i = 0; i < len; i++) {
    v[i] = (i == 0) ? v0 * scale : x[i] * scale;
  }

  return norm;
}

void SVD::ApplyHouseholderLeft(Matrix<5, 5> &W, const float *v,
                               uint8_t startRow, uint8_t endRow) {
  uint8_t len = endRow - startRow + 1;

  // Compute vᵀv (should be 2.0 for our normalized vectors, but compute
  // explicitly)
  float vv = 0.0f;
  for (uint8_t i = 0; i < len; i++) {
    vv += v[i] * v[i];
  }
  if (vv < 1e-30f)
    return;

  float twoOverVv = 2.0f / vv;

  // W = (I - 2vvᵀ) · W
  for (uint8_t col = 0; col < 5; col++) {
    float dot = 0.0f;
    for (uint8_t i = 0; i < len; i++) {
      dot += v[i] * W[startRow + i][col];
    }
    dot *= twoOverVv;
    for (uint8_t i = 0; i < len; i++) {
      W[startRow + i][col] -= dot * v[i];
    }
  }
}

void SVD::ApplyHouseholderRight(Matrix<5, 5> &W, const float *v,
                                uint8_t startCol, uint8_t endCol) {
  uint8_t len = endCol - startCol + 1;

  float vv = 0.0f;
  for (uint8_t i = 0; i < len; i++) {
    vv += v[i] * v[i];
  }
  if (vv < 1e-30f)
    return;

  float twoOverVv = 2.0f / vv;

  // W = W · (I - 2vvᵀ)
  for (uint8_t row = 0; row < 5; row++) {
    float dot = 0.0f;
    for (uint8_t i = 0; i < len; i++) {
      dot += W[row][startCol + i] * v[i];
    }
    dot *= twoOverVv;
    for (uint8_t i = 0; i < len; i++) {
      W[row][startCol + i] -= dot * v[i];
    }
  }
}

void SVD::ComputeGivens(float x, float y, float &c, float &s) {
  float r = sqrtf(x * x + y * y);

  if (r < 1e-30f) {
    c = 1.0f;
    s = 0.0f;
    return;
  }

  c = x / r;
  s = y / r;
}

void SVD::ApplyGivensLeft(Matrix<5, 5> &W, uint8_t i, uint8_t j, float c,
                          float s, uint8_t startCol, uint8_t endCol) {
  // [c  s] [row_i]   = [new_row_i]
  // [-s c] [row_j]     [new_row_j]
  for (uint8_t col = startCol; col <= endCol && col < 5; col++) {
    float t1 = W[i][col];
    float t2 = W[j][col];
    W[i][col] = c * t1 + s * t2;
    W[j][col] = -s * t1 + c * t2;
  }
}

void SVD::ApplyGivensRight(Matrix<5, 5> &W, uint8_t i, uint8_t j, float c,
                           float s, uint8_t startRow, uint8_t endRow) {
  // [col_i  col_j] · [c -s] = [new_col_i  new_col_j]
  //                  [s  c]
  for (uint8_t row = startRow; row <= endRow && row < 5; row++) {
    float t1 = W[row][i];
    float t2 = W[row][j];
    W[row][i] = c * t1 + s * t2;
    W[row][j] = -s * t1 + c * t2;
  }
}

// ============================================================================
// SVD Implementation - Golub-Kahan-Reinsch Algorithm
// ============================================================================
/**
 * @brief SVD for any m×n matrix using Householder bidiagonalization +
 *        implicit QR iteration on the bidiagonal form.
 *
 * Given A (m×n), computes U (m×k), Σ (k×k diagonal), Vᵀ (k×n) where
 * k = min(m,n) and A = U·Σ·Vᵀ.
 *
 * We store results as:
 *   - U: Matrix<m, n> — first k columns are meaningful
 *   - sigma: Matrix<n, 1> — first k entries are non-zero singular values
 *   - Vt: Matrix<n, n> — first k rows are meaningful
 *
 * For m < n (wide matrices), we work with Aᵀ and swap roles of U and V.
 */
template <uint8_t rows, uint8_t columns>
void SVD::SVD(Matrix<rows, columns> &matrixToDecompose,
              Matrix<rows, columns> &U, Matrix<columns, 1> &sigma,
              Matrix<columns, columns> &Vt) {
  static_assert(rows <= 5 && columns <= 5,
                "SVD currently supports matrices up to 5×5");

  uint8_t m = rows;
  uint8_t n = columns;
  uint8_t p = (m < n) ? m : n; // rank = min(m,n)

  // For wide matrices (m < n), work with Aᵀ instead.
  // SVD(A) = U·Σ·Vᵀ ⟺ SVD(Aᵀ) = V·Σ·Uᵀ
  // So if we compute SVD(Aᵀ) = Ũ·Σ·Ṽᵀ, then U = Ṽ and Vt = Ũᵀ.
  bool transposeNeeded = (m < n);

  // Working matrix: always p×p or larger square
  Matrix<5, 5> W{0};
  for (uint8_t i = 0; i < m; i++) {
    for (uint8_t j = 0; j < n; j++) {
      float val = matrixToDecompose.Get(i, j);
      if (transposeNeeded) {
        W[j][i] = val; // store Aᵀ
      } else {
        W[i][j] = val;
      }
    }
  }

  // After bidiagonalization, W holds the bidiagonal matrix B.
  // Q_L and Q_R accumulate the Householder transformations.
  Matrix<5, 5> QL{0}, QR{0};
  for (uint8_t i = 0; i < 5; i++) {
    QL[i][i] = 1;
    QR[i][i] = 1;
  }

  // ---- Phase 1: Householder Bidiagonalization ----
  // Reduce W to upper bidiagonal form using Householder reflections.
  // For a p×q matrix (p ≤ q after transpose), we do p steps:
  //   Step k: zero out subdiagonal in column k, then superdiagonal in row k

  float hhVec[5]; // Householder vector storage

  for (uint8_t k = 0; k < p; k++) {
    // --- Left Householder on column k, rows k..min(m,n)-1 ---
    {
      uint8_t len = (m > n) ? m - k : n - k;
      if (len <= 1)
        continue;

      // Extract the column segment
      float x[5];
      for (uint8_t i = 0; i < len; i++) {
        x[i] = W[k + i][k];
      }

      // Compute Householder reflection: H·x = [α, 0, ..., 0]ᵀ
      float norm = 0;
      for (uint8_t i = 0; i < len; i++)
        norm += x[i] * x[i];
      norm = sqrtf(norm);

      if (norm < 1e-30f)
        continue;

      float alpha = (x[0] >= 0) ? -norm : norm;
      float v0 = x[0] - alpha;
      float vv = v0 * v0 + norm * norm - alpha * x[0];
      if (vv < 1e-30f)
        continue;

      float scale = 1.0f / sqrtf(vv);
      // Store Householder vector (first element is implicit 1, rest in hhVec)
      hhVec[0] =
          v0 * scale; // this is the first element of the reflected vector
      for (uint8_t i = 1; i < len; i++) {
        hhVec[i] = x[i] * scale;
      }

      // Apply H from left to W: W = H·W
      // For each column j, w[k+i][j] -= 2·v_i·(vᵀ·w_col) / (vᵀv)
      float vvNorm = 1.0f + hhVec[0] * hhVec[0];
      for (uint8_t i = 1; i < len; i++) {
        vvNorm += hhVec[i] * hhVec[i];
      }

      for (uint8_t j = k; j < n; j++) {
        float dot = 0;
        for (uint8_t i = 0; i < len; i++) {
          dot += hhVec[i] * W[k + i][j];
        }
        dot *= 2.0f / vvNorm;
        for (uint8_t i = 0; i < len; i++) {
          W[k + i][j] -= dot * hhVec[i];
        }
      }

      // Apply H from right to QL: QL = QL · H
      for (uint8_t j = k; j < m; j++) {
        float dot = 0;
        for (uint8_t i = 0; i < len; i++) {
          dot += hhVec[i] * QL[j][k + i];
        }
        dot *= 2.0f / vvNorm;
        for (uint8_t i = 0; i < len; i++) {
          QL[j][k + i] -= dot * hhVec[i];
        }
      }
    }

    // --- Right Householder on row k, columns k+1..min(m,n)-1 ---
    {
      uint8_t len = (p > 1) ? p - 1 - k : 0;
      if (len <= 0)
        continue;

      // Extract the row segment
      float x[5];
      for (uint8_t i = 0; i < len; i++) {
        x[i] = W[k][k + 1 + i];
      }

      // Compute Householder reflection
      float norm = 0;
      for (uint8_t i = 0; i < len; i++)
        norm += x[i] * x[i];
      norm = sqrtf(norm);

      if (norm < 1e-30f)
        continue;

      float alpha = (x[0] >= 0) ? -norm : norm;
      float v0 = x[0] - alpha;
      float vv = v0 * v0 + norm * norm - alpha * x[0];
      if (vv < 1e-30f)
        continue;

      float scale = 1.0f / sqrtf(vv);
      hhVec[0] = v0 * scale;
      for (uint8_t i = 1; i < len; i++) {
        hhVec[i] = x[i] * scale;
      }

      // Compute vᵀv
      float vvNorm = 1.0f + hhVec[0] * hhVec[0];
      for (uint8_t i = 1; i < len; i++) {
        vvNorm += hhVec[i] * hhVec[i];
      }

      // Apply H from right to W: W = W·H
      for (uint8_t i = 0; i < m; i++) {
        float dot = 0;
        for (uint8_t j = 0; j < len; j++) {
          dot += hhVec[j] * W[i][k + 1 + j];
        }
        dot *= 2.0f / vvNorm;
        for (uint8_t j = 0; j < len; j++) {
          W[i][k + 1 + j] -= dot * hhVec[j];
        }
      }

      // Apply H from right to QR: QR = QR · H
      for (uint8_t i = 0; i < n; i++) {
        float dot = 0;
        for (uint8_t j = 0; j < len; j++) {
          dot += hhVec[j] * QR[i][k + 1 + j];
        }
        dot *= 2.0f / vvNorm;
        for (uint8_t j = 0; j < len; j++) {
          QR[i][k + 1 + j] -= dot * hhVec[j];
        }
      }
    }
  }

  // ---- Phase 2: Implicit QR Iteration on Bidiagonal Matrix ----
  // W now contains the upper bidiagonal matrix B.
  // We apply implicit QR steps to diagonalize it.

  uint32_t maxIter = 1000;
  float tol = 1e-8f;

  for (uint32_t iter = 0; iter < maxIter; iter++) {
    // Deflate: zero out negligible subdiagonal elements
    for (uint8_t i = p - 1; i > 0; i--) {
      float test = fabsf(W[i][i - 1]);
      float scale = fabsf(W[i - 1][i - 1]) + fabsf(W[i][i]);
      if (test < tol * (scale + 1e-30f)) {
        W[i][i - 1] = 0;
      }
    }

    // Find the smallest unreduced block [start..end]
    uint8_t start = 0, end = p - 1;
    for (uint8_t i = 0; i < p - 1; i++) {
      if (fabsf(W[i + 1][i]) >
          tol * (fabsf(W[i][i]) + fabsf(W[i + 1][i + 1]) + 1e-30f)) {
        start = i + 1;
      }
    }
    for (int8_t i = (int8_t)p - 2; i >= 0; i--) {
      if (fabsf(W[i + 1][i]) >
          tol * (fabsf(W[i][i]) + fabsf(W[i + 1][i + 1]) + 1e-30f)) {
        end = (uint8_t)i;
        break;
      }
    }

    // Check convergence of the block
    if (start >= end) {
      continue;
    }

    bool blockConverged = true;
    for (uint8_t i = start; i <= end; i++) {
      if (i > start && fabsf(W[i][i - 1]) > tol * (fabsf(W[i - 1][i - 1]) +
                                                   fabsf(W[i][i]) + 1e-30f)) {
        blockConverged = false;
        break;
      }
      if (i < end &&
          fabsf(W[i][i + 1]) >
              tol * (fabsf(W[i][i]) + fabsf(W[i + 1][i + 1]) + 1e-30f)) {
        blockConverged = false;
        break;
      }
    }
    if (blockConverged)
      continue;

    // Wilkinson shift from bottom 2×2 corner
    float a = W[end - 1][end - 1];
    float b = W[end - 1][end];
    float c = W[end][end - 1];
    float d = W[end][end];

    float trace = a + d;
    float det = a * d - b * c;
    float disc = trace * trace - 4.0f * det;
    float shift;
    if (disc >= 0) {
      float sqrtDisc = sqrtf(disc);
      float e1 = (trace + sqrtDisc) / 2.0f;
      float e2 = (trace - sqrtDisc) / 2.0f;
      shift = (fabsf(e1 - d) < fabsf(e2 - d)) ? e1 : e2;
    } else {
      shift = d;
    }

    // --- Implicit QR step using Givens rotations ---
    // First, apply Givens rotation from the left to zero out (W[start][start-1]
    // - shift) For the bidiagonal structure, we process from top to bottom.

    float x = W[start][start] - shift;
    float y = (start > 0) ? W[start][start - 1] : 0.0f;

    for (uint8_t i = start; i <= end; i++) {
      float r = sqrtf(x * x + y * y);
      if (r < 1e-30f) {
        x = W[i][i];
        y = (i < end) ? W[i + 1][i] : 0.0f;
        continue;
      }

      float cs = x / r;
      float sn = y / r;

      // Apply Givens from left to rows i, i+1 of W (columns i..p-1)
      for (uint8_t j = i; j < p; j++) {
        float t1 = W[i][j];
        float t2 = W[i + 1][j];
        W[i][j] = cs * t1 + sn * t2;
        W[i + 1][j] = -sn * t1 + cs * t2;
      }

      // Apply Givens from right to columns i, i+1 of W (rows 0..i)
      if (i > start) {
        for (uint8_t j = 0; j <= i; j++) {
          float t1 = W[j][i];
          float t2 = W[j][i + 1];
          W[j][i] = cs * t1 + sn * t2;
          W[j][i + 1] = -sn * t1 + cs * t2;
        }
      }

      // Accumulate right transformations into QR
      for (uint8_t j = 0; j < n; j++) {
        float t1 = QR[j][i];
        float t2 = QR[j][i + 1];
        QR[j][i] = cs * t1 + sn * t2;
        QR[j][i + 1] = -sn * t1 + cs * t2;
      }

      // Prepare next Givens rotation
      x = W[i + 1][i];
      y = (i + 1 < end) ? W[i + 1][i + 1] : 0.0f;
    }
  }

  // ---- Phase 3: Extract Results ----
  // Singular values are the absolute values of diagonal elements of W
  for (uint8_t i = 0; i < p; i++) {
    sigma[i][0] = fabsf(W[i][i]);
  }

  // Sort singular values in descending order and reorder U, V accordingly
  for (uint8_t i = 0; i < p - 1; i++) {
    for (uint8_t j = i + 1; j < p; j++) {
      if (sigma[j][0] > sigma[i][0]) {
        float tmpS = sigma[i][0];
        sigma[i][0] = sigma[j][0];
        sigma[j][0] = tmpS;

        // Swap columns of QL
        for (uint8_t k = 0; k < 5; k++) {
          float tmpQ = QL[k][i];
          QL[k][i] = QL[k][j];
          QL[k][j] = tmpQ;
        }

        // Swap columns of QR
        for (uint8_t k = 0; k < 5; k++) {
          float tmpQ = QR[k][i];
          QR[k][i] = QR[k][j];
          QR[k][j] = tmpQ;
        }
      }
    }
  }

  // ---- Phase 4: Compute Final U and Vt ----
  // If transposeNeeded (wide matrix), we computed SVD(Aᵀ) = Ũ·Σ·Ṽᵀ
  //   Then U = Ṽ (= QR[:,0:p]) and Vt = Ũᵀ (= QL[:,0:p]ᵀ)
  // Otherwise, SVD(A) = QL[:,0:p] · Σ · (QR[:,0:p])ᵀ
  //   So U = QL[:,0:p] and Vt = QR[:,0:p]ᵀ

  for (uint8_t i = 0; i < m; i++) {
    for (uint8_t j = 0; j < n; j++) {
      if (j < p) {
        if (transposeNeeded) {
          // U = QR[:, 0:p]ᵀ → U[i][j] = QR[j][i]
          U[i][j] = QR[j][i];
        } else {
          // U = QL[:, 0:p]
          U[i][j] = QL[i][j];
        }
      } else {
        U[i][j] = 0;
      }
    }
  }

  for (uint8_t i = 0; i < n; i++) {
    for (uint8_t j = 0; j < m; j++) {
      if (i < p && j < m) {
        if (transposeNeeded) {
          // Vt = QL[:, 0:p]ᵀ → Vt[i][j] = QL[j][i]
          Vt[i][j] = QL[j][i];
        } else {
          // Vt = QR[:, 0:p]ᵀ → Vt[i][j] = QR[j][i]
          Vt[i][j] = QR[j][i];
        }
      } else {
        Vt[i][j] = 0;
      }
    }
  }
}

#endif