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

#include <gtest/gtest.h>
#include <mkl.h>

#include <cmath>
#include <cstddef>
#include <utility>
#include <vector>

#include "fesa/fem/dof_manager.h"
#include "fesa/math/sparse_matrix.h"

namespace {

fesa::SparseMatrix MakeDenseCsr(const std::size_t size,
                                const std::vector<double>& values) {
  EXPECT_EQ(values.size(), size * size);
  fesa::SparsePattern pattern;
  std::vector<fesa::CooContribution> contributions;
  pattern.row_offsets.push_back(0U);
  for (std::size_t row = 0U; row < size; ++row) {
    for (std::size_t column = 0U; column < size; ++column) {
      pattern.column_indices.push_back(column);
      contributions.push_back(
          {row, column, values[row * size + column], row, column});
    }
    pattern.row_offsets.push_back(pattern.column_indices.size());
  }
  auto matrix = fesa::SparseMatrix::FromCoo(size, size,
                                            std::move(contributions), pattern);
  EXPECT_TRUE(matrix.HasValue());
  return std::move(matrix.Value());
}

void ExpectUnitGram(const fesa::Matrix& vectors) {
  for (std::size_t left = 0U; left < vectors.Columns(); ++left) {
    for (std::size_t right = 0U; right < vectors.Columns(); ++right) {
      double dot = 0.0;
      for (std::size_t row = 0U; row < vectors.Rows(); ++row) {
        dot += vectors(row, left) * vectors(row, right);
      }
      EXPECT_NEAR(dot, left == right ? 1.0 : 0.0, 1.0e-10);
    }
  }
}

void RequireSuccess(const fesa::Result<fesa::FeastEigenResult>& result) {
  if (result.HasValue()) return;
  const auto& diagnostics = result.GetStatus().Diagnostics();
  ASSERT_FALSE(diagnostics.empty());
  ADD_FAILURE() << diagnostics.front().code << " "
                << diagnostics.front().entity_identity << " "
                << diagnostics.front().message;
}

}  // namespace

TEST(MklFeastEigenSolver, SolvesN1AndRetainsBackendEvidence) {
  const auto matrix = MakeDenseCsr(1U, {4.0});
  fesa::MklFeastEigenSolver solver;

  const auto result = solver.Solve(matrix, 1U, 1.0);

  RequireSuccess(result);
  if (!result.HasValue()) return;
  EXPECT_EQ(result.Value().requested_modes, 1U);
  ASSERT_EQ(result.Value().eigenvalues.size(), 1U);
  EXPECT_NEAR(result.Value().eigenvalues[0U], 4.0, 1.0e-10);
  EXPECT_DOUBLE_EQ(result.Value().operator_scale, 4.0);
  EXPECT_EQ(result.Value().eigenvectors.Rows(), 1U);
  EXPECT_EQ(result.Value().eigenvectors.Columns(), 1U);
  EXPECT_EQ(result.Value().thread_count, 1U);
  ASSERT_FALSE(result.Value().backend_calls.empty());
  EXPECT_EQ(result.Value().backend_calls.front().purpose, "total-enclosure");
  EXPECT_EQ(result.Value().backend_calls.back().kind, "solve");
  EXPECT_TRUE(result.Value().backend_calls.back().count_available);
  EXPECT_EQ(result.Value().backend_calls.back().count,
            result.Value().eigenvalues.size());
  EXPECT_TRUE(result.Value().backend_calls.back().loop_available);
  EXPECT_TRUE(result.Value().backend_calls.back().epsout_available);
  EXPECT_EQ(result.Value().backend_calls.back().outcome, "accepted");
  EXPECT_EQ(result.Value().backend_calls.back().decision_reason, "accepted");
  EXPECT_FALSE(result.Value().upper_search_termination.empty());
  ExpectUnitGram(result.Value().eigenvectors);
}

TEST(MklFeastEigenSolver, SolvesAllZeroWithExplicitPhysicalScale) {
  const auto matrix = MakeDenseCsr(2U, {0.0, 0.0, 0.0, 0.0});
  fesa::MklFeastEigenSolver solver;

  const auto result = solver.Solve(matrix, 2U, 7.0);

  RequireSuccess(result);
  if (!result.HasValue()) return;
  ASSERT_EQ(result.Value().eigenvalues.size(), 2U);
  EXPECT_DOUBLE_EQ(result.Value().operator_scale, 7.0);
  EXPECT_DOUBLE_EQ(result.Value().eigenvalues[0U], 0.0);
  EXPECT_DOUBLE_EQ(result.Value().eigenvalues[1U], 0.0);
  ExpectUnitGram(result.Value().eigenvectors);
}

TEST(MklFeastEigenSolver, PreservesRepeatedMultiplicityAndStableOrder) {
  const auto matrix =
      MakeDenseCsr(3U, {2.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 5.0});
  fesa::MklFeastEigenSolver solver;

  const auto first = solver.Solve(matrix, 3U, 1.0);
  const auto second = solver.Solve(matrix, 3U, 1.0);

  RequireSuccess(first);
  RequireSuccess(second);
  if (!first.HasValue() || !second.HasValue()) return;
  ASSERT_EQ(first.Value().eigenvalues.size(), 3U);
  EXPECT_NEAR(first.Value().eigenvalues[0U], 2.0, 1.0e-10);
  EXPECT_NEAR(first.Value().eigenvalues[1U], 2.0, 1.0e-10);
  EXPECT_NEAR(first.Value().eigenvalues[2U], 5.0, 1.0e-10);
  EXPECT_EQ(first.Value().column_ids, second.Value().column_ids);
  ExpectUnitGram(first.Value().eigenvectors);
}

TEST(MklFeastEigenSolver, UsesGuardedPrefixWithoutDroppingRepeatedCluster) {
  constexpr std::size_t kSize = 64U;
  std::vector<double> values(kSize * kSize, 0.0);
  for (std::size_t index = 0U; index < kSize; ++index) {
    values[index * kSize + index] = index < 40U ? 1.0 : index - 38U;
  }
  const auto matrix = MakeDenseCsr(kSize, values);
  fesa::MklFeastEigenSolver solver;

  const auto one_mode = solver.Solve(matrix, 1U, 1.0);
  const auto forty_first = solver.Solve(matrix, 41U, 1.0);

  RequireSuccess(one_mode);
  RequireSuccess(forty_first);
  if (!one_mode.HasValue() || !forty_first.HasValue()) return;
  EXPECT_GE(one_mode.Value().eigenvalues.size(), 40U);
  EXPECT_GE(forty_first.Value().eigenvalues.size(), 41U);
  for (std::size_t index = 0U; index < 40U; ++index) {
    EXPECT_NEAR(one_mode.Value().eigenvalues[index], 1.0, 1.0e-10);
  }
  EXPECT_NEAR(forty_first.Value().eigenvalues[40U], 2.0, 1.0e-10);
  EXPECT_EQ(one_mode.Value().backend_calls.back().kind, "solve");
  EXPECT_EQ(forty_first.Value().backend_calls.back().kind, "solve");
}

TEST(MklFeastEigenSolver, RejectsSignificantNegativeSpectrum) {
  const auto matrix = MakeDenseCsr(1U, {-2.0});
  fesa::MklFeastEigenSolver solver;

  const auto result = solver.Solve(matrix, 1U, 1.0);

  ASSERT_FALSE(result.HasValue());
  EXPECT_EQ(result.GetStatus().Category(), fesa::FailureCategory::kSolver);
  ASSERT_EQ(result.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(result.GetStatus().Diagnostics()[0U].code,
            "negative-eigenvalue-spectrum");
}

TEST(MklFeastEigenSolver, KeepsCloseDistinctModesSeparated) {
  const auto matrix = MakeDenseCsr(2U, {1.0, 0.0, 0.0, 1.001});
  fesa::MklFeastEigenSolver solver;

  const auto result = solver.Solve(matrix, 2U, 1.0);

  RequireSuccess(result);
  if (!result.HasValue()) return;
  ASSERT_EQ(result.Value().eigenvalues.size(), 2U);
  EXPECT_NEAR(result.Value().eigenvalues[0U], 1.0, 1.0e-10);
  EXPECT_NEAR(result.Value().eigenvalues[1U], 1.001, 1.0e-10);
  EXPECT_NEAR(std::abs(result.Value().eigenvectors(0U, 0U)), 1.0, 1.0e-10);
  EXPECT_NEAR(std::abs(result.Value().eigenvectors(1U, 1U)), 1.0, 1.0e-10);
}

TEST(MklFeastEigenSolver, RestoresCallingThreadLocalSetting) {
  const auto matrix = MakeDenseCsr(1U, {4.0});
  const int prior = mkl_set_num_threads_local(2);
  fesa::MklFeastEigenSolver solver;

  const auto result = solver.Solve(matrix, 1U, 1.0);
  const int restored = mkl_set_num_threads_local(2);
  static_cast<void>(mkl_set_num_threads_local(prior));

  RequireSuccess(result);
  if (!result.HasValue()) return;
  EXPECT_EQ(restored, 2);
}

TEST(MklFeastEigenSolver, RejectsInvalidRequestAndZeroScale) {
  const auto matrix = MakeDenseCsr(1U, {0.0});
  fesa::MklFeastEigenSolver solver;

  const auto zero_request = solver.Solve(matrix, 0U, 1.0);
  ASSERT_FALSE(zero_request.HasValue());
  EXPECT_EQ(zero_request.GetStatus().Diagnostics()[0U].code,
            "feast-invalid-dimension");
  const auto invalid_scale = solver.Solve(matrix, 1U, 0.0);
  ASSERT_FALSE(invalid_scale.HasValue());
  EXPECT_EQ(invalid_scale.GetStatus().Diagnostics()[0U].code,
            "feast-invalid-zero-scale");
}

TEST(MklFeastEigenSolver, PreservesLowestPrefixForZeroAndPositiveSpectra) {
  const std::vector<std::vector<double>> spectra{{2.0},
                                                 {0.0, 0.0, 0.0, 2.0, 5.0, 9.0},
                                                 {1.0, 2.0, 4.0, 8.0},
                                                 {0.0, 0.0, 0.0, 0.0}};
  const std::vector<std::vector<std::size_t>> requests{
      {1U}, {1U, 4U, 6U}, {1U, 2U, 4U}, {1U, 4U}};
  const fesa::MklFeastEigenSolver solver;
  for (std::size_t case_index = 0U; case_index < spectra.size(); ++case_index) {
    const auto& spectrum = spectra[case_index];
    const auto n = spectrum.size();
    std::vector<double> values(n * n, 0.0);
    for (std::size_t i = 0U; i < n; ++i) values[i * n + i] = spectrum[i];
    const auto matrix = MakeDenseCsr(n, values);
    for (const auto requested : requests[case_index]) {
      const auto result = solver.Solve(matrix, requested, 1.0);
      RequireSuccess(result);
      ASSERT_TRUE(result.HasValue());
      ASSERT_GE(result.Value().eigenvalues.size(), requested);
      for (std::size_t i = 0U; i < requested; ++i)
        EXPECT_NEAR(result.Value().eigenvalues[i], spectrum[i], 1e-10);
      ExpectUnitGram(result.Value().eigenvectors);
    }
  }
}

TEST(MklFeastEigenSolver, ResolvesCloseDistinctSpectrumAfterGivensCongruence) {
  const std::array<double, 4> spectrum{1.0, 1.0 + std::ldexp(1.0, -28),
                                       1.0 + std::ldexp(1.0, -27), 8.0};
  const std::array<std::array<double, 4>, 4> q{{{{0.6, 0.0, 0.0, -0.8}},
                                                {{0.0, 0.6, -0.8, 0.0}},
                                                {{0.0, 0.8, 0.6, 0.0}},
                                                {{0.8, 0.0, 0.0, 0.6}}}};
  std::vector<double> values(16U, 0.0);
  for (std::size_t i = 0U; i < 4U; ++i)
    for (std::size_t j = 0U; j < 4U; ++j)
      for (std::size_t k = 0U; k < 4U; ++k)
        values[4U * i + j] += q[i][k] * spectrum[k] * q[j][k];
  const auto result =
      fesa::MklFeastEigenSolver{}.Solve(MakeDenseCsr(4U, values), 2U, 1.0);
  RequireSuccess(result);
  ASSERT_TRUE(result.HasValue());
  ASSERT_EQ(result.Value().eigenvalues.size(), 4U);
  for (std::size_t i = 0U; i < 4U; ++i)
    EXPECT_NEAR(result.Value().eigenvalues[i], spectrum[i], 1e-12);
  EXPECT_GT(result.Value().eigenvalues[1U] - result.Value().eigenvalues[0U],
            1e-9);
  ExpectUnitGram(result.Value().eigenvectors);
}