#include <fesa/solvers/linear/pardiso_linear_solver.hpp>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <string>
#include <string_view>
#include <type_traits>
#include <utility>
#include <vector>

#include <mkl.h>

namespace fesa {
namespace {

static_assert(
    std::is_same_v<MKL_INT, std::int32_t>,
    "FESA requires the oneMKL LP64 interface.");

Diagnostic solver_error(std::string code, std::string message) {
    return {
        DiagnosticStage::solver,
        Severity::error,
        std::move(code),
        std::move(message),
        std::nullopt,
    };
}

std::optional<std::string> validate_matrix(
    const SymmetricCsr& matrix) {
    if (matrix.order == 0) {
        return "PARDISO requires a nonempty reduced system.";
    }
    if (matrix.order >
        static_cast<std::size_t>(
            std::numeric_limits<MKL_INT>::max())) {
        return "Matrix order exceeds the oneMKL LP64 index range.";
    }
    if (matrix.row_offsets.size() != matrix.order + 1 ||
        matrix.row_offsets.front() != 0) {
        return "CSR row offsets must contain order + 1 entries "
               "starting at zero.";
    }
    if (matrix.column_indices.size() != matrix.values.size()) {
        return "CSR column and value counts must match.";
    }

    MKL_INT previous_offset = 0;
    for (const MKL_INT offset : matrix.row_offsets) {
        if (offset < previous_offset || offset < 0 ||
            static_cast<std::size_t>(offset) >
                matrix.column_indices.size()) {
            return "CSR row offsets must be nondecreasing and in range.";
        }
        previous_offset = offset;
    }
    if (static_cast<std::size_t>(matrix.row_offsets.back()) !=
        matrix.column_indices.size()) {
        return "The final CSR row offset must equal the entry count.";
    }

    for (std::size_t row = 0; row < matrix.order; ++row) {
        MKL_INT previous_column = -1;
        bool has_diagonal = false;
        const std::size_t begin =
            static_cast<std::size_t>(matrix.row_offsets[row]);
        const std::size_t end =
            static_cast<std::size_t>(matrix.row_offsets[row + 1]);
        for (std::size_t entry = begin; entry < end; ++entry) {
            const MKL_INT column = matrix.column_indices[entry];
            if (column < static_cast<MKL_INT>(row) ||
                column >= static_cast<MKL_INT>(matrix.order) ||
                column <= previous_column) {
                return "CSR rows must contain sorted unique "
                       "upper-triangle columns.";
            }
            if (!std::isfinite(matrix.values[entry])) {
                return "CSR values must be finite.";
            }
            has_diagonal =
                has_diagonal || column == static_cast<MKL_INT>(row);
            previous_column = column;
        }
        if (!has_diagonal) {
            return "Every CSR row must contain its diagonal entry.";
        }
    }
    return std::nullopt;
}

class PardisoSession final {
public:
    PardisoSession() {
        pardisoinit(handles_.data(), &matrix_type_, parameters_.data());
        parameters_[26] = 1;
        parameters_[34] = 1;
    }

    PardisoSession(const PardisoSession&) = delete;
    PardisoSession& operator=(const PardisoSession&) = delete;

    ~PardisoSession() noexcept {
        static_cast<void>(release());
    }

    MKL_INT release() noexcept {
        if (!active_) {
            return 0;
        }
        active_ = false;

        constexpr MKL_INT release_all = -1;
        MKL_INT error = 0;
        pardiso(
            handles_.data(),
            &max_factorizations_,
            &matrix_number_,
            &matrix_type_,
            &release_all,
            &order_,
            matrix_->values.data(),
            matrix_->row_offsets.data(),
            matrix_->column_indices.data(),
            permutation_.data(),
            &right_hand_side_count_,
            parameters_.data(),
            &message_level_,
            right_hand_side_->data(),
            solution_->data(),
            &error);
        return error;
    }

    MKL_INT execute(
        const MKL_INT phase,
        const SymmetricCsr& matrix,
        std::vector<double>& right_hand_side,
        std::vector<double>& solution) {
        order_ = static_cast<MKL_INT>(matrix.order);
        matrix_ = &matrix;
        right_hand_side_ = &right_hand_side;
        solution_ = &solution;
        permutation_.resize(matrix.order);
        active_ = true;

        MKL_INT error = 0;
        pardiso(
            handles_.data(),
            &max_factorizations_,
            &matrix_number_,
            &matrix_type_,
            &phase,
            &order_,
            matrix.values.data(),
            matrix.row_offsets.data(),
            matrix.column_indices.data(),
            permutation_.data(),
            &right_hand_side_count_,
            parameters_.data(),
            &message_level_,
            right_hand_side.data(),
            solution.data(),
            &error);
        return error;
    }

private:
    std::array<void*, 64> handles_{};
    std::array<MKL_INT, 64> parameters_{};
    std::vector<MKL_INT> permutation_;
    const SymmetricCsr* matrix_{};
    std::vector<double>* right_hand_side_{};
    std::vector<double>* solution_{};
    MKL_INT order_{};
    MKL_INT max_factorizations_{1};
    MKL_INT matrix_number_{1};
    MKL_INT matrix_type_{2};
    MKL_INT right_hand_side_count_{1};
    MKL_INT message_level_{};
    bool active_{};
};

double relative_residual(
    const SymmetricCsr& matrix,
    const std::span<const double> rhs,
    const std::span<const double> solution) {
    std::vector<double> residual(rhs.begin(), rhs.end());
    for (double& value : residual) {
        value = -value;
    }

    for (std::size_t row = 0; row < matrix.order; ++row) {
        const std::size_t begin =
            static_cast<std::size_t>(matrix.row_offsets[row]);
        const std::size_t end =
            static_cast<std::size_t>(matrix.row_offsets[row + 1]);
        for (std::size_t entry = begin; entry < end; ++entry) {
            const std::size_t column =
                static_cast<std::size_t>(matrix.column_indices[entry]);
            const double value = matrix.values[entry];
            residual[row] += value * solution[column];
            if (column != row) {
                residual[column] += value * solution[row];
            }
        }
    }

    double residual_squared = 0.0;
    double rhs_squared = 0.0;
    for (std::size_t index = 0; index < rhs.size(); ++index) {
        residual_squared += residual[index] * residual[index];
        rhs_squared += rhs[index] * rhs[index];
    }
    const double residual_norm = std::sqrt(residual_squared);
    const double rhs_norm = std::sqrt(rhs_squared);
    return rhs_norm == 0.0 ? residual_norm : residual_norm / rhs_norm;
}

std::string pardiso_failure(
    const std::string_view phase,
    const MKL_INT error) {
    return "PARDISO " + std::string{phase} +
           " failed with error " + std::to_string(error) + ".";
}

}  // namespace

PardisoLinearSolver::PardisoLinearSolver() = default;

PardisoLinearSolver::~PardisoLinearSolver() = default;

LinearSolveResult PardisoLinearSolver::solve(
    const SymmetricCsr& matrix,
    const std::span<const double> rhs) {
    LinearSolveResult result;
    if (const auto error = validate_matrix(matrix);
        error.has_value()) {
        result.diagnostics.push_back(
            solver_error("solver.invalid_csr", *error));
        return result;
    }
    if (rhs.size() != matrix.order) {
        result.diagnostics.push_back(solver_error(
            "solver.dimension_mismatch",
            "Right-hand side size must equal the matrix order."));
        return result;
    }
    if (!std::ranges::all_of(rhs, [](const double value) {
            return std::isfinite(value);
        })) {
        result.diagnostics.push_back(solver_error(
            "solver.invalid_rhs",
            "Right-hand side values must be finite."));
        return result;
    }

    std::vector<double> right_hand_side(rhs.begin(), rhs.end());
    std::vector<double> solution(matrix.order, 0.0);
    {
        PardisoSession session;

        constexpr MKL_INT analyze = 11;
        MKL_INT error = session.execute(
            analyze, matrix, right_hand_side, solution);
        if (error != 0) {
            result.diagnostics.push_back(solver_error(
                "solver.analysis_failed",
                pardiso_failure("analysis", error)));
        }

        if (error == 0) {
            constexpr MKL_INT factorize = 22;
            error = session.execute(
                factorize, matrix, right_hand_side, solution);
            if (error != 0) {
                result.diagnostics.push_back(solver_error(
                    "solver.factorization_failed",
                    pardiso_failure("factorization", error)));
            }
        }

        if (error == 0) {
            constexpr MKL_INT solve_system = 33;
            error = session.execute(
                solve_system, matrix, right_hand_side, solution);
            if (error != 0) {
                result.diagnostics.push_back(solver_error(
                    "solver.solve_failed",
                    pardiso_failure("solve", error)));
            }
        }

        if (const MKL_INT release_error = session.release();
            release_error != 0) {
            result.diagnostics.push_back(solver_error(
                "solver.release_failed",
                pardiso_failure("release", release_error)));
        }
    }
    if (!result.diagnostics.empty()) {
        return result;
    }
    if (!std::ranges::all_of(solution, [](const double value) {
            return std::isfinite(value);
        })) {
        result.diagnostics.push_back(solver_error(
            "solver.nonfinite_solution",
            "PARDISO produced a nonfinite solution."));
        return result;
    }

    result.relative_residual =
        relative_residual(matrix, rhs, solution);
    if (!std::isfinite(result.relative_residual)) {
        result.diagnostics.push_back(solver_error(
            "solver.nonfinite_residual",
            "The independently computed relative residual is nonfinite."));
        return result;
    }
    result.solution = std::move(solution);
    return result;
}

}  // namespace fesa
