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

#include <gtest/gtest.h>

#include <limits>
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

#include "fesa/analysis/linear_frequency_procedure.h"
#include "fesa/analysis/solid_frequency_procedure.h"
#include "fesa/solvers/eigen/eigen_solver.h"
#include "fesa/solvers/eigen/mkl_feast_eigen_solver.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/math/sparse_matrix.h"
#include "fesa/solvers/linear/mkl_pardiso_solver.h"
#include "fesa/solvers/linear/mkl_pardiso_tangent_solver.h"
#include "solvers/linear/pardiso_session_test_seam.h"

namespace {

fesa::SparseMatrix MakeDenseCsr(const std::size_t rows,
                                const std::size_t columns,
                                const std::vector<double>& values) {
  EXPECT_EQ(values.size(), rows * columns);

  fesa::SparsePattern pattern;
  std::vector<fesa::CooContribution> contributions;
  pattern.row_offsets.reserve(rows + 1U);
  pattern.row_offsets.push_back(0U);
  for (std::size_t row = 0U; row < rows; ++row) {
    for (std::size_t column = 0U; column < columns; ++column) {
      pattern.column_indices.push_back(column);
      contributions.push_back(
          {row, column, values[row * columns + column], row, column});
    }
    pattern.row_offsets.push_back(pattern.column_indices.size());
  }

  auto matrix = fesa::SparseMatrix::FromCoo(rows, columns,
                                            std::move(contributions), pattern);
  EXPECT_TRUE(matrix.HasValue());
  return std::move(matrix.Value());
}

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

}  // namespace

TEST(MklPardisoSolver, RejectsInvalidCsrStateAndDimensions) {
  static_assert(std::is_base_of_v<fesa::LinearSolver, fesa::MklPardisoSolver>);
  static_assert(std::has_virtual_destructor_v<fesa::LinearSolver>);

  fesa::MklPardisoSolver solver;
  fesa::Vector untouched{2U};
  untouched[0U] = 17.0;
  untouched[1U] = -4.0;
  const auto before_factorize = solver.Solve(fesa::Vector{2U, 1.0}, untouched);
  ExpectSolverFailure(before_factorize);
  EXPECT_EQ(before_factorize.Diagnostics().front().code,
            "solver-not-factorized");
  EXPECT_DOUBLE_EQ(untouched[0U], 17.0);
  EXPECT_DOUBLE_EQ(untouched[1U], -4.0);

  // A fully constrained model has a valid 0x0 Kff. It still observes the
  // factorize-then-solve lifecycle without invoking a numerical backend.
  const auto empty = MakeDenseCsr(0U, 0U, {});
  ASSERT_TRUE(solver.Factorize(empty).IsOk());
  fesa::Vector empty_solution{0U};
  EXPECT_TRUE(solver.Solve(fesa::Vector{0U}, empty_solution).IsOk());
  EXPECT_EQ(empty_solution.Size(), 0U);

  // Refactorization from the trivial state must establish ordinary PARDISO
  // state rather than retaining a zero-equation shortcut.
  const auto spd = MakeDenseCsr(2U, 2U, {4.0, 1.0, 1.0, 3.0});
  ASSERT_TRUE(solver.Factorize(spd).IsOk());
  fesa::Vector solution{2U};
  ASSERT_TRUE(solver.Solve(fesa::Vector{2U, 1.0}, solution).IsOk());
  EXPECT_NEAR(solution[0U], 2.0 / 11.0, 1.0e-14);
  EXPECT_NEAR(solution[1U], 3.0 / 11.0, 1.0e-14);

  const double solved_first = solution[0U];
  const double solved_second = solution[1U];
  ExpectSolverFailure(solver.Solve(fesa::Vector{1U, 1.0}, solution));
  EXPECT_DOUBLE_EQ(solution[0U], solved_first);
  EXPECT_DOUBLE_EQ(solution[1U], solved_second);

  fesa::Vector wrong_solution{1U};
  wrong_solution[0U] = 41.0;
  ExpectSolverFailure(solver.Solve(fesa::Vector{2U, 1.0}, wrong_solution));
  EXPECT_DOUBLE_EQ(wrong_solution[0U], 41.0);

  const auto rectangular = MakeDenseCsr(2U, 3U, {2.0, 0.0, 0.0, 0.0, 3.0, 0.0});
  const auto rectangular_status = solver.Factorize(rectangular);
  ExpectSolverFailure(rectangular_status);
  EXPECT_EQ(rectangular_status.Diagnostics().front().code,
            "solver-matrix-not-square");

  fesa::SparsePattern invalid_pattern{{0U, 2U}, {0U}};
  auto invalid_csr = fesa::SparseMatrix::FromCoo(
      1U, 1U, {{0U, 0U, 1.0, 0U, 0U}}, invalid_pattern);
  EXPECT_FALSE(invalid_csr.HasValue());

  const auto nonsymmetric = MakeDenseCsr(2U, 2U, {2.0, 1.0, 0.0, 3.0});
  const auto nonsymmetric_status = solver.Factorize(nonsymmetric);
  ExpectSolverFailure(nonsymmetric_status);
  EXPECT_EQ(nonsymmetric_status.Diagnostics().front().code,
            "solver-matrix-not-symmetric");

  const auto scaled_nonsymmetric =
      MakeDenseCsr(2U, 2U, {2.0e-20, 1.0e-20, 1.1e-20, 3.0e-20});
  const auto scaled_nonsymmetric_status = solver.Factorize(scaled_nonsymmetric);
  // Stop this case before inspecting diagnostics when the production code
  // incorrectly accepts the matrix; this keeps the RED failure deterministic.
  ASSERT_FALSE(scaled_nonsymmetric_status.IsOk());
  ExpectSolverFailure(scaled_nonsymmetric_status);
  EXPECT_EQ(scaled_nonsymmetric_status.Diagnostics().front().code,
            "solver-matrix-not-symmetric");

  fesa::SparsePattern no_diagonal_pattern{{0U, 1U, 2U}, {1U, 0U}};
  auto no_diagonal = fesa::SparseMatrix::FromCoo(
      2U, 2U, {{0U, 1U, 1.0, 0U, 0U}, {1U, 0U, 1.0, 1U, 0U}},
      no_diagonal_pattern);
  ASSERT_TRUE(no_diagonal.HasValue());
  const auto no_diagonal_status = solver.Factorize(no_diagonal.Value());
  ExpectSolverFailure(no_diagonal_status);
  EXPECT_EQ(no_diagonal_status.Diagnostics().front().code,
            "solver-missing-diagonal");
}

namespace {

class RecordingEigenSolver final : public fesa::EigenSolver {
 public:
  explicit RecordingEigenSolver(bool& destroyed) : destroyed_(destroyed) {}
  ~RecordingEigenSolver() override { destroyed_ = true; }

  fesa::Result<fesa::FeastEigenResult> Solve(
      const fesa::SparseMatrix& matrix, std::size_t modes,
      double scale) const override {
    EXPECT_EQ(matrix.Rows(), 0U);
    EXPECT_EQ(modes, 3U);
    EXPECT_DOUBLE_EQ(scale, 7.0);
    ++standard_calls;
    return fesa::Result<fesa::FeastEigenResult>::Failure(
        fesa::Status::Failure(fesa::FailureCategory::kSolver, {}));
  }

  fesa::Result<fesa::FeastEigenResult> SolveGeneralized(
      const fesa::SparseMatrix& stiffness, const fesa::SparseMatrix& mass,
      std::size_t modes, double lower_bound,
      double scale) const override {
    EXPECT_EQ(stiffness.Rows(), 0U);
    EXPECT_EQ(mass.Rows(), 0U);
    EXPECT_EQ(modes, 4U);
    EXPECT_DOUBLE_EQ(lower_bound, 2.0);
    EXPECT_DOUBLE_EQ(scale, 11.0);
    ++generalized_calls;
    return fesa::Result<fesa::FeastEigenResult>::Failure(
        fesa::Status::Failure(fesa::FailureCategory::kSolver, {}));
  }

  mutable int standard_calls{0};
  mutable int generalized_calls{0};

 private:
  bool& destroyed_;
};

}  // namespace

TEST(EigenSolverContract, DispatchesBothPencilsAndDestroysThroughBase) {
  static_assert(std::is_abstract_v<fesa::EigenSolver>);
  static_assert(std::has_virtual_destructor_v<fesa::EigenSolver>);
  static_assert(std::is_base_of_v<fesa::EigenSolver,
                                fesa::MklFeastEigenSolver>);
  static_assert(std::is_constructible_v<fesa::LinearFrequencyProcedureExecutor,
                                        const fesa::EigenSolver&>);
  static_assert(std::is_constructible_v<fesa::SolidFrequencyProcedureExecutor,
                                        const fesa::EigenSolver&>);
  bool destroyed = false;
  {
    auto recording = std::make_unique<RecordingEigenSolver>(destroyed);
    const auto* observations = recording.get();
    std::unique_ptr<fesa::EigenSolver> solver = std::move(recording);
    const fesa::EigenSolver& contract = *solver;
    const auto empty = MakeDenseCsr(0U, 0U, {});
    EXPECT_FALSE(contract.Solve(empty, 3U, 7.0).HasValue());
    EXPECT_FALSE(contract.SolveGeneralized(empty, empty, 4U, 2.0, 11.0)
                     .HasValue());
    EXPECT_EQ(observations->standard_calls, 1);
    EXPECT_EQ(observations->generalized_calls, 1);
    EXPECT_FALSE(destroyed);
  }
  EXPECT_TRUE(destroyed);
}

TEST(PardisoSession, RetainsDistinctStoragePoliciesAndFreshPhaseOrder) {
  std::vector<int> phases;
  fesa::internal::ScopedPardisoBackendTestHook hook{
      [&](const fesa::internal::PardisoBackendTestCall& call,
          MKL_INT&, MKL_INT&) {
        phases.push_back(static_cast<int>(call.phase));
        if (call.phase == 11) {
          EXPECT_EQ(call.equation_count, 2);
          if (call.matrix_type == 2) {
            EXPECT_EQ(call.row_offsets[2], 3);
            EXPECT_EQ(call.column_indices[0], 0);
            EXPECT_EQ(call.column_indices[1], 1);
            EXPECT_EQ(call.column_indices[2], 1);
          } else {
            EXPECT_EQ(call.matrix_type, 11);
            EXPECT_EQ(call.row_offsets[2], 4);
            EXPECT_DOUBLE_EQ(call.values[2], 0.0);
          }
        }
      }};
  {
    fesa::MklPardisoSolver solver;
    const auto matrix = MakeDenseCsr(2U, 2U, {4.0, 1.0, 1.0, 3.0});
    ASSERT_TRUE(solver.Factorize(matrix).IsOk());
    fesa::Vector solution{2U};
    ASSERT_TRUE(solver.Solve(fesa::Vector{2U, 1.0}, solution).IsOk());
    EXPECT_NEAR(solution[0], 2.0 / 11.0, 1e-14);
    ASSERT_TRUE(solver.Solve(fesa::Vector{2U, 2.0}, solution).IsOk());
    EXPECT_NEAR(solution[1], 6.0 / 11.0, 1e-14);
    ASSERT_TRUE(solver.Factorize(matrix).IsOk());
  }
  EXPECT_EQ(phases, (std::vector<int>{11, 22, 33, 33, -1, 11, 22, -1}));
  phases.clear();
  {
    fesa::MklPardisoTangentSolver solver;
    const auto matrix = MakeDenseCsr(2U, 2U, {2.0, 1.0, 0.0, 3.0});
    ASSERT_TRUE(solver.FactorizeFresh(matrix).IsOk());
    fesa::Vector correction{2U};
    ASSERT_TRUE(solver.SolveCorrection(fesa::Vector{2U, 6.0}, correction).IsOk());
    EXPECT_NEAR(correction[0], 2.0, 1e-14);
    EXPECT_NEAR(correction[1], 2.0, 1e-14);
    ASSERT_TRUE(solver.FactorizeFresh(matrix).IsOk());
  }
  EXPECT_EQ(phases, (std::vector<int>{11, 22, 33, -1, 11, 22, -1}));
}

TEST(PardisoSession, FailedSolvePreservesOutputAndRetainedFactorization) {
  int solve_calls = 0;
  fesa::internal::ScopedPardisoBackendTestHook hook{
      [&](const fesa::internal::PardisoBackendTestCall& call,
          MKL_INT& error, MKL_INT&) {
        if (call.phase != 33) return;
        ++solve_calls;
        if (solve_calls == 1) error = -3;
        if (solve_calls == 2)
          call.solution[0] = (std::numeric_limits<double>::infinity)();
      }};
  fesa::MklPardisoSolver solver;
  ASSERT_TRUE(solver.Factorize(MakeDenseCsr(1U, 1U, {2.0})).IsOk());
  fesa::Vector solution{1U, 17.0};
  const auto failed = solver.Solve(fesa::Vector{1U, 6.0}, solution);
  ASSERT_FALSE(failed.IsOk());
  EXPECT_EQ(failed.Diagnostics().front().code, "pardiso-solve-failed");
  EXPECT_DOUBLE_EQ(solution[0], 17.0);
  const auto nonfinite = solver.Solve(fesa::Vector{1U, 6.0}, solution);
  ASSERT_FALSE(nonfinite.IsOk());
  EXPECT_EQ(nonfinite.Diagnostics().front().code, "nonfinite-solver-solution");
  EXPECT_DOUBLE_EQ(solution[0], 17.0);
  ASSERT_TRUE(solver.Solve(fesa::Vector{1U, 6.0}, solution).IsOk());
  EXPECT_DOUBLE_EQ(solution[0], 3.0);
}

TEST(PardisoSession, ReleaseFailurePrecedesNewValidationAndInvalidatesState) {
  int releases = 0;
  fesa::internal::ScopedPardisoBackendTestHook hook{
      [&](const fesa::internal::PardisoBackendTestCall& call,
          MKL_INT& error, MKL_INT&) {
        if (call.phase == -1) { ++releases; error = -3; }
      }};
  fesa::MklPardisoSolver solver;
  ASSERT_TRUE(solver.Factorize(MakeDenseCsr(1U, 1U, {2.0})).IsOk());
  const auto status = solver.Factorize(MakeDenseCsr(1U, 2U, {1.0, 2.0}));
  ASSERT_FALSE(status.IsOk());
  EXPECT_EQ(status.Diagnostics().front().code, "pardiso-release-failed");
  fesa::Vector output{1U, 19.0};
  const auto unavailable = solver.Solve(fesa::Vector{1U, 1.0}, output);
  ASSERT_FALSE(unavailable.IsOk());
  EXPECT_EQ(unavailable.Diagnostics().front().code, "solver-not-factorized");
  EXPECT_DOUBLE_EQ(output[0], 19.0);
  EXPECT_EQ(releases, 1);
}

TEST(PardisoSession, FactorizationFailuresAndExceptionsReleaseExactlyOnce) {
  const auto matrix = MakeDenseCsr(1U, 1U, {2.0});
  for (const int failing_phase : {11, 22}) {
    std::vector<int> phases;
    fesa::internal::ScopedPardisoBackendTestHook hook{
        [&](const fesa::internal::PardisoBackendTestCall& call,
            MKL_INT& error, MKL_INT&) {
          phases.push_back(static_cast<int>(call.phase));
          if (call.phase == failing_phase) error = -8;
        }};
    {
      fesa::MklPardisoSolver solver;
      const auto status = solver.Factorize(matrix);
      ASSERT_FALSE(status.IsOk());
      EXPECT_EQ(status.Diagnostics().front().code, "pardiso-integer-overflow");
      fesa::Vector output{1U, 29.0};
      EXPECT_FALSE(solver.Solve(fesa::Vector{1U, 1.0}, output).IsOk());
      EXPECT_DOUBLE_EQ(output[0], 29.0);
    }
    EXPECT_EQ(phases, failing_phase == 11 ? (std::vector<int>{11, -1})
                                         : (std::vector<int>{11, 22, -1}));
  }
  std::vector<int> phases;
  fesa::internal::ScopedPardisoBackendTestHook hook{
      [&](const fesa::internal::PardisoBackendTestCall& call,
          MKL_INT&, MKL_INT&) {
        phases.push_back(static_cast<int>(call.phase));
        if (call.phase == 11) throw std::runtime_error("injected unwind");
      }};
  {
    fesa::MklPardisoTangentSolver solver;
    EXPECT_THROW(solver.FactorizeFresh(matrix), std::runtime_error);
    EXPECT_EQ(phases, (std::vector<int>{11, -1}));
  }
  EXPECT_EQ(phases, (std::vector<int>{11, -1}));
}

TEST(PardisoSession, FailedTangentSolveRetainsStateButFailedFreshInvalidatesIt) {
  bool fail_solve = true;
  fesa::internal::ScopedPardisoBackendTestHook hook{
      [&](const fesa::internal::PardisoBackendTestCall& call,
          MKL_INT& error, MKL_INT&) {
        if (call.phase == 33 && fail_solve) error = -3;
      }};
  fesa::MklPardisoTangentSolver solver;
  ASSERT_TRUE(solver.FactorizeFresh(
      MakeDenseCsr(2U, 2U, {2.0, 1.0, 0.0, 3.0})).IsOk());
  fesa::Vector correction{2U, 31.0};
  auto status = solver.SolveCorrection(fesa::Vector{2U, 6.0}, correction);
  ASSERT_FALSE(status.IsOk());
  EXPECT_EQ(status.Diagnostics().front().code, "pardiso-solve-failed");
  EXPECT_DOUBLE_EQ(correction[0], 31.0);
  EXPECT_DOUBLE_EQ(correction[1], 31.0);
  fail_solve = false;
  ASSERT_TRUE(solver.SolveCorrection(fesa::Vector{2U, 6.0}, correction).IsOk());
  EXPECT_NEAR(correction[0], 2.0, 1e-14);
  EXPECT_NEAR(correction[1], 2.0, 1e-14);
  status = solver.FactorizeFresh(MakeDenseCsr(1U, 2U, {1.0, 2.0}));
  ASSERT_FALSE(status.IsOk());
  EXPECT_EQ(status.Diagnostics().front().code, "solver-matrix-not-square");
  status = solver.SolveCorrection(fesa::Vector{2U, 6.0}, correction);
  ASSERT_FALSE(status.IsOk());
  EXPECT_EQ(status.Diagnostics().front().code, "solver-not-factorized");
  EXPECT_NEAR(correction[0], 2.0, 1e-14);
  EXPECT_NEAR(correction[1], 2.0, 1e-14);
}

TEST(PardisoSession, EmptyLifecycleNeverInvokesBackend) {
  int calls = 0;
  fesa::internal::ScopedPardisoBackendTestHook hook{
      [&](const fesa::internal::PardisoBackendTestCall&, MKL_INT&, MKL_INT&) {
        ++calls;
      }};
  {
    fesa::MklPardisoSolver solver;
    fesa::MklPardisoTangentSolver tangent;
    const auto empty = MakeDenseCsr(0U, 0U, {});
    fesa::Vector solution{0U};
    ASSERT_TRUE(solver.Factorize(empty).IsOk());
    ASSERT_TRUE(solver.Solve(fesa::Vector{0U}, solution).IsOk());
    ASSERT_TRUE(solver.Factorize(empty).IsOk());
    ASSERT_TRUE(tangent.FactorizeFresh(empty).IsOk());
    ASSERT_TRUE(tangent.SolveCorrection(fesa::Vector{0U}, solution).IsOk());
    ASSERT_TRUE(tangent.FactorizeFresh(empty).IsOk());
  }
  EXPECT_EQ(calls, 0);
}

TEST(PardisoSession, RestoresDefaultThreadPolicyAndPreservesCallerOverride) {
  const int previous = mkl_set_num_threads_local(0);
  struct RestoreThreads {
    int previous;
    ~RestoreThreads() { mkl_set_num_threads_local(previous); }
  } restore{previous};
  const auto matrix = MakeDenseCsr(1U, 1U, {2.0});
  for (const int caller_threads : {0, 2}) {
    mkl_set_num_threads_local(caller_threads);
    int release_calls = 0;
    fesa::internal::ScopedPardisoBackendTestHook hook{
        [&](const fesa::internal::PardisoBackendTestCall& call,
            MKL_INT& error, MKL_INT&) {
          EXPECT_EQ(mkl_get_max_threads(), caller_threads == 0 ? 1 : 2);
          if (call.phase == 22) error = -8;
          if (call.phase == -1) ++release_calls;
        }};
    {
      fesa::MklPardisoSolver solver;
      EXPECT_FALSE(solver.Factorize(matrix).IsOk());
      EXPECT_EQ(mkl_set_num_threads_local(caller_threads), caller_threads);
    }
    EXPECT_EQ(release_calls, 1);
    EXPECT_EQ(mkl_set_num_threads_local(caller_threads), caller_threads);
  }
}