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

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

#include <mkl.h>

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

namespace fesa {
namespace {

Status solverFailure(
    const std::string& code,
    const std::string& identity,
    const std::string& message) {
    return Status::failure(
        FailureCategory::solver,
        {{Severity::error,
          code,
          {{}, 0U},
          "PARDISO",
          identity,
          message}});
}

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 phaseText = std::to_string(phase);
    const std::string errorText = std::to_string(error);
    return solverFailure(
        code,
        "phase=" + phaseText + ",error=" + errorText,
        "oneMKL PARDISO phase " + phaseText + " failed with error " +
            errorText + " (" + reason + ").");
}

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());
    }

    Status factorize(const SparseMatrix& matrix) {
        const MKL_INT releaseError = release();
        if (releaseError != 0) {
            return pardisoFailure(-1, releaseError);
        }

        const Status csrStatus = matrix.validate();
        if (!csrStatus.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) {
            return solverFailure(
                "solver-empty-matrix",
                "matrix-shape",
                "PARDISO factorization requires at least one equation.");
        }
        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 copyStatus = copyValidatedUpperTriangle(matrix);
        if (!copyStatus.isOk()) {
            clearOwnedArrays();
            return copyStatus;
        }

        // 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>(equationCount_), 0);
        ownsPardisoState_ = 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();
    }

    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>(equationCount_);
        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.");
            }
        }

        std::vector<double> rhsCopy(rhs.data(), rhs.data() + rhs.size());
        Vector candidate{size};
        MKL_INT phase = 33;
        MKL_INT error = 0;
        callPardiso(phase, rhsCopy.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:
    Status copyValidatedUpperTriangle(const SparseMatrix& matrix) {
        const auto& publicOffsets = matrix.rowOffsets();
        const auto& publicColumns = matrix.columnIndices();
        const auto& publicValues = matrix.values();

        for (std::size_t row = 0U; row < matrix.rows(); ++row) {
            for (std::size_t position = publicOffsets[row];
                 position < publicOffsets[row + 1U];
                 ++position) {
                const std::size_t column = publicColumns[position];
                const auto reverseBegin = publicColumns.begin() +
                    static_cast<std::ptrdiff_t>(publicOffsets[column]);
                const auto reverseEnd = publicColumns.begin() +
                    static_cast<std::ptrdiff_t>(publicOffsets[column + 1U]);
                const auto reverse =
                    std::lower_bound(reverseBegin, reverseEnd, row);
                if (reverse == reverseEnd || *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 reversePosition = static_cast<std::size_t>(
                    std::distance(publicColumns.begin(), reverse));
                const double left = publicValues[position];
                const double right = publicValues[reversePosition];
                const double scale = (std::max)({1.0, std::abs(left), std::abs(right)});
                if (std::abs(left - right) > 1.0e-12 * scale) {
                    return solverFailure(
                        "solver-matrix-not-symmetric",
                        std::to_string(row) + ":" + std::to_string(column),
                        "The full public CSR values violate the approved symmetry tolerance.");
                }
            }
        }

        equationCount_ = static_cast<MKL_INT>(matrix.rows());
        rowOffsets_.clear();
        columnIndices_.clear();
        values_.clear();
        rowOffsets_.reserve(matrix.rows() + 1U);
        rowOffsets_.push_back(0);

        for (std::size_t row = 0U; row < matrix.rows(); ++row) {
            bool hasDiagonal = false;
            for (std::size_t position = publicOffsets[row];
                 position < publicOffsets[row + 1U];
                 ++position) {
                const std::size_t column = publicColumns[position];
                if (column < row) {
                    continue;
                }
                if (!convertsToMklInt(column) ||
                    !convertsToMklInt(columnIndices_.size())) {
                    return solverFailure(
                        "solver-dimension-overflow",
                        std::to_string(row) + ":" + std::to_string(column),
                        "CSR indices exceed the oneMKL integer range.");
                }
                hasDiagonal = hasDiagonal || column == row;
                columnIndices_.push_back(static_cast<MKL_INT>(column));
                values_.push_back(publicValues[position]);
            }
            if (!hasDiagonal) {
                return solverFailure(
                    "solver-missing-diagonal",
                    std::to_string(row),
                    "Every PARDISO SPD row must retain its diagonal slot.");
            }
            if (!convertsToMklInt(columnIndices_.size())) {
                return solverFailure(
                    "solver-dimension-overflow",
                    std::to_string(row),
                    "CSR row offsets exceed the oneMKL integer range.");
            }
            rowOffsets_.push_back(
                static_cast<MKL_INT>(columnIndices_.size()));
        }
        return Status::ok();
    }

    void callPardiso(
        const MKL_INT phase,
        double* rhs,
        double* solution,
        MKL_INT& error) {
        pardiso(
            pt_.data(),
            &maxFactorizations_,
            &matrixNumber_,
            &mtype_,
            &phase,
            &equationCount_,
            values_.data(),
            rowOffsets_.data(),
            columnIndices_.data(),
            permutation_.data(),
            &rhsCount_,
            iparm_.data(),
            &messageLevel_,
            rhs,
            solution,
            &error);
    }

    MKL_INT release() noexcept {
        MKL_INT error = 0;
        if (ownsPardisoState_) {
            const MKL_INT phase = -1;
            double placeholder = 0.0;
            callPardiso(phase, &placeholder, &placeholder, error);
        }
        ownsPardisoState_ = false;
        factorized_ = false;
        pt_.fill(nullptr);
        iparm_.fill(0);
        permutation_.clear();
        clearOwnedArrays();
        return error;
    }

    void clearOwnedArrays() noexcept {
        equationCount_ = 0;
        rowOffsets_.clear();
        columnIndices_.clear();
        values_.clear();
    }

    std::array<void*, 64U> pt_{};
    std::array<MKL_INT, 64U> iparm_{};
    std::vector<MKL_INT> rowOffsets_;
    std::vector<MKL_INT> columnIndices_;
    std::vector<MKL_INT> permutation_;
    std::vector<double> values_;
    MKL_INT equationCount_{0};
    MKL_INT maxFactorizations_{1};
    MKL_INT matrixNumber_{1};
    MKL_INT mtype_{2};
    MKL_INT rhsCount_{1};
    MKL_INT messageLevel_{0};
    bool ownsPardisoState_{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