#include "Matrix.hpp"
#include "SVD.hpp"
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include <iostream>

// Generic helper functions for any matrix size
template <uint8_t rows, uint8_t columns>
static float frobeniusNorm(const Matrix<rows, columns> &M) {
  float sum = 0.0f;
  for (int i = 0; i < rows; i++)
    for (int j = 0; j < columns; j++) {
      float v = M.Get(i, j);
      sum += v * v;
    }
  return sqrtf(sum);
}

template <uint8_t n>
static bool isOrthogonal(const Matrix<n, n> &M, float tol = 1e-4f) {
  Matrix<n, n> Mt = M.Transpose();
  Matrix<n, n> MtM{0};
  Mt.Mult(M, MtM);
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++) {
      float expected = (i == j) ? 1.0f : 0.0f;
      if (fabsf(MtM.Get(i, j) - expected) > tol)
        return false;
    }
  return true;
}

TEST_CASE("SVD Integration: 2x2 [[1,2],[3,4]]", "[Matrix][SVD][Integration]") {
  Matrix<2, 2> A{1, 2, 3, 4};
  Matrix<2, 2> U{0};
  Matrix<2, 1> sigma{0};
  Matrix<2, 2> Vt{0};

  SVD::SVD(A, U, sigma, Vt);

  // Reference singular values from scipy: [5.464985704219, 0.365966190626]
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(5.4649857f, 1e-3f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(0.3659662f, 1e-3f));

  // Check orthogonality of U and Vt (first 2x2 blocks)
  REQUIRE(isOrthogonal<2>(U));
  REQUIRE(isOrthogonal<2>(Vt));

  // Check reconstruction: A ≈ U · diag(sigma) · Vt
  Matrix<2, 2> recon{0};
  Matrix<2, 2> Usig{0};
  for (int i = 0; i < 2; i++)
    for (int j = 0; j < 2; j++)
      Usig[i][j] = U.Get(i, j) * sigma.Get(j, 0);

  Usig.Mult(Vt, recon);

  float err = 0.0f;
  for (int i = 0; i < 2; i++)
    for (int j = 0; j < 2; j++) {
      float diff = recon.Get(i, j) - A.Get(i, j);
      err += diff * diff;
    }
  err = sqrtf(err);
  REQUIRE_THAT(err, Catch::Matchers::WithinAbs(0.0f, 1e-3f));

  std::cout << "SVD 2x2 [[1,2],[3,4]]:\n";
  std::cout << "Sigma: [" << sigma.Get(0, 0) << ", " << sigma.Get(1, 0)
            << "]\n";
}

TEST_CASE("SVD Integration: 3x3 diagonal [10,5,2]",
          "[Matrix][SVD][Integration]") {
  Matrix<3, 3> A{10, 0, 0, 0, 5, 0, 0, 0, 2};
  Matrix<3, 3> U{0};
  Matrix<3, 1> sigma{0};
  Matrix<3, 3> Vt{0};

  SVD::SVD(A, U, sigma, Vt);

  // Singular values should be [10, 5, 2] (already diagonal)
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(10.0f, 1e-3f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(5.0f, 1e-3f));
  REQUIRE_THAT(sigma.Get(2, 0), Catch::Matchers::WithinRel(2.0f, 1e-3f));

  // U and Vt should be identity (or close) for diagonal matrix
  float uErr = frobeniusNorm(U - Matrix<3, 3>{1, 0, 0, 0, 1, 0, 0, 0, 1});
  float vtErr = frobeniusNorm(Vt - Matrix<3, 3>{1, 0, 0, 0, 1, 0, 0, 0, 1});
  REQUIRE_THAT(uErr, Catch::Matchers::WithinAbs(0.0f, 1e-2f));
  REQUIRE_THAT(vtErr, Catch::Matchers::WithinAbs(0.0f, 1e-2f));
}

TEST_CASE("SVD Integration: 3x3 rank-deficient [[1,2,3],[4,5,6],[7,8,9]]",
          "[Matrix][SVD][Integration]") {
  Matrix<3, 3> A{1, 2, 3, 4, 5, 6, 7, 8, 9};
  Matrix<3, 3> U{0};
  Matrix<3, 1> sigma{0};
  Matrix<3, 3> Vt{0};

  SVD::SVD(A, U, sigma, Vt);

  // Reference: [16.848103352614, 1.068369514555, 0.0]
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(16.8481f, 1e-2f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(1.06837f, 1e-2f));
  // Third singular value should be ~0 (rank-deficient)
  REQUIRE_THAT(sigma.Get(2, 0), Catch::Matchers::WithinAbs(0.0f, 1e-2f));

  // Check reconstruction
  Matrix<3, 3> recon{0};
  Matrix<3, 3> Usig{0};
  for (int i = 0; i < 3; i++)
    for (int j = 0; j < 3; j++)
      Usig[i][j] = U.Get(i, j) * sigma.Get(j, 0);
  Usig.Mult(Vt, recon);

  float err = 0.0f;
  for (int i = 0; i < 3; i++)
    for (int j = 0; j < 3; j++) {
      float diff = recon.Get(i, j) - A.Get(i, j);
      err += diff * diff;
    }
  err = sqrtf(err);
  REQUIRE_THAT(err, Catch::Matchers::WithinAbs(0.0f, 1e-2f));

  std::cout << "SVD 3x3 rank-deficient:\n";
  std::cout << "Sigma: [" << sigma.Get(0, 0) << ", " << sigma.Get(1, 0) << ", "
            << sigma.Get(2, 0) << "]\n";
}

TEST_CASE("SVD Integration: tall 4x3 matrix", "[Matrix][SVD][Integration]") {
  Matrix<4, 3> A{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  Matrix<4, 3> U{0};
  Matrix<3, 1> sigma{0};
  Matrix<3, 3> Vt{0};

  SVD::SVD(A, U, sigma, Vt);

  // Reference: [25.462407436036, 1.290661675761, 0.0]
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(25.4624f, 1e-2f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(1.29066f, 1e-2f));
  REQUIRE_THAT(sigma.Get(2, 0), Catch::Matchers::WithinAbs(0.0f, 1e-2f));

  // Check reconstruction
  Matrix<4, 3> recon{0};
  Matrix<4, 3> Usig{0};
  for (int i = 0; i < 4; i++)
    for (int j = 0; j < 3; j++)
      Usig[i][j] = U.Get(i, j) * sigma.Get(j, 0);
  Usig.Mult(Vt, recon);

  float err = 0.0f;
  for (int i = 0; i < 4; i++)
    for (int j = 0; j < 3; j++) {
      float diff = recon.Get(i, j) - A.Get(i, j);
      err += diff * diff;
    }
  err = sqrtf(err);
  REQUIRE_THAT(err, Catch::Matchers::WithinAbs(0.0f, 1e-2f));

  std::cout << "SVD tall 4x3:\n";
  std::cout << "Sigma: [" << sigma.Get(0, 0) << ", " << sigma.Get(1, 0) << ", "
            << sigma.Get(2, 0) << "]\n";
}

TEST_CASE("SVD Integration: wide 3x5 matrix", "[Matrix][SVD][Integration]") {
  Matrix<3, 5> A{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
  Matrix<3, 5> U{0};
  Matrix<5, 1> sigma{0}; // sigma is columns x 1 = 5x1 for wide matrix
  Matrix<5, 5> Vt{0};    // Vt is columns x columns = 5x5

  SVD::SVD(A, U, sigma, Vt);

  // Reference: [35.127223333575, 2.465396696917, 0.0]
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(35.1272f, 1e-2f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(2.46540f, 1e-2f));
  REQUIRE_THAT(sigma.Get(2, 0), Catch::Matchers::WithinAbs(0.0f, 1e-2f));

  // Check reconstruction: A (3x5) = U * Sigma * Vt, where U (3x5) has
  // its meaningful part in the first 3 columns, sigma (5x1) in the
  // first 3 entries, and Vt (5x5) in its first 3 rows (right
  // singular vectors as rows). So:
  //   A[i][j] = sum_k U[i][k] * sigma[k] * Vt[k][j]

  float err2 = 0.0f;
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 5; j++) {
      float recon_val = 0.0f;
      for (int k = 0; k < 3; k++) {
        recon_val += U.Get(i, k) * sigma.Get(k, 0) * Vt.Get(k, j);
      }
      float diff = recon_val - A.Get(i, j);
      err2 += diff * diff;
    }
  }
  err2 = sqrtf(err2);
  REQUIRE_THAT(err2, Catch::Matchers::WithinAbs(0.0f, 1e-2f));

  std::cout << "SVD wide 3x5:\n";
  std::cout << "Sigma: [" << sigma.Get(0, 0) << ", " << sigma.Get(1, 0) << ", "
            << sigma.Get(2, 0) << "]\n";
}

TEST_CASE("SVD Integration: identity 3x3", "[Matrix][SVD][Integration]") {
  Matrix<3, 3> A{1, 0, 0, 0, 1, 0, 0, 0, 1};
  Matrix<3, 3> U{0};
  Matrix<3, 1> sigma{0};
  Matrix<3, 3> Vt{0};

  SVD::SVD(A, U, sigma, Vt);

  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(1.0f, 1e-3f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(1.0f, 1e-3f));
  REQUIRE_THAT(sigma.Get(2, 0), Catch::Matchers::WithinRel(1.0f, 1e-3f));

  float err = frobeniusNorm(U - Matrix<3, 3>{1, 0, 0, 0, 1, 0, 0, 0, 1});
  REQUIRE_THAT(err, Catch::Matchers::WithinAbs(0.0f, 1e-2f));
}

TEST_CASE("SVD Integration: symmetric positive definite 2x2 [[5,3],[3,5]]",
          "[Matrix][SVD][Integration]") {
  Matrix<2, 2> A{5, 3, 3, 5};
  Matrix<2, 2> U{0};
  Matrix<2, 1> sigma{0};
  Matrix<2, 2> Vt{0};

  SVD::SVD(A, U, sigma, Vt);

  // For SPD matrix, singular values = eigenvalues: [8, 2]
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(8.0f, 1e-3f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(2.0f, 1e-3f));

  // Check reconstruction
  Matrix<2, 2> recon{0};
  Matrix<2, 2> Usig{0};
  for (int i = 0; i < 2; i++)
    for (int j = 0; j < 2; j++)
      Usig[i][j] = U.Get(i, j) * sigma.Get(j, 0);
  Usig.Mult(Vt, recon);

  float err = 0.0f;
  for (int i = 0; i < 2; i++)
    for (int j = 0; j < 2; j++) {
      float diff = recon.Get(i, j) - A.Get(i, j);
      err += diff * diff;
    }
  err = sqrtf(err);
  REQUIRE_THAT(err, Catch::Matchers::WithinAbs(0.0f, 1e-3f));

  std::cout << "SVD SPD 2x2 [[5,3],[3,5]]:\n";
  std::cout << "Sigma: [" << sigma.Get(0, 0) << ", " << sigma.Get(1, 0)
            << "]\n";
}

// ----------------------------------------------------------------------------
// Matrix::SVD member wrapper (delegates to SVD::SVD)
// ----------------------------------------------------------------------------

/**
 * Reconstruction error ‖U·diag(sigma)·Vᵀ − A‖_F. Zero-padded entries of
 * U/sigma/Vt (wide/tall cases) are zero by the output conventions, so the
 * full product equals U[:, :k]·diag(sigma[:k])·Vt[:k, :].
 */
template <uint8_t rows, uint8_t columns>
static float svdReconstructionError(const Matrix<rows, columns> &A,
                                    const Matrix<rows, columns> &U,
                                    const Matrix<columns, 1> &sigma,
                                    const Matrix<columns, columns> &Vt) {
  Matrix<rows, columns> recon{0};
  Matrix<rows, columns> Usig{0};
  for (int i = 0; i < rows; i++)
    for (int j = 0; j < columns; j++)
      Usig[i][j] = U.Get(i, j) * sigma.Get(j, 0);
  Usig.Mult(Vt, recon);
  float err = 0.0f;
  for (int i = 0; i < rows; i++)
    for (int j = 0; j < columns; j++) {
      float diff = recon.Get(i, j) - A.Get(i, j);
      err += diff * diff;
    }
  return sqrtf(err);
}

/**
 * Orthonormality of the first k columns of M: the k×k leading block of
 * MᵀM must equal I_k. (For a tall SVD, U has k = min(rows, cols)
 * meaningful columns and this is the full UᵀU.)
 */
template <uint8_t r, uint8_t c>
static bool leadingColumnsOrthonormal(const Matrix<r, c> &M, uint8_t k,
                                      float tol = 1e-4f) {
  Matrix<c, r> Mt = M.Transpose();
  Matrix<c, c> MtM{0};
  Mt.Mult(M, MtM);
  for (int i = 0; i < k; i++)
    for (int j = 0; j < k; j++) {
      float expected = (i == j) ? 1.0f : 0.0f;
      if (fabsf(MtM.Get(i, j) - expected) > tol)
        return false;
    }
  return true;
}

/**
 * Orthonormality of the first k rows of M: the k×k leading block of
 * M·Mᵀ must equal I_k. (Vᵀ may have zero-padded trailing rows in the
 * wide case, so check only the meaningful leading block.)
 */
template <uint8_t r, uint8_t c>
static bool leadingRowsOrthonormal(const Matrix<r, c> &M, uint8_t k,
                                   float tol = 1e-4f) {
  Matrix<c, r> Mt = M.Transpose();
  Matrix<r, r> MMt{0};
  M.Mult(Mt, MMt);
  for (int i = 0; i < k; i++)
    for (int j = 0; j < k; j++) {
      float expected = (i == j) ? 1.0f : 0.0f;
      if (fabsf(MMt.Get(i, j) - expected) > tol)
        return false;
    }
  return true;
}

TEST_CASE("Matrix::SVD wrapper: 3x2 tall [[1,2],[3,4],[5,6]]",
          "[Matrix][SVD][Wrapper]") {
  Matrix<3, 2> A{1, 2, 3, 4, 5, 6};
  Matrix<3, 2> U{0};
  Matrix<2, 1> sigma{0};
  Matrix<2, 2> Vt{0};

  A.SVD(U, sigma, Vt);

  // Reference singular values from numpy: [9.52552, 0.514301]
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(9.52552f, 1e-3f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(0.514301f, 1e-3f));

  REQUIRE(leadingColumnsOrthonormal(U, 2));
  REQUIRE(leadingRowsOrthonormal(Vt, 2));

  float err = svdReconstructionError(A, U, sigma, Vt);
  REQUIRE_THAT(err, Catch::Matchers::WithinAbs(0.0f, 1e-3f));
}

TEST_CASE("Matrix::SVD wrapper: 2x3 wide [[1,2,3],[4,5,6]]",
          "[Matrix][SVD][Wrapper]") {
  Matrix<2, 3> A{1, 2, 3, 4, 5, 6};
  Matrix<2, 3> U{0};
  Matrix<3, 1> sigma{0};
  Matrix<3, 3> Vt{0};

  A.SVD(U, sigma, Vt);

  // Reference singular values from numpy: [9.50803, 0.77287]; the third
  // entry (wide-matrix padding) must be zero.
  REQUIRE_THAT(sigma.Get(0, 0), Catch::Matchers::WithinRel(9.50803f, 1e-3f));
  REQUIRE_THAT(sigma.Get(1, 0), Catch::Matchers::WithinRel(0.77287f, 1e-3f));
  REQUIRE_THAT(sigma.Get(2, 0), Catch::Matchers::WithinAbs(0.0f, 1e-6f));

  REQUIRE(leadingColumnsOrthonormal(U, 2));
  REQUIRE(leadingRowsOrthonormal(Vt, 2));

  float err = svdReconstructionError(A, U, sigma, Vt);
  REQUIRE_THAT(err, Catch::Matchers::WithinAbs(0.0f, 1e-3f));
}