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

#include <mkl.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "fesa/math/sparse_matrix.h"
#include "fesa/math/vector.h"

namespace fesa {
namespace {

/// @brief Builds a structured linear-solver failure.
Status SolverFailure(const std::string& code, const std::string& identity,
                     const std::string& message) {
  return Status::Failure(
      FailureCategory::kSolver,
      {{Severity::kError, code, {{}, 0U}, "PARDISO", identity, message}});
}

/// @brief Translates a PARDISO phase error into the stable solver taxonomy.
Status PardisoFailure(const MKL_INT phase, const MKL_INT error) {
  std::string code;
  std::string reason;
  switch (error) {
    case -4:
      code = "pardiso-zero-or-negative-pivot";
      reason = "zero or negative pivot";
      break;
    case -7:
      code = "pardiso-singular-diagonal";
      reason = "singular diagonal";
      break;
    case -8:
      code = "pardiso-integer-overflow";
      reason = "32-bit backend integer overflow";
      break;
    case 21:
    case 22:
    case 23:
    case 24:
      code = "pardiso-invalid-csr";
      reason = "matrix checker rejected the CSR indices";
      break;
    default:
      code = phase == 11   ? "pardiso-analysis-failed"
             : phase == 22 ? "pardiso-factorization-failed"
             : phase == 33 ? "pardiso-solve-failed"
                           : "pardiso-release-failed";
      reason = "backend error";
      break;
  }

  const std::string phase_text = std::to_string(phase);
  const std::string error_text = std::to_string(error);
  return SolverFailure(code, "phase=" + phase_text + ",error=" + error_text,
                       "oneMKL PARDISO phase " + phase_text +
                           " failed with error " + error_text + " (" + reason +
                           ").");
}

/// @brief Reports whether a public size fits the private MKL integer type.
bool ConvertsToMklInt(const std::size_t value) {
  return value <=
         static_cast<std::size_t>((std::numeric_limits<MKL_INT>::max)());
}

}  // namespace

class MklPardisoSolver::Impl {
 public:
  Impl() = default;

  ~Impl() { static_cast<void>(Release()); }

  /// @brief Validates and factorizes one public CSR matrix.
  Status Factorize(const SparseMatrix& matrix) {
    const MKL_INT release_error = Release();
    if (release_error != 0) {
      return PardisoFailure(-1, release_error);
    }

    const Status csr_status = matrix.Validate();
    if (!csr_status.IsOk()) {
      return SolverFailure("solver-invalid-csr", "public-csr",
                           "The public sparse matrix failed CSR validation.");
    }
    if (matrix.Rows() != matrix.Columns()) {
      return SolverFailure("solver-matrix-not-square", "matrix-shape",
                           "PARDISO factorization requires a square matrix.");
    }
    if (matrix.Rows() == 0U) {
      // A fully constrained model has no free equations. Preserve the
      // observable factorize/solve lifecycle without creating backend
      // state or calling PARDISO with its invalid n=0 input.
      factorized_ = true;
      return Status::Ok();
    }
    if (!ConvertsToMklInt(matrix.Rows()) ||
        !ConvertsToMklInt(matrix.Values().size())) {
      return SolverFailure(
          "solver-dimension-overflow", "matrix-shape",
          "Sparse matrix dimensions exceed the oneMKL integer range.");
    }

    const Status copy_status = CopyValidatedUpperTriangle(matrix);
    if (!copy_status.IsOk()) {
      ClearOwnedArrays();
      return copy_status;
    }

    // PARDISO owns internal memory behind pt after phase 11. Initialize
    // once per factorization and retain it until refactorization/destruction.
    pt_.fill(nullptr);
    iparm_.fill(0);
    pardisoinit(pt_.data(), &mtype_, iparm_.data());
    iparm_[26] = 1;  // Validate sorted CSR integer arrays.
    iparm_[34] = 1;  // Consume the project's native zero-based CSR.
    permutation_.assign(static_cast<std::size_t>(equation_count_), 0);
    owns_pardiso_state_ = true;

    MKL_INT phase = 11;
    MKL_INT error = 0;
    CallPardiso(phase, nullptr, nullptr, error);
    if (error != 0) {
      const Status failure = PardisoFailure(phase, error);
      static_cast<void>(Release());
      return failure;
    }

    phase = 22;
    error = 0;
    CallPardiso(phase, nullptr, nullptr, error);
    if (error != 0) {
      const Status failure = PardisoFailure(phase, error);
      static_cast<void>(Release());
      return failure;
    }

    factorized_ = true;
    return Status::Ok();
  }

  /// @brief Substitutes one RHS while preserving solution on failure.
  Status Solve(const Vector& rhs, Vector& solution) {
    if (!factorized_) {
      return SolverFailure(
          "solver-not-factorized", "factorization-state",
          "Substitution requires a successful retained factorization.");
    }
    const std::size_t size = static_cast<std::size_t>(equation_count_);
    if (rhs.Size() != size || solution.Size() != size) {
      return SolverFailure(
          "solver-vector-dimension-mismatch", "rhs-or-solution",
          "RHS and solution dimensions must match the factorized matrix.");
    }
    for (std::size_t index = 0U; index < rhs.Size(); ++index) {
      if (!std::isfinite(rhs[index])) {
        return SolverFailure("nonfinite-solver-rhs", std::to_string(index),
                             "PARDISO RHS values must be finite.");
      }
    }
    if (size == 0U) {
      return Status::Ok();
    }

    std::vector<double> rhs_copy(rhs.Data(), rhs.Data() + rhs.Size());
    Vector candidate{size};
    MKL_INT phase = 33;
    MKL_INT error = 0;
    CallPardiso(phase, rhs_copy.data(), candidate.Data(), error);
    if (error != 0) {
      return PardisoFailure(phase, error);
    }
    for (std::size_t index = 0U; index < candidate.Size(); ++index) {
      if (!std::isfinite(candidate[index])) {
        return SolverFailure(
            "nonfinite-solver-solution", std::to_string(index),
            "PARDISO substitution produced a nonfinite solution.");
      }
    }

    solution = std::move(candidate);
    return Status::Ok();
  }

 private:
  /// @brief Validates symmetry and copies the upper triangle for PARDISO.
  Status CopyValidatedUpperTriangle(const SparseMatrix& matrix) {
    const auto& public_offsets = matrix.RowOffsets();
    const auto& public_columns = matrix.ColumnIndices();
    const auto& public_values = matrix.Values();

    double matrix_scale = 0.0;
    for (const double value : public_values) {
      matrix_scale = (std::max)(matrix_scale, std::abs(value));
    }

    for (std::size_t row = 0U; row < matrix.Rows(); ++row) {
      for (std::size_t position = public_offsets[row];
           position < public_offsets[row + 1U]; ++position) {
        const std::size_t column = public_columns[position];
        const auto reverse_begin =
            public_columns.begin() +
            static_cast<std::ptrdiff_t>(public_offsets[column]);
        const auto reverse_end =
            public_columns.begin() +
            static_cast<std::ptrdiff_t>(public_offsets[column + 1U]);
        const auto reverse = std::lower_bound(reverse_begin, reverse_end, row);
        if (reverse == reverse_end || *reverse != row) {
          return SolverFailure(
              "solver-matrix-not-symmetric",
              std::to_string(row) + ":" + std::to_string(column),
              "The full public CSR must contain both symmetric entries.");
        }

        const std::size_t reverse_position = static_cast<std::size_t>(
            std::distance(public_columns.begin(), reverse));
        const double left = public_values[position];
        const double right = public_values[reverse_position];
        const double difference = std::abs(left - right);
        // The approved symmetry test is normalized by the matrix's
        // actual nonzero scale, without an absolute unit-size floor.
        const bool is_symmetric = matrix_scale == 0.0
                                      ? difference == 0.0
                                      : difference <= 1.0e-12 * matrix_scale;
        if (!is_symmetric) {
          return SolverFailure(
              "solver-matrix-not-symmetric",
              std::to_string(row) + ":" + std::to_string(column),
              "The full public CSR values violate the approved symmetry "
              "tolerance.");
        }
      }
    }

    equation_count_ = static_cast<MKL_INT>(matrix.Rows());
    row_offsets_.clear();
    column_indices_.clear();
    values_.clear();
    row_offsets_.reserve(matrix.Rows() + 1U);
    row_offsets_.push_back(0);

    for (std::size_t row = 0U; row < matrix.Rows(); ++row) {
      bool has_diagonal = false;
      for (std::size_t position = public_offsets[row];
           position < public_offsets[row + 1U]; ++position) {
        const std::size_t column = public_columns[position];
        if (column < row) {
          continue;
        }
        if (!ConvertsToMklInt(column) ||
            !ConvertsToMklInt(column_indices_.size())) {
          return SolverFailure(
              "solver-dimension-overflow",
              std::to_string(row) + ":" + std::to_string(column),
              "CSR indices exceed the oneMKL integer range.");
        }
        has_diagonal = has_diagonal || column == row;
        column_indices_.push_back(static_cast<MKL_INT>(column));
        values_.push_back(public_values[position]);
      }
      if (!has_diagonal) {
        return SolverFailure(
            "solver-missing-diagonal", std::to_string(row),
            "Every PARDISO SPD row must retain its diagonal slot.");
      }
      if (!ConvertsToMklInt(column_indices_.size())) {
        return SolverFailure(
            "solver-dimension-overflow", std::to_string(row),
            "CSR row offsets exceed the oneMKL integer range.");
      }
      row_offsets_.push_back(static_cast<MKL_INT>(column_indices_.size()));
    }
    return Status::Ok();
  }

  /// @brief Calls PARDISO with retained private arrays and phase state.
  void CallPardiso(const MKL_INT phase, double* rhs, double* solution,
                   MKL_INT& error) {
    pardiso(pt_.data(), &max_factorizations_, &matrix_number_, &mtype_, &phase,
            &equation_count_, values_.data(), row_offsets_.data(),
            column_indices_.data(), permutation_.data(), &rhs_count_,
            iparm_.data(), &message_level_, rhs, solution, &error);
  }

  /// @brief Releases backend state and resets every retained array.
  MKL_INT Release() noexcept {
    MKL_INT error = 0;
    if (owns_pardiso_state_) {
      const MKL_INT phase = -1;
      double placeholder = 0.0;
      CallPardiso(phase, &placeholder, &placeholder, error);
    }
    owns_pardiso_state_ = false;
    factorized_ = false;
    pt_.fill(nullptr);
    iparm_.fill(0);
    permutation_.clear();
    ClearOwnedArrays();
    return error;
  }

  /// @brief Clears FESA-owned CSR arrays without touching backend state.
  void ClearOwnedArrays() noexcept {
    equation_count_ = 0;
    row_offsets_.clear();
    column_indices_.clear();
    values_.clear();
  }

  std::array<void*, 64U> pt_{};
  std::array<MKL_INT, 64U> iparm_{};
  std::vector<MKL_INT> row_offsets_;
  std::vector<MKL_INT> column_indices_;
  std::vector<MKL_INT> permutation_;
  std::vector<double> values_;
  MKL_INT equation_count_{0};
  MKL_INT max_factorizations_{1};
  MKL_INT matrix_number_{1};
  MKL_INT mtype_{2};
  MKL_INT rhs_count_{1};
  MKL_INT message_level_{0};
  bool owns_pardiso_state_{false};
  bool factorized_{false};
};

MklPardisoSolver::MklPardisoSolver() : impl_{std::make_unique<Impl>()} {}

MklPardisoSolver::~MklPardisoSolver() = default;

Status MklPardisoSolver::Factorize(const SparseMatrix& matrix) {
  return impl_->Factorize(matrix);
}

Status MklPardisoSolver::Solve(const Vector& rhs, Vector& solution) const {
  return impl_->Solve(rhs, solution);
}

}  // namespace fesa