#include "fesa/assembly/sparse_assembler.h"

#include <gtest/gtest.h>

#include <algorithm>
#include <array>
#include <cstring>
#include <filesystem>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/elements/mitc4_shell.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/model/domain.h"

namespace {

fesa::ModelDefinition MakeDefinition() {
  const std::filesystem::path source{"models/sparse-assembly.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "fnv1a64:0123456789abcdef";
  definition.nodes = {{{"Beam-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 10U}},
                      {{"Beam-1", 2, "2"}, {2.0, 0.0, 0.0}, {source, 11U}},
                      {{"Beam-1", 3, "3"}, {5.0, 0.0, 0.0}, {source, 12U}}};
  definition.materials = {{"Material", 120.0, 0.25, {source, 20U}}};
  definition.sections = {{"Section",
                          2.0,
                          1.5,
                          0.0,
                          0.75,
                          0.5,
                          {0.0, 1.0, 0.0},
                          {},
                          {source, 30U}}};
  definition.elements = {
      {{"Beam-1", 10, "10"}, {0U, 1U}, 0U, 0U, {source, 40U}},
      {{"Beam-1", 20, "20"}, {1U, 2U}, 0U, 0U, {source, 41U}}};
  definition.steps = {{"Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0, {source, 50U}}};
  return definition;
}

fesa::ModelDefinition MakeShellDefinition(
    const fesa::ShellSourceElementType source_type,
    const bool two_elements = false) {
  const std::filesystem::path source{"models/shell-sparse-assembly.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "fnv1a64:fedcba9876543210";
  if (two_elements) {
    definition.nodes = {{{"Shell-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 10U}},
                        {{"Shell-1", 2, "2"}, {1.0, 0.0, 0.0}, {source, 11U}},
                        {{"Shell-1", 3, "3"}, {2.0, 0.0, 0.0}, {source, 12U}},
                        {{"Shell-1", 4, "4"}, {0.0, 1.0, 0.0}, {source, 13U}},
                        {{"Shell-1", 5, "5"}, {1.0, 1.0, 0.0}, {source, 14U}},
                        {{"Shell-1", 6, "6"}, {2.0, 1.0, 0.0}, {source, 15U}}};
    for (std::size_t node = 0U; node < definition.nodes.size(); ++node) {
      definition.shell_node_initial_frames.push_back(
          {static_cast<fesa::EntityIndex>(node),
           {0.0, 0.0, 1.0},
           {1.0, 0.0, 0.0},
           {0.0, 1.0, 0.0}});
    }
    definition.shell_elements = {{{"Shell-1", 10, "10"},
                                  source_type,
                                  {0U, 1U, 4U, 3U},
                                  0U,
                                  0U,
                                  {source, 40U}},
                                 {{"Shell-1", 20, "20"},
                                  source_type,
                                  {1U, 2U, 5U, 4U},
                                  0U,
                                  0U,
                                  {source, 41U}}};
  } else {
    // A YZ-plane fixture catches any accidental global-Z director assumption.
    definition.nodes = {{{"Shell-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 10U}},
                        {{"Shell-1", 2, "2"}, {0.0, 1.0, 0.0}, {source, 11U}},
                        {{"Shell-1", 3, "3"}, {0.0, 1.0, 1.0}, {source, 12U}},
                        {{"Shell-1", 4, "4"}, {0.0, 0.0, 1.0}, {source, 13U}}};
    for (std::size_t node = 0U; node < definition.nodes.size(); ++node) {
      definition.shell_node_initial_frames.push_back(
          {static_cast<fesa::EntityIndex>(node),
           {1.0, 0.0, 0.0},
           {0.0, 1.0, 0.0},
           {0.0, 0.0, 1.0}});
    }
    definition.shell_elements = {{{"Shell-1", 10, "10"},
                                  source_type,
                                  {0U, 1U, 2U, 3U},
                                  0U,
                                  0U,
                                  {source, 40U}}};
  }
  definition.materials = {{"Material", 120.0, 0.25, {source, 20U}}};
  definition.shell_sections = {{"ShellSection", 0.2, 0U, {source, 30U}}};
  definition.steps = {{"Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0, {source, 50U}}};
  return definition;
}

fesa::Result<fesa::Mitc4Stiffness> DirectShellStiffness(
    const fesa::Domain& domain, const fesa::EntityIndex element_index) {
  const auto& definition = domain.ShellElements().at(element_index);
  std::array<const fesa::Node*, 4> nodes{};
  std::array<std::array<double, 3>, 4> directors{};
  for (std::size_t node = 0U; node < definition.node_indices.size(); ++node) {
    const fesa::EntityIndex node_index = definition.node_indices[node];
    nodes[node] = &domain.Nodes().at(node_index);
    directors[node] = domain.ShellNodeInitialFrames().at(node_index).director;
  }
  auto shell = fesa::Mitc4Shell::Create(
      nodes, directors, domain.ShellSections().at(definition.section_index),
      domain.Materials().at(definition.material_index));
  if (!shell.HasValue()) {
    return fesa::Result<fesa::Mitc4Stiffness>::Failure(shell.GetStatus());
  }
  return shell.Value().Stiffness();
}

fesa::Result<fesa::SparseMatrix> AssembleShell(
    const fesa::ShellSourceElementType source_type,
    const fesa::ParallelFor& parallel_for, const bool two_elements = false) {
  auto domain =
      fesa::Domain::Create(MakeShellDefinition(source_type, two_elements));
  if (!domain.HasValue()) {
    return fesa::Result<fesa::SparseMatrix>::Failure(domain.GetStatus());
  }
  auto model = fesa::AnalysisModel::Create(domain.Value());
  if (!model.HasValue()) {
    return fesa::Result<fesa::SparseMatrix>::Failure(model.GetStatus());
  }
  auto dofs = fesa::DofManager::Create(model.Value());
  if (!dofs.HasValue()) {
    return fesa::Result<fesa::SparseMatrix>::Failure(dofs.GetStatus());
  }
  return fesa::SparseAssembler::AssembleStiffness(model.Value(), dofs.Value(),
                                                  parallel_for);
}

template <class T>
bool ByteIdentical(const std::vector<T>& left, const std::vector<T>& right) {
  return left.size() == right.size() &&
         (left.empty() ||
          std::memcmp(left.data(), right.data(), left.size() * sizeof(T)) == 0);
}

double Entry(const fesa::SparseMatrix& matrix, const std::size_t row,
             const std::size_t column) {
  const auto begin = matrix.ColumnIndices().begin() + matrix.RowOffsets()[row];
  const auto end =
      matrix.ColumnIndices().begin() + matrix.RowOffsets()[row + 1U];
  const auto found = std::lower_bound(begin, end, column);
  if (found == end || *found != column) {
    return 0.0;
  }
  return matrix.Values()[static_cast<std::size_t>(
      std::distance(matrix.ColumnIndices().begin(), found))];
}

class ReverseParallelFor final : public fesa::ParallelFor {
 public:
  void Execute(const std::size_t count,
               const std::function<void(std::size_t)>& body) const override {
    ++calls_;
    observed_count_ = count;
    for (std::size_t index = count; index > 0U; --index) {
      body(index - 1U);
    }
  }

  std::size_t Calls() const noexcept { return calls_; }

  std::size_t ObservedCount() const noexcept { return observed_count_; }

 private:
  mutable std::size_t calls_{0U};
  mutable std::size_t observed_count_{0U};
};

void ExpectByteIdentical(const fesa::SparseMatrix& actual,
                         const fesa::SparseMatrix& expected) {
  EXPECT_TRUE(ByteIdentical(actual.RowOffsets(), expected.RowOffsets()));
  EXPECT_TRUE(ByteIdentical(actual.ColumnIndices(), expected.ColumnIndices()));
  EXPECT_TRUE(ByteIdentical(actual.Values(), expected.Values()));
}

TEST(SparseAssembly, SerialTbbAndRepeatedRunsAreByteIdentical) {
  auto domain_result = fesa::Domain::Create(MakeDefinition());
  ASSERT_TRUE(domain_result.HasValue());
  auto model_result = fesa::AnalysisModel::Create(domain_result.Value());
  ASSERT_TRUE(model_result.HasValue());
  auto dofs_result = fesa::DofManager::Create(model_result.Value());
  ASSERT_TRUE(dofs_result.HasValue());

  fesa::SerialParallelFor serial_executor;
  fesa::TbbParallelFor tbb_executor;
  ReverseParallelFor reverse_executor;
  auto serial = fesa::SparseAssembler::AssembleStiffness(
      model_result.Value(), dofs_result.Value(), serial_executor);
  auto tbb = fesa::SparseAssembler::AssembleStiffness(
      model_result.Value(), dofs_result.Value(), tbb_executor);
  auto reversed = fesa::SparseAssembler::AssembleStiffness(
      model_result.Value(), dofs_result.Value(), reverse_executor);
  ASSERT_TRUE(serial.HasValue());
  ASSERT_TRUE(tbb.HasValue());
  ASSERT_TRUE(reversed.HasValue());

  EXPECT_EQ(reverse_executor.Calls(), 1U);
  EXPECT_EQ(reverse_executor.ObservedCount(), 2U);
  EXPECT_EQ(serial.Value().Rows(), 18U);
  EXPECT_EQ(serial.Value().Columns(), 18U);
  EXPECT_EQ(serial.Value().RowOffsets(),
            dofs_result.Value().GetSparsePattern().row_offsets);
  EXPECT_EQ(serial.Value().ColumnIndices(),
            dofs_result.Value().GetSparsePattern().column_indices);
  EXPECT_TRUE(serial.Value().Validate().IsOk());
  ExpectByteIdentical(tbb.Value(), serial.Value());
  ExpectByteIdentical(reversed.Value(), serial.Value());

  for (std::size_t repetition = 0U; repetition < 8U; ++repetition) {
    auto repeated = fesa::SparseAssembler::AssembleStiffness(
        model_result.Value(), dofs_result.Value(), tbb_executor);
    ASSERT_TRUE(repeated.HasValue());
    ExpectByteIdentical(repeated.Value(), serial.Value());
  }

  for (std::size_t row = 0U; row < serial.Value().Rows(); ++row) {
    for (std::size_t column = 0U; column < serial.Value().Columns(); ++column) {
      EXPECT_DOUBLE_EQ(Entry(serial.Value(), row, column),
                       Entry(serial.Value(), column, row));
    }
  }

  EXPECT_NEAR(Entry(serial.Value(), 0U, 0U), 120.0, 1.0e-12);
  EXPECT_NEAR(Entry(serial.Value(), 0U, 6U), -120.0, 1.0e-12);
  EXPECT_NEAR(Entry(serial.Value(), 6U, 6U), 200.0, 1.0e-12);
  EXPECT_NEAR(Entry(serial.Value(), 6U, 12U), -80.0, 1.0e-12);
  EXPECT_NEAR(Entry(serial.Value(), 12U, 12U), 80.0, 1.0e-12);
}

TEST(SparseAssembly,
     AssemblesFourNodeTwentyFourDofKernelAndPreservesDiagonalSlots) {
  auto domain = fesa::Domain::Create(
      MakeShellDefinition(fesa::ShellSourceElementType::kS4));
  ASSERT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());
  auto dofs = fesa::DofManager::Create(model.Value());
  ASSERT_TRUE(dofs.HasValue());
  fesa::SerialParallelFor serial_executor;

  auto assembled = fesa::SparseAssembler::AssembleStiffness(
      model.Value(), dofs.Value(), serial_executor);
  auto expected = DirectShellStiffness(domain.Value(), 0U);
  ASSERT_TRUE(assembled.HasValue());
  ASSERT_TRUE(expected.HasValue());

  EXPECT_EQ(assembled.Value().Rows(), 24U);
  EXPECT_EQ(assembled.Value().Columns(), 24U);
  EXPECT_EQ(assembled.Value().Values().size(), 24U * 24U);
  EXPECT_EQ(assembled.Value().RowOffsets(),
            dofs.Value().GetSparsePattern().row_offsets);
  EXPECT_EQ(assembled.Value().ColumnIndices(),
            dofs.Value().GetSparsePattern().column_indices);
  for (std::size_t row = 0U; row < 24U; ++row) {
    const auto begin = assembled.Value().ColumnIndices().begin() +
                       assembled.Value().RowOffsets()[row];
    const auto end = assembled.Value().ColumnIndices().begin() +
                     assembled.Value().RowOffsets()[row + 1U];
    EXPECT_NE(std::lower_bound(begin, end, row), end);
    for (std::size_t column = 0U; column < 24U; ++column) {
      EXPECT_DOUBLE_EQ(Entry(assembled.Value(), row, column),
                       expected.Value().stabilized_global24(row, column));
    }
  }
}

TEST(SparseAssembly, ShellSerialTbbReverseAndRepeatedRunsAreByteIdentical) {
  fesa::SerialParallelFor serial_executor;
  fesa::TbbParallelFor tbb_executor;
  ReverseParallelFor reverse_executor;
  auto serial =
      AssembleShell(fesa::ShellSourceElementType::kS4, serial_executor, true);
  auto tbb =
      AssembleShell(fesa::ShellSourceElementType::kS4, tbb_executor, true);
  auto reversed =
      AssembleShell(fesa::ShellSourceElementType::kS4, reverse_executor, true);
  ASSERT_TRUE(serial.HasValue());
  ASSERT_TRUE(tbb.HasValue());
  ASSERT_TRUE(reversed.HasValue());

  EXPECT_EQ(reverse_executor.Calls(), 1U);
  EXPECT_EQ(reverse_executor.ObservedCount(), 2U);
  ExpectByteIdentical(tbb.Value(), serial.Value());
  ExpectByteIdentical(reversed.Value(), serial.Value());
  for (std::size_t repetition = 0U; repetition < 8U; ++repetition) {
    auto repeated =
        AssembleShell(fesa::ShellSourceElementType::kS4, tbb_executor, true);
    ASSERT_TRUE(repeated.HasValue());
    ExpectByteIdentical(repeated.Value(), serial.Value());
  }
}

TEST(SparseAssembly, S4AndS4rSemanticFixturesAssembleIdenticalStiffness) {
  fesa::SerialParallelFor serial_executor;
  auto s4 = AssembleShell(fesa::ShellSourceElementType::kS4, serial_executor);
  auto s4r = AssembleShell(fesa::ShellSourceElementType::kS4r, serial_executor);
  ASSERT_TRUE(s4.HasValue());
  ASSERT_TRUE(s4r.HasValue());
  EXPECT_TRUE(std::any_of(s4.Value().Values().begin(),
                          s4.Value().Values().end(),
                          [](const double value) { return value != 0.0; }));
  ExpectByteIdentical(s4r.Value(), s4.Value());
}

}  // namespace