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

#include "fesa/fem/dof_manager.hpp"
#include "fesa/math/sparse_matrix.hpp"

#include <gtest/gtest.h>

#include <cmath>
#include <initializer_list>
#include <limits>
#include <string>
#include <utility>
#include <vector>

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.rowOffsets.reserve(size + 1U);
    pattern.rowOffsets.push_back(0U);
    for (std::size_t row = 0U; row < size; ++row) {
        for (std::size_t column = 0U; column < size; ++column) {
            pattern.columnIndices.push_back(column);
            contributions.push_back({
                row,
                column,
                values[row * size + column],
                row,
                column});
        }
        pattern.rowOffsets.push_back(pattern.columnIndices.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.failureCategory(), fesa::FailureCategory::solver);
    ASSERT_EQ(status.diagnostics().size(), 1U);
    EXPECT_EQ(status.diagnostics()[0U].severity, fesa::Severity::error);
    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 concreteSolver;
    fesa::LinearSolver& solver = concreteSolver;
    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 expectedFirst = makeVector({1.0, 2.0});
    const auto expectedSecond = makeVector({-2.0, 0.5});
    const auto rhsFirst = matrix.multiply(expectedFirst);
    const auto rhsSecond = matrix.multiply(expectedSecond);

    fesa::MklPardisoSolver solver;
    ASSERT_TRUE(solver.factorize(matrix).isOk());
    fesa::Vector first{2U};
    fesa::Vector second{2U};
    ASSERT_TRUE(solver.solve(rhsFirst, first).isOk());
    ASSERT_TRUE(solver.solve(rhsSecond, second).isOk());

    EXPECT_LE(relativeError(first, expectedFirst), 1.0e-9);
    EXPECT_LE(relativeError(second, expectedSecond), 1.0e-9);
    EXPECT_LE(normalizedResidual(matrix, first, rhsFirst), 1.0e-10);
    EXPECT_LE(normalizedResidual(matrix, second, rhsSecond), 1.0e-10);
}

TEST(MklPardisoSolver, RefactorizesWithoutLeakingState) {
    const auto firstMatrix = makeDenseCsr(2U, {4.0, 1.0, 1.0, 3.0});
    const auto secondMatrix = makeDenseCsr(2U, {2.0, 0.0, 0.0, 5.0});
    const auto firstExpected = makeVector({1.0, 2.0});
    const auto secondExpected = makeVector({-3.0, 4.0});

    fesa::MklPardisoSolver solver;
    ASSERT_TRUE(solver.factorize(firstMatrix).isOk());
    fesa::Vector firstSolution{2U};
    ASSERT_TRUE(
        solver.solve(firstMatrix.multiply(firstExpected), firstSolution).isOk());
    EXPECT_LE(relativeError(firstSolution, firstExpected), 1.0e-9);

    ASSERT_TRUE(solver.factorize(secondMatrix).isOk());
    fesa::Vector secondSolution{2U};
    const auto secondRhs = secondMatrix.multiply(secondExpected);
    ASSERT_TRUE(solver.solve(secondRhs, secondSolution).isOk());
    EXPECT_LE(relativeError(secondSolution, secondExpected), 1.0e-9);
    EXPECT_LE(
        normalizedResidual(secondMatrix, secondSolution, secondRhs), 1.0e-10);
}

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

    const auto singular = makeDenseCsr(2U, {1.0, 1.0, 1.0, 1.0});
    const auto singularStatus = solver.factorize(singular);
    expectStructuredSolverFailure(singularStatus);
    EXPECT_TRUE(
        singularStatus.diagnostics()[0U].code ==
            "pardiso-zero-or-negative-pivot" ||
        singularStatus.diagnostics()[0U].code ==
            "pardiso-singular-diagonal");
    EXPECT_NE(
        singularStatus.diagnostics()[0U].entityIdentity.find(
            "phase=22,error="),
        std::string::npos);

    const auto indefinite = makeDenseCsr(2U, {1.0, 2.0, 2.0, 1.0});
    const auto indefiniteStatus = solver.factorize(indefinite);
    expectStructuredSolverFailure(indefiniteStatus);
    EXPECT_TRUE(
        indefiniteStatus.diagnostics()[0U].code ==
            "pardiso-zero-or-negative-pivot" ||
        indefiniteStatus.diagnostics()[0U].code ==
            "pardiso-singular-diagonal");
    EXPECT_NE(
        indefiniteStatus.diagnostics()[0U].entityIdentity.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 rhsStatus = solver.solve(rhs, solution);
    expectStructuredSolverFailure(rhsStatus);
    EXPECT_EQ(rhsStatus.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 nonfiniteMatrix = fesa::SparseMatrix::fromCoo(
        1U,
        1U,
        {{0U,
          0U,
          (std::numeric_limits<double>::quiet_NaN)(),
          0U,
          0U}},
        pattern);
    EXPECT_FALSE(nonfiniteMatrix.hasValue());
    EXPECT_EQ(
        nonfiniteMatrix.status().diagnostics()[0U].code,
        "nonfinite-sparse-value");
}

TEST(MklPardisoSolver, ConditioningSweepPassesResolvedCasesAndFailsUnresolvedCasesExplicitly) {
    const std::vector<double> commonScales{1.0e-12, 1.0, 1.0e12};
    for (const double scale : commonScales) {
        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 resolvedRatio = 1.0e-8;
    const auto resolvedMatrix = makeDenseCsr(
        2U, {1.0, 0.0, 0.0, resolvedRatio});
    const auto resolvedExpected = makeVector({0.5, -2.0});
    const auto resolvedRhs = resolvedMatrix.multiply(resolvedExpected);
    fesa::MklPardisoSolver resolvedSolver;
    ASSERT_TRUE(resolvedSolver.factorize(resolvedMatrix).isOk());
    fesa::Vector resolvedSolution{2U};
    ASSERT_TRUE(resolvedSolver.solve(resolvedRhs, resolvedSolution).isOk());
    EXPECT_LE(
        normalizedResidual(resolvedMatrix, resolvedSolution, resolvedRhs),
        1.0e-10);
    EXPECT_LE(relativeError(resolvedSolution, resolvedExpected), 1.0e-9);

    const std::vector<double> unresolvedCandidates{1.0e-16, 1.0e-300};
    for (const double ratio : unresolvedCandidates) {
        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 factorStatus = solver.factorize(matrix);
        if (!factorStatus.isOk()) {
            expectStructuredSolverFailure(factorStatus);
            continue;
        }

        fesa::Vector solution{2U};
        const auto solveStatus = solver.solve(rhs, solution);
        if (!solveStatus.isOk()) {
            expectStructuredSolverFailure(solveStatus);
            continue;
        }

        EXPECT_LE(normalizedResidual(matrix, solution, rhs), 1.0e-10);
        EXPECT_LE(relativeError(solution, expected), 1.0e-9);
    }
}