#include "analysis/eigen_reduction.h"

#include <gtest/gtest.h>

#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <utility>
#include <vector>

#include "fesa/fem/dof_manager.h"
#include "fesa/solvers/eigen/mkl_feast_eigen_solver.h"

namespace {

fesa::SparseMatrix Sparse(const fesa::Matrix& dense) {
  fesa::SparsePattern pattern;
  std::vector<fesa::CooContribution> entries;
  pattern.row_offsets.push_back(0U);
  for (std::size_t i = 0; i < dense.Rows(); ++i) {
    for (std::size_t j = 0; j < dense.Columns(); ++j) {
      pattern.column_indices.push_back(j);
      entries.push_back({i, j, dense(i, j), 0U, i * dense.Columns() + j});
    }
    pattern.row_offsets.push_back(pattern.column_indices.size());
  }
  auto result = fesa::SparseMatrix::FromCoo(dense.Rows(), dense.Columns(),
                                            std::move(entries), pattern);
  EXPECT_TRUE(result.HasValue());
  return std::move(result.Value());
}

fesa::Matrix NodalMass(const std::array<double, 3>& d) {
  fesa::Matrix mass(6U, 6U);
  for (std::size_t i = 0; i < 3U; ++i) {
    mass(i, i) = 2.0;
    for (std::size_t j = 0; j < 3U; ++j) {
      mass(i + 3U, j + 3U) = (static_cast<double>(i == j) - d[i] * d[j]) / 6.0;
    }
  }
  return mass;
}

fesa::ShellNodeInitialFrame TiltedFrame() {
  return {0U, {0.6, 0.0, 0.8}, {0.0, 1.0, 0.0}, {-0.8, 0.0, 0.6}};
}

}  // namespace

TEST(EigenReduction, KeepsPositivePartialGlobalRotationAndOriginalPencil) {
  const auto frame = TiltedFrame();
  fesa::Matrix stiffness(6U, 6U);
  stiffness(5U, 5U) = 10369.0 / 7500.0;
  std::vector<std::uint8_t> mask(6U, 1U);
  mask[5U] = 0U;
  const auto result = fesa::eigen_internal::BuildEigenReduction(
      Sparse(stiffness), {NodalMass(frame.director)}, {frame}, mask,
      {19.0 / 4800.0}, 4.0);
  ASSERT_TRUE(result.HasValue());
  EXPECT_EQ(result.Value().null_count, 0U);
  EXPECT_EQ(result.Value().whitened_operator.Rows(), 1U);
  EXPECT_NEAR(result.Value().whitened_operator.Values()[0], 10369.0 / 450.0,
              1e-12);
  fesa::Matrix x(1U, 1U, 1.0);
  const auto y = fesa::eigen_internal::RestoreScaledMode(result.Value(), x, 0U);
  for (std::size_t i = 0; i < 5U; ++i) EXPECT_DOUBLE_EQ(y[i], 0.0);
  EXPECT_NEAR(y[5U] * y[5U] * (3.0 / 50.0), 1.0, 1e-12);
}

TEST(EigenReduction, RemovesOnlyActualFreeDirectorNull) {
  const fesa::ShellNodeInitialFrame frame{0U, {0, 0, 1}, {1, 0, 0}, {0, 1, 0}};
  auto mass = NodalMass(frame.director);
  fesa::Matrix stiffness(6U, 6U);
  for (std::size_t i = 0; i < 5U; ++i) stiffness(i, i) = mass(i, i) * 7.0;
  stiffness(5U, 5U) = 0.25;
  const auto result = fesa::eigen_internal::BuildEigenReduction(
      Sparse(stiffness), {mass}, {frame}, std::vector<std::uint8_t>(6U, 0U),
      {0.25}, 2.0);
  ASSERT_TRUE(result.HasValue());
  EXPECT_EQ(result.Value().free_count, 6U);
  EXPECT_EQ(result.Value().null_count, 1U);
  EXPECT_EQ(result.Value().whitened_operator.Rows(), 5U);
  fesa::Vector ones(5U, 1.0);
  const auto action = result.Value().whitened_operator.Multiply(ones);
  for (std::size_t i = 0; i < 5U; ++i) EXPECT_NEAR(action[i], 7.0, 1e-12);
  EXPECT_LE(result.Value().eta_kz, 1e-12);
  stiffness(0U, 5U) = stiffness(5U, 0U) = 0.1;
  EXPECT_FALSE(fesa::eigen_internal::BuildEigenReduction(
                   Sparse(stiffness), {mass}, {frame},
                   std::vector<std::uint8_t>(6U, 0U), {0.25}, 2.0)
                   .HasValue());
}

TEST(EigenReduction, RejectsNonfiniteUsedBasisAndPreservesPriorCandidate) {
  const fesa::ShellNodeInitialFrame frame{0U, {0, 0, 1}, {1, 0, 0}, {0, 1, 0}};
  const auto mass = NodalMass(frame.director);
  fesa::Matrix stiffness(6U, 6U);
  for (std::size_t i = 0; i < 5U; ++i) stiffness(i, i) = mass(i, i) * 7.0;
  stiffness(5U, 5U) = .25;
  const auto matrix = Sparse(stiffness);
  const std::vector<std::uint8_t> mask(6U, 0U);
  const auto prior = fesa::eigen_internal::BuildEigenReduction(
      matrix, {mass}, {frame}, mask, {.25}, 2.0);
  ASSERT_TRUE(prior.HasValue());
  const auto prior_values = prior.Value().whitened_operator.Values();
  for (const double invalid : {std::numeric_limits<double>::quiet_NaN(),
                                std::numeric_limits<double>::infinity()}) {
    for (const bool first_tangent : {false, true}) {
      auto bad_frame = frame;
      (first_tangent ? bad_frame.tangent_a : bad_frame.tangent_b)[0] = invalid;
      const auto rejected = fesa::eigen_internal::BuildEigenReduction(
          matrix, {mass}, {bad_frame}, mask, {.25}, 2.0);
      ASSERT_FALSE(rejected.HasValue());
      EXPECT_EQ(rejected.GetStatus().Category(), fesa::FailureCategory::kSolver);
      EXPECT_EQ(rejected.GetStatus().Diagnostics()[0].code,
                "invalid-eigen-mass-space");
      EXPECT_EQ(prior.Value().whitened_operator.Values(), prior_values);
    }
  }
  const auto rerun = fesa::eigen_internal::BuildEigenReduction(
      matrix, {mass}, {frame}, mask, {.25}, 2.0);
  ASSERT_TRUE(rerun.HasValue());
  EXPECT_EQ(rerun.Value().whitened_operator.Values(), prior_values);
}

TEST(EigenReduction, RejectsPositiveButUnresolvableMassWithoutDeletingDofs) {
  const double p = 1e-8;
  const std::array<double, 3> d{2 * p / (1 + p * p), 0,
                                (1 - p * p) / (1 + p * p)};
  const fesa::ShellNodeInitialFrame frame{0U, d, {0, 1, 0}, {-d[2], 0, d[0]}};
  fesa::Matrix stiffness(6U, 6U);
  for (std::size_t i = 0; i < 6U; ++i) stiffness(i, i) = 1.0;
  std::vector<std::uint8_t> mask(6U, 0U);
  mask[3U] = 1U;
  const auto result = fesa::eigen_internal::BuildEigenReduction(
      Sparse(stiffness), {NodalMass(d)}, {frame}, mask, {0.25}, 1.0);
  ASSERT_FALSE(result.HasValue());
  EXPECT_EQ(result.GetStatus().Diagnostics()[0].code,
            "eigen-mass-conditioning");
}

TEST(EigenReduction, WhitensOffDiagonalMassAndRestoresOriginalPencil) {
  const fesa::ShellNodeInitialFrame frame{0U, {0, 0, 1}, {1, 0, 0}, {0, 1, 0}};
  auto mass = NodalMass(frame.director);
  mass(0, 1) = mass(1, 0) = 1.0;
  fesa::Matrix stiffness(6U, 6U);
  stiffness(0, 0) = stiffness(1, 1) = 6.0;
  stiffness(0, 1) = stiffness(1, 0) = 4.0;
  std::vector<std::uint8_t> mask{0U, 0U, 1U, 1U, 1U, 1U};
  const auto reduction = fesa::eigen_internal::BuildEigenReduction(
      Sparse(stiffness), {mass}, {frame}, mask, {0.25}, 3.0);
  ASSERT_TRUE(reduction.HasValue());
  const auto eigen = fesa::MklFeastEigenSolver{}.Solve(
      reduction.Value().whitened_operator, 2U, 1.0);
  ASSERT_TRUE(eigen.HasValue());
  ASSERT_EQ(eigen.Value().eigenvalues.size(), 2U);
  EXPECT_NEAR(eigen.Value().eigenvalues[0], 2.0, 1e-12);
  EXPECT_NEAR(eigen.Value().eigenvalues[1], 10.0 / 3.0, 1e-12);
  for (std::size_t mode = 0; mode < 2U; ++mode) {
    auto y = fesa::eigen_internal::RestoreScaledMode(
        reduction.Value(), eigen.Value().eigenvectors, mode);
    fesa::Vector q = y;
    for (std::size_t i = 0; i < 3U; ++i) q[i] *= 3.0;
    const auto mq = mass.Multiply(q);
    EXPECT_NEAR(q.Dot(mq), 1.0, 1e-12);
    auto residual = stiffness.Multiply(q);
    residual.Axpy(-eigen.Value().eigenvalues[mode], mq);
    EXPECT_LE(residual.Norm(), 1e-12);
  }
}

TEST(EigenReduction, PreservesSmallResolvablePositivePartialRotationMass) {
  const double p = 1.0 / 2000.0;
  const std::array<double, 3> d{2 * p / (1 + p * p), 0,
                                (1 - p * p) / (1 + p * p)};
  const fesa::ShellNodeInitialFrame frame{0U, d, {0, 1, 0}, {-d[2], 0, d[0]}};
  fesa::Matrix stiffness(6U, 6U);
  for (std::size_t i = 0; i < 6U; ++i) stiffness(i, i) = 1.0;
  std::vector<std::uint8_t> mask(6U, 0U);
  mask[3U] = 1U;
  const auto result = fesa::eigen_internal::BuildEigenReduction(
      Sparse(stiffness), {NodalMass(d)}, {frame}, mask, {0.25}, 1.0);
  ASSERT_TRUE(result.HasValue());
  EXPECT_EQ(result.Value().null_count, 0U);
  EXPECT_EQ(result.Value().whitened_operator.Rows(), 5U);
  EXPECT_GE(result.Value().nodes[0].record.rcond_inf, 0x1p-32);
}

TEST(EigenReduction, HandlesTwoFreeRotationsAndRejectsEmptyFiniteSpace) {
  const fesa::ShellNodeInitialFrame frame{0U, {0, 0, 1}, {1, 0, 0}, {0, 1, 0}};
  auto mass = NodalMass(frame.director);
  fesa::Matrix stiffness(6U, 6U);
  stiffness(3U, 3U) = 1.0;
  stiffness(5U, 5U) = 0.25;
  std::vector<std::uint8_t> mask{1U, 1U, 1U, 0U, 1U, 0U};
  auto result = fesa::eigen_internal::BuildEigenReduction(
      Sparse(stiffness), {mass}, {frame}, mask, {0.25}, 1.0);
  ASSERT_TRUE(result.HasValue());
  EXPECT_EQ(result.Value().null_count, 1U);
  EXPECT_EQ(result.Value().whitened_operator.Rows(), 1U);
  mass(5U, 5U) = 0.1;
  EXPECT_FALSE(fesa::eigen_internal::BuildEigenReduction(
                   Sparse(stiffness), {mass}, {frame}, mask, {0.25}, 1.0)
                   .HasValue());
  mass(5U, 5U) = 0.0;
  for (const bool drill_free : {false, true}) {
    mask.assign(6U, 1U);
    mask[5U] = drill_free ? 0U : 1U;
    const auto empty = fesa::eigen_internal::BuildEigenReduction(
        Sparse(stiffness), {mass}, {frame}, mask, {0.25}, 1.0);
    ASSERT_FALSE(empty.HasValue());
    EXPECT_EQ(empty.GetStatus().Category(), fesa::FailureCategory::kModel);
    EXPECT_EQ(empty.GetStatus().Diagnostics()[0].code, "empty-eigenproblem");
  }
}