#include "fesa/solvers/linear/mkl_pardiso_tangent_solver.h"

#include <gtest/gtest.h>

#include <cstddef>
#include <future>
#include <initializer_list>
#include <limits>
#include <type_traits>
#include <utility>
#include <vector>

#include "fesa/fem/dof_manager.h"
#include "fesa/math/sparse_matrix.h"
#include "fesa/math/vector.h"
#include "fesa/solvers/linear/tangent_solver.h"
#include "math/dense_blas_internal.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.reserve(size + 1U);
  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());
}

fesa::Vector MakeVector(const std::initializer_list<double> values) {
  fesa::Vector result{values.size()};
  std::size_t index = 0U;
  for (const double value : values) {
    result[index++] = value;
  }
  return result;
}

void ExpectSolverFailure(const fesa::Status& status) {
  EXPECT_FALSE(status.IsOk());
  EXPECT_EQ(status.Category(), fesa::FailureCategory::kSolver);
  ASSERT_EQ(status.Diagnostics().size(), 1U);
  EXPECT_EQ(status.Diagnostics().front().severity, fesa::Severity::kError);
}

TEST(TangentSolver, SolvesNonsymmetricKnownCorrection) {
  static_assert(
      std::is_base_of_v<fesa::TangentSolver, fesa::MklPardisoTangentSolver>);
  static_assert(std::has_virtual_destructor_v<fesa::TangentSolver>);

  const auto matrix = MakeDenseCsr(2U, {2.0, 1.0, 0.0, 3.0});
  const auto rhs = MakeVector({4.0, 6.0});
  fesa::MklPardisoTangentSolver concrete_solver;
  fesa::TangentSolver& solver = concrete_solver;

  ASSERT_TRUE(solver.FactorizeFresh(matrix).IsOk());
  fesa::Vector correction{2U};
  ASSERT_TRUE(solver.SolveCorrection(rhs, correction).IsOk());
  EXPECT_NEAR(correction[0U], 1.0, 1.0e-14);
  EXPECT_NEAR(correction[1U], 2.0, 1.0e-14);
}

TEST(TangentSolver, SolvesTrivialZeroEquationSystem) {
  const auto matrix = MakeDenseCsr(0U, {});
  fesa::MklPardisoTangentSolver solver;

  ASSERT_TRUE(solver.FactorizeFresh(matrix).IsOk());
  fesa::Vector correction{0U};
  EXPECT_TRUE(solver.SolveCorrection(fesa::Vector{0U}, correction).IsOk());
  EXPECT_EQ(correction.Size(), 0U);
}

TEST(TangentSolver, RejectsNonfiniteInputsWithoutModifyingCorrection) {
  const fesa::SparsePattern pattern{{0U, 1U}, {0U}};
  const auto nonfinite_matrix = fesa::SparseMatrix::FromCoo(
      1U, 1U,
      {{0U, 0U, (std::numeric_limits<double>::quiet_NaN)(), 0U, 0U}},
      pattern);
  EXPECT_FALSE(nonfinite_matrix.HasValue());
  EXPECT_EQ(nonfinite_matrix.GetStatus().Diagnostics().front().code,
            "nonfinite-sparse-value");

  fesa::MklPardisoTangentSolver solver;
  ASSERT_TRUE(solver.FactorizeFresh(MakeDenseCsr(2U, {2.0, 1.0, 0.0, 3.0}))
                  .IsOk());
  auto rhs = MakeVector({4.0, 6.0});
  rhs[1U] = (std::numeric_limits<double>::infinity)();
  fesa::Vector correction = MakeVector({17.0, -4.0});

  const auto status = solver.SolveCorrection(rhs, correction);
  ExpectSolverFailure(status);
  EXPECT_EQ(status.Diagnostics().front().code, "nonfinite-solver-rhs");
  EXPECT_DOUBLE_EQ(correction[0U], 17.0);
  EXPECT_DOUBLE_EQ(correction[1U], -4.0);
}

TEST(TangentSolver, ReportsSingularFactorizationDeterministically) {
  fesa::MklPardisoTangentSolver solver;

  const auto status =
      solver.FactorizeFresh(MakeDenseCsr(2U, {0.0, 0.0, 0.0, 0.0}));
  ASSERT_FALSE(status.IsOk());
  ExpectSolverFailure(status);
  EXPECT_EQ(status.Diagnostics().front().code, "pardiso-perturbed-pivot");
  EXPECT_EQ(status.Diagnostics().front().entity_identity, "phase=22");
}

TEST(TangentSolver, RejectsSolveBeforeSuccessfulFactorization) {
  fesa::MklPardisoTangentSolver solver;
  fesa::Vector correction = MakeVector({23.0, -9.0});

  const auto status =
      solver.SolveCorrection(MakeVector({1.0, 2.0}), correction);
  ExpectSolverFailure(status);
  EXPECT_EQ(status.Diagnostics().front().code, "solver-not-factorized");
  EXPECT_DOUBLE_EQ(correction[0U], 23.0);
  EXPECT_DOUBLE_EQ(correction[1U], -9.0);
}

TEST(TangentSolver, ReusesRhsButFreshValuesAndFailureCannotReuseStaleState) {
  fesa::MklPardisoTangentSolver solver;
  ASSERT_TRUE(solver.FactorizeFresh(MakeDenseCsr(2U, {2.0, 1.0, 0.0, 3.0}))
                  .IsOk());
  fesa::Vector correction{2U};
  for (double scale : {1.0, -2.0, 0.5}) {
    ASSERT_TRUE(solver.SolveCorrection(MakeVector({4.0 * scale, 6.0 * scale}),
                                      correction).IsOk());
    EXPECT_NEAR(correction[0U], scale, 1.0e-14);
    EXPECT_NEAR(correction[1U], 2.0 * scale, 1.0e-14);
  }
  const auto bad_rhs = MakeVector({std::numeric_limits<double>::quiet_NaN(), 0});
  correction = MakeVector({17, -4});
  EXPECT_FALSE(solver.SolveCorrection(bad_rhs, correction).IsOk());
  EXPECT_DOUBLE_EQ(correction[0U], 17);
  EXPECT_DOUBLE_EQ(correction[1U], -4);
  ASSERT_TRUE(solver.SolveCorrection(MakeVector({4, 6}), correction).IsOk());
  ASSERT_TRUE(solver.FactorizeFresh(MakeDenseCsr(2U, {4, 0, 1, 2})).IsOk());
  ASSERT_TRUE(solver.SolveCorrection(MakeVector({4, 5}), correction).IsOk());
  EXPECT_NEAR(correction[0U], 1, 1.0e-14);
  EXPECT_NEAR(correction[1U], 2, 1.0e-14);
  auto diagonal = fesa::SparseMatrix::FromCoo(
      2U, 2U, {{0U, 0U, 5, 0U, 0U}, {1U, 1U, 7, 1U, 0U}},
      {{0U, 1U, 2U}, {0U, 1U}});
  ASSERT_TRUE(diagonal.HasValue());
  ASSERT_TRUE(solver.FactorizeFresh(diagonal.Value()).IsOk());
  ASSERT_TRUE(solver.SolveCorrection(MakeVector({10, -7}), correction).IsOk());
  EXPECT_NEAR(correction[0U], 2, 1.0e-14);
  EXPECT_NEAR(correction[1U], -1, 1.0e-14);
  EXPECT_FALSE(solver.FactorizeFresh(MakeDenseCsr(2U, {0, 0, 0, 0})).IsOk());
  correction = MakeVector({19, -8});
  const auto stale = solver.SolveCorrection(MakeVector({10, -7}), correction);
  ASSERT_FALSE(stale.IsOk());
  EXPECT_EQ(stale.Diagnostics().front().code, "solver-not-factorized");
  EXPECT_DOUBLE_EQ(correction[0U], 19);
  EXPECT_DOUBLE_EQ(correction[1U], -8);
  ASSERT_TRUE(solver.FactorizeFresh(diagonal.Value()).IsOk());
  EXPECT_TRUE(solver.SolveCorrection(MakeVector({10, -7}), correction).IsOk());
}

TEST(TangentSolver, SingleEquationScopeRestoresTlsAcrossFreshFailureAndReuse) {
  const auto worker = [] {
    const int previous = mkl_set_num_threads_local(0);
    struct Restore {
      int previous;
      ~Restore() { static_cast<void>(mkl_set_num_threads_local(previous)); }
    } restore{previous};
    const int default_threads = mkl_get_max_threads();
    fesa::MklPardisoTangentSolver solver;
    for (int repeat = 0; repeat < 10; ++repeat) {
      auto correction = MakeVector({17.0});
      if (!solver.FactorizeFresh(MakeDenseCsr(1U, {2.0})).IsOk()) return false;
      if (!solver.SolveCorrection(MakeVector({4.0}), correction).IsOk() ||
          correction[0U] != 2.0) return false;
      if (!solver.FactorizeFresh(MakeDenseCsr(1U, {4.0})).IsOk()) return false;
      if (!solver.SolveCorrection(MakeVector({4.0}), correction).IsOk() ||
          correction[0U] != 1.0) return false;
      if (solver.SolveCorrection(MakeVector({
              std::numeric_limits<double>::infinity()}), correction).IsOk() ||
          correction[0U] != 1.0) return false;
      if (solver.FactorizeFresh(MakeDenseCsr(1U, {0.0})).IsOk()) return false;
      if (solver.SolveCorrection(MakeVector({4.0}), correction).IsOk() ||
          correction[0U] != 1.0) return false;
      if (mkl_get_max_threads() != default_threads ||
          mkl_set_num_threads_local(0) != 0) return false;
    }
    return true;
  };
  auto first = std::async(std::launch::async, worker);
  auto second = std::async(std::launch::async, worker);
  EXPECT_TRUE(first.get());
  EXPECT_TRUE(second.get());
}

}  // namespace