#include "fesa/math/sparse_matrix.h"

#include <gtest/gtest.h>

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

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

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 reversed_contributions = contributions;
  std::reverse(reversed_contributions.begin(), reversed_contributions.end());
  auto second =
      SparseMatrix::FromCoo(1U, 1U, std::move(reversed_contributions), 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 one_entry{{0U, 1U}, {0U}};
  const auto expect_failure = [](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.GetStatus().IsOk());
      EXPECT_EQ(result.GetStatus().Category(), fesa::FailureCategory::kModel);
      EXPECT_FALSE(result.GetStatus().Diagnostics().empty());
    }
  };

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

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