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

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

namespace QR {

// ============================================================================
// QR Building Block Implementations (fully templated, heap-free)
// ============================================================================

/**
 * GivensRotation: R * (a, b)^T = (r, 0)^T with R = [[c, s], [-s, c]],
 * r = +hypot(a, b), c = a/r, s = b/r.
 */
// [[maybe_unused]]: this helper is only referenced from template
// (EigenQR/Tridiagonalize), so in translation units that include this file
// but never instantiate those templates, the definition is legitimately
// unused. The attribute silences -Wunused-function there without hiding
// real dead code in TUs that do use the algorithm.
[[maybe_unused]] static void GivensRotation(float a, float b, float &c,
                                            float &s) {
  float r = sqrtf(a * a + b * b);
  if (r == 0.0f) {
    c = 1.0f;
    s = 0.0f;
    return;
  }
  c = a / r;
  s = b / r;
}

/**
 * ApplyRotationBothSides: A <- G A G^T (similarity transform) with
 * G = [[c, s], [-s, c]] on the (i, i+1) block, i.e. G is the ZEROING
 * rotation G*(x, y)^T = (r, 0)^T (the orientation used by the implicit QR
 * chase: A = Q R with Q = G^T gives the next iterate R Q = G A G^T).
 * With (c, s) = GivensRotation(A[i][i], A[i+1][i]) this zeroes
 * A[i+1][i] after the LEFT multiplication; the right multiplication then
 * chases the bulge along the superdiagonal (tridiagonal chase).
 *
 * A must be symmetric on entry; the result stays symmetric, so both
 * triangles are written.
 *
 * Block updates (with a00 = A[i][i], a01 = A[i][i+1], a11 = A[i+1][i+1]):
 *   A[i][i]     = c^2 a00 + 2 c s a01 + s^2 a11
 *   A[i][i+1]   = (c^2 - s^2) a01 + c s (a11 - a00)
 *   A[i+1][i+1] = s^2 a00 - 2 c s a01 + c^2 a11
 * Off-block updates (uniform for both sides, since the left factor G and
 * the right factor G^T mix each side with the pattern (a, b) -> (c a + s b,
 * -s a + c b) after transposition):
 *   for j not in {i, i+1}:
 *     A[i][j]   = A[j][i]   = c A[i][j]   + s A[i+1][j]
 *     A[i+1][j] = A[j][i+1] = -s A[i][j]  + c A[i+1][j]
 */
template <uint8_t N>
static void ApplyRotationBothSides(Matrix<N, N> &A, uint8_t i, float c,
                                   float s) {
  float a00 = A.Get(i, i);
  float a01 = A.Get(i, i + 1);
  float a11 = A.Get(i + 1, i + 1);
  float c2 = c * c;
  float s2 = s * s;
  float cs = c * s;

  A[i][i] = c2 * a00 + 2.0f * cs * a01 + s2 * a11;
  A[i][i + 1] = (c2 - s2) * a01 + cs * (a11 - a00);
  A[i + 1][i + 1] = s2 * a00 - 2.0f * cs * a01 + c2 * a11;
  A[i + 1][i] = A[i][i + 1]; // keep both triangles in sync

  for (uint8_t j = 0; j < N; ++j) {
    if (j == i || j == i + 1)
      continue;
    float x = A.Get(i, j);
    float y = A.Get(i + 1, j);
    A[i][j] = c * x + s * y;
    A[j][i] = A[i][j];
    A[i + 1][j] = -s * x + c * y;
    A[j][i + 1] = A[i + 1][j];
  }
}

/**
 * ApplyRotationToVectors: V <- V G^T with G = [[c, s], [-s, c]] on columns
 * (i, i+1), applied to every row. G^T = [[c, -s], [s, c]], so
 *   V[r][i]     <- c V[r][i]   + s V[r][i+1]
 *   V[r][i+1]   <- -s V[r][i]  + c V[r][i+1]
 *
 * Convention pairing: if A evolves as A <- G A G^T (ApplyRotationBothSides
 * with the SAME c, s), then V accumulates V <- V G^T. With V0 = I the
 * invariant A0 = V A V^T is preserved at every step, so at convergence
 * A0 = V D V^T and the columns of V are the eigenvectors. (Rationale:
 * each chase step is A <- R Q with R = G A the upper-triangular factor and
 * Q = G^T the orthogonal factor of A = Q R, so A = G^T A' G and the
 * orthogonal factors multiply as G1^T G2^T ... in application order.)
 */
template <uint8_t N>
static void ApplyRotationToVectors(Matrix<N, N> &V, uint8_t i, float c,
                                   float s) {
  for (uint8_t r = 0; r < N; ++r) {
    float x = V.Get(r, i);
    float y = V.Get(r, i + 1);
    V[r][i] = c * x + s * y;
    V[r][i + 1] = -s * x + c * y;
  }
}

/**
 * WilkinsonShift: eigenvalue of [[a, b], [b, d]] closest to d.
 * mu = (a+d)/2 - sign(a-d) * sqrt(((a-d)/2)^2 + b^2), sign(0) = +1.
 */
[[maybe_unused]] static float WilkinsonShift(float a, float b, float d) {
  float delta = 0.5f * (a - d);
  float spread = sqrtf(delta * delta + b * b);
  return 0.5f * (a + d) - (delta >= 0.0f ? spread : -spread);
}

/**
 * Solve2x2Eigen: closed-form eigen-decomposition of the 2x2 block at
 * (lo, lo+1). Works for symmetric blocks and for general 2x2 blocks with
 * real eigenvalues (used by the N == 2 entry point).
 *
 * lambdaHi/lambdaLo come from the characteristic polynomial
 *   lambda^2 - trace*lambda + det = 0.
 * The eigenvector for lambdaHi is v = (b, lambdaHi - a) (from the first
 * row of (A - lambda*I)v = 0), normalized to unit length. If b == 0 the
 * block is triangular and the eigenvectors are coordinate vectors:
 * e1 for the larger of {a, d}, e2 for the other.
 */
template <uint8_t N>
static void Solve2x2Eigen(const Matrix<N, N> &A, uint8_t lo, float &lambdaHi,
                          float &lambdaLo, float &c, float &s) {
  float a = A.Get(lo, lo);
  float b = A.Get(lo, lo + 1);
  float e = A.Get(lo + 1, lo);
  float d = A.Get(lo + 1, lo + 1);

  float trace = a + d;
  float det = a * d - b * e;
  float disc = trace * trace - 4.0f * det;
  if (disc < 0.0f)
    disc = 0.0f; // round-off clamp: real 2x2 blocks have disc >= 0
  float sqrtDisc = sqrtf(disc);
  lambdaHi = 0.5f * (trace + sqrtDisc);
  lambdaLo = 0.5f * (trace - sqrtDisc);

  if (b != 0.0f) {
    float v1 = lambdaHi - a;
    float n = sqrtf(b * b + v1 * v1);
    c = b / n;
    s = v1 / n;
  } else if (a >= d) {
    c = 1.0f; // e1 is the eigenvector of a = lambdaHi
    s = 0.0f;
  } else {
    c = 0.0f; // e2 is the eigenvector of d = lambdaHi
    s = 1.0f;
  }
}

/**
 * Deflate: zero subdiagonal entries i in [lo, hi) whose magnitude is at or
 * below tolerance * (|A[i][i]| + |A[i+1][i+1]|).
 */
template <uint8_t N>
static void Deflate(Matrix<N, N> &A, uint8_t lo, uint8_t hi, float tolerance) {
  for (uint8_t i = lo; i < hi; ++i) {
    float t = A.Get(i + 1, i);
    float scale = fabsf(A.Get(i, i)) + fabsf(A.Get(i + 1, i + 1));
    if (fabsf(t) <= tolerance * scale) {
      A[i + 1][i] = 0.0f;
      A[i][i + 1] = 0.0f;
    }
  }
}

// ============================================================================
// QR::EigenQR driver (implicit Wilkinson-shifted QR, bulge chasing)
// ============================================================================

/**
 * Tridiagonalize: Givens tridiagonalization (Golub & Van Loan 8.3.1).
 *
 * For column k = 0..N-3 the entries A[k+2..N-1, k] are eliminated by
 * rotations on (i, i+1) applied BOTTOM-UP, i = N-2 down to k+1, each
 * formed from the CURRENT (already-updated) pair (A[i][k], A[i+1][k]).
 * Bottom-up is essential: a top-down pass zeros A[i+1][k] with a rotation
 * that would later be undone when the next rotation (i+1, i+2) is formed
 * from an entry below, reviving A[i][k]. Each bottom-up rotation zeros the
 * bottom of the remaining nonzero pair and the entries below stay zero
 * (they are not mixed again, only rows i-1/i are mixed next).
 *
 * Already-tridiagonalized leading columns j < k are untouched: the mixed
 * rows are both >= k+1 > j+1, so A[i][j] and A[i+1][j] are both zero there.
 * The rotation on (i, i+1) also keeps column k+1..k+2 structure intact and
 * does not destroy earlier columns, so after column k is done the leading
 * (k+1)x(k+1) block is tridiagonal forever.
 *
 * On return: A is symmetric tridiagonal and A_orig = U A U^T (U = product
 * of every rotation applied, in application order, as U <- U G^T).
 */
template <uint8_t N>
static void Tridiagonalize(Matrix<N, N> &A, Matrix<N, N> &U) {
  U = Matrix<N, N>{0};
  for (uint8_t i = 0; i < N; ++i) {
    U[i][i] = 1.0f;
  }
  float c = 0.0f, s = 0.0f;
  for (uint8_t k = 0; k + 2 < N; ++k) {
    for (int i = (int)N - 2; i >= (int)k + 1; --i) {
      GivensRotation(A.Get(i, k), A.Get(i + 1, k), c, s);
      ApplyRotationBothSides(A, (uint8_t)i, c, s);
      ApplyRotationToVectors(U, (uint8_t)i, c, s);
    }
  }
}

/**
 * See QR.hpp for the full contract. Implementation sketch:
 *
 *   Phase 0 (N >= 3): Tridiagonalize(A, U)  // A_orig = U A U^T
 *   V = I.
 *   while (hi > 0):
 *     Deflate(A, 0, hi, tol); peel exact-zero trailing subdiagonals (hi--)
 *     lo = top of the trailing unreduced block (scan down, stop at first
 *          exact zero subdiagonal)
 *     if lo == hi - 1:  closed-form 2x2 eigen-solve; fold Vblock into V
 *     else:             one implicit Wilkinson-shifted QR step:
 *                        mu = WilkinsonShift(A[hi-1][hi-1], A[hi][hi-1], A[hi][hi])
 *                        A[lo..hi diagonal] -= mu          // whole block!
 *                        G1 = Givens(A[lo][lo], A[lo+1][lo])
 *                        for i = lo..hi-1:
 *                          (i > lo: Gi = Givens(A[i][i], A[i+1][i]))
 *                          ApplyRotationBothSides(A, i, Gi)   // A <- Gi A Gi^T
 *                          ApplyRotationToVectors(V, i, Gi)   // V <- V Gi^T
 *                        A[lo..hi diagonal] += mu
 *   eigenvalues = diag(A), sorted descending with matching V column swaps.
 *   eigenvectors = U * V.
 *
 * Invariant maintained for N >= 3 (symmetric input): A is symmetric
 * tridiagonal (up to deflated zeros and ~1e-7 float roundoff in the
 * off-tridiagonal corners) at the top of every loop iteration, and
 * A_orig = U A U^T = (U V) A (U V)^T throughout (V = product of every
 * rotation applied so far, in application order, as V <- V Gi^T). At
 * convergence A = V D V^T and therefore A_orig = (U V) D (U V)^T.
 *
 * Orientation note: each chase rotation Gi is the ZEROING rotation
 * (Gi * (x, y)^T = (r, 0)^T). The step A <- Gi A Gi^T equals R Q with
 * R = Gi A upper-triangular (on the block) and Q = Gi^T -- i.e. it IS the
 * standard QR update Q(A - mu I)Q^T with Q the orthogonal QR factor. The
 * eigenvector accumulator therefore collects the Q factors: V <- V Gi^T.
 */
template <uint8_t N>
void EigenQR(Matrix<N, N> &matrixToDecompose, Matrix<N, N> &eigenVectors,
             Matrix<N, 1> &eigenValues, uint32_t maxIterations, float tolerance) {
  static_assert(N >= 2, "QR::EigenQR requires N >= 2 (N = 1 is trivial)");

  Matrix<N, N> A = matrixToDecompose; // input is not modified
  Matrix<N, N> V{0};
  // NB: Matrix::Identity() is a static factory that returns by value; a
  // bare call would be a no-op. Set the diagonal explicitly.
  for (uint8_t i = 0; i < N; ++i) {
    V[i][i] = 1.0f;
  }

  // ------------------------------------------------------------------
  // N == 2: closed-form solution (works for nonsymmetric input too)
  // ------------------------------------------------------------------
  if (N == 2) {
    float l1 = 0.0f, l2 = 0.0f, c = 0.0f, s = 0.0f;
    Solve2x2Eigen(A, 0, l1, l2, c, s);
    // V = I * Vblock = [[c, -s], [s, c]]
    V[0][0] = c;
    V[0][1] = -s;
    V[1][0] = s;
    V[1][1] = c;
    eigenValues[0][0] = l1;
    eigenValues[1][0] = l2;
    for (uint8_t r = 0; r < N; ++r)
      for (uint8_t col = 0; col < N; ++col)
        eigenVectors[r][col] = V.Get(r, col);
    return;
  }

  // ------------------------------------------------------------------
  // N >= 3: implicit shifted QR iteration (symmetric input required)
  // ------------------------------------------------------------------
  // Phase 0: general symmetric -> symmetric tridiagonal. The implicit
  // QR bulge chase only preserves a tridiagonal structure, so the input
  // must be reduced first: A_orig = U A U^T with A tridiagonal.
  Matrix<N, N> U{};
  Tridiagonalize(A, U);

  uint32_t iter = 0;
  uint8_t hi = N - 1;
  while (hi > 0) {
    Deflate(A, 0, hi, tolerance);
    // Peel trailing rows whose subdiagonal is exactly zero (deflated or
    // already solved). Must be re-done every iteration: a peel is only
    // meaningful once the subdiagonal beneath it has converged.
    while (hi > 0 && A.Get(hi, hi - 1) == 0.0f) {
      --hi;
    }
    if (hi == 0) {
      break; // fully diagonal (within tolerance)
    }

    // Find the top of the trailing unreduced block: scan down from hi-1
    // and stop at the first exact zero subdiagonal. A[hi][hi-1] != 0 here
    // (just peeled), so lo < hi.
    uint8_t lo = hi;
    for (int i = (int)hi - 1; i >= 0; --i) {
      if (A.Get(i + 1, i) == 0.0f) {
        break;
      }
      lo = (uint8_t)i;
    }

    if (lo + 1 == hi) {
      // Trailing unreduced block is 2x2: solve in closed form.
      float l1 = 0.0f, l2 = 0.0f, c = 0.0f, s = 0.0f;
      Solve2x2Eigen(A, lo, l1, l2, c, s);
      A[lo][lo] = l1;
      A[lo + 1][lo + 1] = l2;
      A[lo][lo + 1] = 0.0f;
      A[lo + 1][lo] = 0.0f;
      // Fold Vblock = [[c, -s], [s, c]] into V: V <- V * Vblock on
      // columns (lo, lo+1). NOTE the sign convention differs from
      // ApplyRotationToVectors (which applies [[c, s], [-s, c]]):
      // here column 0 of Vblock is (c, s)^T, column 1 is (-s, c)^T.
      for (uint8_t r = 0; r < N; ++r) {
        float x = V.Get(r, lo);
        float y = V.Get(r, lo + 1);
        V[r][lo] = c * x + s * y;
        V[r][lo + 1] = -s * x + c * y;
      }
      if (lo == 0) {
        break; // block reached the top: matrix is fully solved
      }
      hi = (uint8_t)(lo - 1);
      continue;
    }

    // One implicit Wilkinson-shifted QR step on block [lo, hi].
    float mu = WilkinsonShift(A.Get(hi - 1, hi - 1), A.Get(hi, hi - 1),
                              A.Get(hi, hi));

    // The shift applies to the ENTIRE active block: bulge chasing
    // triangularizes (A - mu*I), and the first Givens rotation is formed
    // from (A[lo][lo] - mu, A[lo+1][lo]).
    for (uint8_t i = lo; i <= hi; ++i) {
      A[i][i] -= mu;
    }

    float c = 0.0f, s = 0.0f;
    for (uint8_t i = lo; i < hi; ++i) {
      if (i == lo) {
        GivensRotation(A.Get(lo, lo), A.Get(lo + 1, lo), c, s);
      } else {
        GivensRotation(A.Get(i, i), A.Get(i + 1, i), c, s);
      }
      ApplyRotationBothSides(A, i, c, s);
      ApplyRotationToVectors(V, i, c, s);
    }

    for (uint8_t i = lo; i <= hi; ++i) {
      A[i][i] += mu;
    }

    if (++iter >= maxIterations) {
      // Best-effort: fall through with the partially diagonalized A.
      break;
    }
  }

  // ------------------------------------------------------------------
  // Collect eigenvalues and sort DESCENDING (swap eigenvectors to match)
  // ------------------------------------------------------------------
  for (uint8_t i = 0; i < N; ++i) {
    eigenValues[i][0] = A.Get(i, i);
  }
  for (uint8_t i = 0; i < N - 1; ++i) {
    uint8_t k = i;
    for (uint8_t j = i + 1; j < N; ++j) {
      if (eigenValues.Get(j, 0) > eigenValues.Get(k, 0)) {
        k = j;
      }
    }
    if (k != i) {
      float t = eigenValues[i][0];
      eigenValues[i][0] = eigenValues[k][0];
      eigenValues[k][0] = t;
      for (uint8_t r = 0; r < N; ++r) {
        float x = V.Get(r, i);
        V[r][i] = V.Get(r, k);
        V[r][k] = x;
      }
    }
  }

  // True eigenvectors of the original matrix: U * V. Reuse the A buffer
  // (its diagonal has already been collected into eigenValues).
  U.Mult(V, A);

  for (uint8_t r = 0; r < N; ++r) {
    for (uint8_t col = 0; col < N; ++col) {
      eigenVectors[r][col] = A.Get(r, col);
    }
  }
}

} // namespace QR

#endif // QR_H_