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

#include <gtest/gtest.h>

#include <algorithm>
#include <limits>
#include <type_traits>
#include <utility>
#include <vector>

namespace {

using fesa::CooContribution;
using fesa::SparseMatrix;
using fesa::SparsePattern;
using fesa::Vector;

static_assert(
    !std::is_base_of_v<fesa::Matrix, SparseMatrix>,
    "SparseMatrix must own CSR storage independently of dense Matrix.");

TEST(SparseAssembly, ValidatesKnownCsrAndMultiply) {
    const SparsePattern pattern{
        {0U, 2U, 2U, 4U, 5U},
        {0U, 2U, 1U, 3U, 3U}};
    std::vector<CooContribution> contributions{
        {2U, 3U, 4.0, 2U, 0U},
        {0U, 2U, 2.0, 0U, 1U},
        {3U, 3U, 5.0, 3U, 0U},
        {0U, 0U, 1.0, 0U, 0U},
        {2U, 1U, 3.0, 1U, 0U}};

    auto result = SparseMatrix::fromCoo(
        4U, 4U, std::move(contributions), pattern);
    ASSERT_TRUE(result.hasValue());
    const SparseMatrix& matrix = result.value();

    EXPECT_EQ(matrix.rows(), 4U);
    EXPECT_EQ(matrix.columns(), 4U);
    EXPECT_EQ(matrix.rowOffsets(), pattern.rowOffsets);
    EXPECT_EQ(matrix.columnIndices(), pattern.columnIndices);
    EXPECT_EQ(matrix.values(), (std::vector<double>{1.0, 2.0, 3.0, 4.0, 5.0}));
    EXPECT_TRUE(matrix.validate().isOk());

    Vector rhs{4U};
    rhs[0U] = 1.0;
    rhs[1U] = 2.0;
    rhs[2U] = 3.0;
    rhs[3U] = 4.0;
    const Vector product = matrix.multiply(rhs);
    ASSERT_EQ(product.size(), 4U);
    EXPECT_DOUBLE_EQ(product[0U], 7.0);
    EXPECT_DOUBLE_EQ(product[1U], 0.0);
    EXPECT_DOUBLE_EQ(product[2U], 22.0);
    EXPECT_DOUBLE_EQ(product[3U], 20.0);
    EXPECT_THROW(static_cast<void>(matrix.multiply(Vector{3U})), std::invalid_argument);
}

TEST(SparseAssembly, ReducesDuplicatesInFixedTupleOrder) {
    const SparsePattern pattern{{0U, 1U}, {0U}};
    const std::vector<CooContribution> contributions{
        {0U, 0U, 1.0, 2U, 0U},
        {0U, 0U, -1.0e16, 1U, 0U},
        {0U, 0U, 1.0e16, 0U, 0U}};

    auto first = SparseMatrix::fromCoo(1U, 1U, contributions, pattern);
    ASSERT_TRUE(first.hasValue());
    ASSERT_EQ(first.value().values().size(), 1U);
    EXPECT_DOUBLE_EQ(first.value().values()[0U], 1.0);

    auto reversedContributions = contributions;
    std::reverse(reversedContributions.begin(), reversedContributions.end());
    auto second = SparseMatrix::fromCoo(
        1U, 1U, std::move(reversedContributions), pattern);
    ASSERT_TRUE(second.hasValue());
    EXPECT_EQ(second.value().rowOffsets(), first.value().rowOffsets());
    EXPECT_EQ(second.value().columnIndices(), first.value().columnIndices());
    EXPECT_EQ(second.value().values(), first.value().values());
}

TEST(SparseAssembly, RejectsInvalidIndexPatternAndShape) {
    const SparsePattern oneEntry{{0U, 1U}, {0U}};
    const auto expectFailure = [](
                                   std::size_t rows,
                                   std::size_t columns,
                                   std::vector<CooContribution> contributions,
                                   const SparsePattern& pattern) {
        auto result = SparseMatrix::fromCoo(
            rows, columns, std::move(contributions), pattern);
        EXPECT_FALSE(result.hasValue());
        if (!result.hasValue()) {
            EXPECT_FALSE(result.status().isOk());
            EXPECT_EQ(result.status().failureCategory(), fesa::FailureCategory::model);
            EXPECT_FALSE(result.status().diagnostics().empty());
        }
    };

    expectFailure(2U, 2U, {}, {{0U, 0U}, {}});
    expectFailure(1U, 1U, {}, {{1U, 1U}, {0U}});
    expectFailure(2U, 2U, {}, {{0U, 1U, 0U}, {0U}});
    expectFailure(1U, 2U, {}, {{0U, 2U}, {1U, 0U}});
    expectFailure(1U, 1U, {}, {{0U, 2U}, {0U, 0U}});
    expectFailure(1U, 1U, {}, {{0U, 1U}, {1U}});
    expectFailure(1U, 1U, {{1U, 0U, 1.0, 0U, 0U}}, oneEntry);
    expectFailure(1U, 1U, {{0U, 1U, 1.0, 0U, 0U}}, oneEntry);
    expectFailure(1U, 2U, {{0U, 1U, 1.0, 0U, 0U}}, {{0U, 1U}, {0U}});
    expectFailure(
        1U,
        1U,
        {{0U, 0U, (std::numeric_limits<double>::infinity)(), 0U, 0U}},
        oneEntry);
    expectFailure(
        1U,
        1U,
        {{0U, 0U, (std::numeric_limits<double>::quiet_NaN)(), 0U, 0U}},
        oneEntry);
    expectFailure(
        1U,
        1U,
        {{0U, 0U, (std::numeric_limits<double>::max)(), 0U, 0U},
         {0U, 0U, (std::numeric_limits<double>::max)(), 1U, 0U}},
        oneEntry);
}

TEST(SparseAssembly, PreservesExpectedStructuralZeros) {
    const SparsePattern pattern{
        {0U, 2U, 4U, 5U},
        {0U, 2U, 1U, 2U, 0U}};
    std::vector<CooContribution> contributions{
        {0U, 0U, 2.0, 0U, 0U},
        {1U, 1U, 4.0, 0U, 1U},
        {1U, 1U, -4.0, 1U, 0U}};

    auto result = SparseMatrix::fromCoo(
        3U, 3U, std::move(contributions), pattern);
    ASSERT_TRUE(result.hasValue());
    EXPECT_EQ(result.value().rowOffsets(), pattern.rowOffsets);
    EXPECT_EQ(result.value().columnIndices(), pattern.columnIndices);
    EXPECT_EQ(
        result.value().values(),
        (std::vector<double>{2.0, 0.0, 0.0, 0.0, 0.0}));
    EXPECT_EQ(
        std::count(result.value().values().begin(), result.value().values().end(), 0.0),
        4);
}

} // namespace