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

#include <gtest/gtest.h>

#include <cmath>
#include <initializer_list>
#include <limits>
#include <string>
#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.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;
}

double NormalizedResidual(const fesa::SparseMatrix& matrix,
                          const fesa::Vector& solution,
                          const fesa::Vector& rhs) {
  auto residual = matrix.Multiply(solution);
  residual.Axpy(-1.0, rhs);
  const double numerator = residual.Norm();
  const double denominator = rhs.Norm();
  if (!std::isfinite(numerator) || !std::isfinite(denominator)) {
    return (std::numeric_limits<double>::infinity)();
  }
  if (denominator == 0.0) {
    return numerator == 0.0 ? 0.0 : (std::numeric_limits<double>::infinity)();
  }
  return numerator / denominator;
}

double RelativeError(const fesa::Vector& actual, const fesa::Vector& expected) {
  auto difference = actual;
  difference.Axpy(-1.0, expected);
  const double numerator = difference.Norm();
  const double denominator = expected.Norm();
  if (!std::isfinite(numerator) || !std::isfinite(denominator)) {
    return (std::numeric_limits<double>::infinity)();
  }
  if (denominator == 0.0) {
    return numerator == 0.0 ? 0.0 : (std::numeric_limits<double>::infinity)();
  }
  return numerator / denominator;
}

void ExpectStructuredSolverFailure(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()[0U].severity, fesa::Severity::kError);
  EXPECT_FALSE(status.Diagnostics()[0U].code.empty());
  EXPECT_FALSE(status.Diagnostics()[0U].message.empty());
}

}  // namespace

TEST(MklPardisoSolver, SolvesKnownSpdWithNormalizedResidual) {
  const auto matrix =
      MakeDenseCsr(3U, {6.0, 2.0, 1.0, 2.0, 5.0, 2.0, 1.0, 2.0, 4.0});
  const auto expected = MakeVector({1.0, -2.0, 3.0});
  const auto rhs = matrix.Multiply(expected);

  fesa::MklPardisoSolver concrete_solver;
  fesa::LinearSolver& solver = concrete_solver;
  ASSERT_TRUE(solver.Factorize(matrix).IsOk());

  fesa::Vector solution{3U};
  ASSERT_TRUE(solver.Solve(rhs, solution).IsOk());
  EXPECT_LE(NormalizedResidual(matrix, solution, rhs), 1.0e-10);
  EXPECT_LE(RelativeError(solution, expected), 1.0e-9);
}

TEST(MklPardisoSolver, ReusesOneFactorizationForRepeatedRhs) {
  const auto matrix = MakeDenseCsr(2U, {4.0, 1.0, 1.0, 3.0});
  const auto expected_first = MakeVector({1.0, 2.0});
  const auto expected_second = MakeVector({-2.0, 0.5});
  const auto rhs_first = matrix.Multiply(expected_first);
  const auto rhs_second = matrix.Multiply(expected_second);

  fesa::MklPardisoSolver solver;
  ASSERT_TRUE(solver.Factorize(matrix).IsOk());
  fesa::Vector first{2U};
  fesa::Vector second{2U};
  ASSERT_TRUE(solver.Solve(rhs_first, first).IsOk());
  ASSERT_TRUE(solver.Solve(rhs_second, second).IsOk());

  EXPECT_LE(RelativeError(first, expected_first), 1.0e-9);
  EXPECT_LE(RelativeError(second, expected_second), 1.0e-9);
  EXPECT_LE(NormalizedResidual(matrix, first, rhs_first), 1.0e-10);
  EXPECT_LE(NormalizedResidual(matrix, second, rhs_second), 1.0e-10);
}

TEST(MklPardisoSolver, RefactorizesWithoutLeakingState) {
  const auto first_matrix = MakeDenseCsr(2U, {4.0, 1.0, 1.0, 3.0});
  const auto second_matrix = MakeDenseCsr(2U, {2.0, 0.0, 0.0, 5.0});
  const auto first_expected = MakeVector({1.0, 2.0});
  const auto second_expected = MakeVector({-3.0, 4.0});

  fesa::MklPardisoSolver solver;
  ASSERT_TRUE(solver.Factorize(first_matrix).IsOk());
  fesa::Vector first_solution{2U};
  ASSERT_TRUE(
      solver.Solve(first_matrix.Multiply(first_expected), first_solution)
          .IsOk());
  EXPECT_LE(RelativeError(first_solution, first_expected), 1.0e-9);

  ASSERT_TRUE(solver.Factorize(second_matrix).IsOk());
  fesa::Vector second_solution{2U};
  const auto second_rhs = second_matrix.Multiply(second_expected);
  ASSERT_TRUE(solver.Solve(second_rhs, second_solution).IsOk());
  EXPECT_LE(RelativeError(second_solution, second_expected), 1.0e-9);
  EXPECT_LE(NormalizedResidual(second_matrix, second_solution, second_rhs),
            1.0e-10);
}

TEST(MklPardisoSolver, ClassifiesSingularIndefiniteAndNonfiniteFailures) {
  fesa::MklPardisoSolver solver;

  const auto singular = MakeDenseCsr(2U, {1.0, 1.0, 1.0, 1.0});
  const auto singular_status = solver.Factorize(singular);
  ExpectStructuredSolverFailure(singular_status);
  EXPECT_TRUE(singular_status.Diagnostics()[0U].code ==
                  "pardiso-zero-or-negative-pivot" ||
              singular_status.Diagnostics()[0U].code ==
                  "pardiso-singular-diagonal");
  EXPECT_NE(
      singular_status.Diagnostics()[0U].entity_identity.find("phase=22,error="),
      std::string::npos);

  const auto indefinite = MakeDenseCsr(2U, {1.0, 2.0, 2.0, 1.0});
  const auto indefinite_status = solver.Factorize(indefinite);
  ExpectStructuredSolverFailure(indefinite_status);
  EXPECT_TRUE(indefinite_status.Diagnostics()[0U].code ==
                  "pardiso-zero-or-negative-pivot" ||
              indefinite_status.Diagnostics()[0U].code ==
                  "pardiso-singular-diagonal");
  EXPECT_NE(indefinite_status.Diagnostics()[0U].entity_identity.find(
                "phase=22,error="),
            std::string::npos);

  const auto spd = MakeDenseCsr(2U, {3.0, 1.0, 1.0, 2.0});
  ASSERT_TRUE(solver.Factorize(spd).IsOk());
  auto rhs = MakeVector({1.0, 2.0});
  rhs[1U] = (std::numeric_limits<double>::infinity)();
  fesa::Vector solution{2U};
  solution[0U] = 23.0;
  solution[1U] = -9.0;
  const auto rhs_status = solver.Solve(rhs, solution);
  ExpectStructuredSolverFailure(rhs_status);
  EXPECT_EQ(rhs_status.Diagnostics()[0U].code, "nonfinite-solver-rhs");
  EXPECT_DOUBLE_EQ(solution[0U], 23.0);
  EXPECT_DOUBLE_EQ(solution[1U], -9.0);

  fesa::SparsePattern pattern{{0U, 1U}, {0U}};
  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()[0U].code,
            "nonfinite-sparse-value");
}

TEST(MklPardisoSolver,
     ConditioningSweepPassesResolvedCasesAndFailsUnresolvedCasesExplicitly) {
  const std::vector<double> common_scales{1.0e-12, 1.0, 1.0e12};
  for (const double scale : common_scales) {
    const auto matrix =
        MakeDenseCsr(2U, {4.0 * scale, 1.0 * scale, 1.0 * scale, 3.0 * scale});
    const auto expected = MakeVector({1.25, -0.75});
    const auto rhs = matrix.Multiply(expected);

    fesa::MklPardisoSolver solver;
    ASSERT_TRUE(solver.Factorize(matrix).IsOk()) << "scale=" << scale;
    fesa::Vector solution{2U};
    ASSERT_TRUE(solver.Solve(rhs, solution).IsOk()) << "scale=" << scale;
    EXPECT_LE(NormalizedResidual(matrix, solution, rhs), 1.0e-10);
    EXPECT_LE(RelativeError(solution, expected), 1.0e-9);
  }

  const double resolved_ratio = 1.0e-8;
  const auto resolved_matrix =
      MakeDenseCsr(2U, {1.0, 0.0, 0.0, resolved_ratio});
  const auto resolved_expected = MakeVector({0.5, -2.0});
  const auto resolved_rhs = resolved_matrix.Multiply(resolved_expected);
  fesa::MklPardisoSolver resolved_solver;
  ASSERT_TRUE(resolved_solver.Factorize(resolved_matrix).IsOk());
  fesa::Vector resolved_solution{2U};
  ASSERT_TRUE(resolved_solver.Solve(resolved_rhs, resolved_solution).IsOk());
  EXPECT_LE(
      NormalizedResidual(resolved_matrix, resolved_solution, resolved_rhs),
      1.0e-10);
  EXPECT_LE(RelativeError(resolved_solution, resolved_expected), 1.0e-9);

  const std::vector<double> unresolved_candidates{1.0e-16, 1.0e-300};
  for (const double ratio : unresolved_candidates) {
    const auto matrix = MakeDenseCsr(2U, {1.0, 0.0, 0.0, ratio});
    const auto expected = MakeVector({0.5, -2.0});
    const auto rhs = matrix.Multiply(expected);
    fesa::MklPardisoSolver solver;

    const auto factor_status = solver.Factorize(matrix);
    if (!factor_status.IsOk()) {
      ExpectStructuredSolverFailure(factor_status);
      continue;
    }

    fesa::Vector solution{2U};
    const auto solve_status = solver.Solve(rhs, solution);
    if (!solve_status.IsOk()) {
      ExpectStructuredSolverFailure(solve_status);
      continue;
    }

    EXPECT_LE(NormalizedResidual(matrix, solution, rhs), 1.0e-10);
    EXPECT_LE(RelativeError(solution, expected), 1.0e-9);
  }
}