#include "fesa/fem/dof_manager.h"

#include <gtest/gtest.h>

#include <algorithm>
#include <array>
#include <cstddef>
#include <filesystem>
#include <functional>
#include <optional>
#include <string>
#include <utility>
#include <vector>

namespace fesa {

class DofManagerTestAccess {
 public:
  static void DuplicateOwnership(DofManager& dofs) {
    dofs.constrained_dofs_[3U] = dofs.free_dofs_[7U];
  }

  static void RemoveFullMapping(DofManager& dofs) {
    dofs.free_equations_.pop_back();
  }

  static void ReverseFreeMapping(DofManager& dofs) {
    std::swap(dofs.free_dofs_[0U], dofs.free_dofs_[1U]);
  }

  static void ReverseConstrainedMapping(DofManager& dofs) {
    std::swap(dofs.constrained_dofs_[0U], dofs.constrained_dofs_[1U]);
  }

  static void CorruptEquationMapping(DofManager& dofs) {
    dofs.free_equations_[dofs.free_dofs_[0U]] = dofs.free_dofs_.size();
  }

  static void RemoveSparsePatternEntry(DofManager& dofs) {
    const std::size_t position = dofs.sparse_pattern_.row_offsets[1U] - 1U;
    dofs.sparse_pattern_.column_indices.erase(
        dofs.sparse_pattern_.column_indices.begin() +
        static_cast<std::ptrdiff_t>(position));
    for (std::size_t row = 1U; row < dofs.sparse_pattern_.row_offsets.size();
         ++row) {
      --dofs.sparse_pattern_.row_offsets[row];
    }
  }

  static void AddSparsePatternEntry(DofManager& dofs) {
    const std::size_t position = dofs.sparse_pattern_.row_offsets[1U];
    dofs.sparse_pattern_.column_indices.insert(
        dofs.sparse_pattern_.column_indices.begin() +
            static_cast<std::ptrdiff_t>(position),
        12U);
    for (std::size_t row = 1U; row < dofs.sparse_pattern_.row_offsets.size();
         ++row) {
      ++dofs.sparse_pattern_.row_offsets[row];
    }
  }
};

}  // namespace fesa

namespace {

class GenericLayoutElement final : public fesa::Element {
 public:
  explicit GenericLayoutElement(fesa::ElementDofLayout layout)
      : layout_{std::move(layout)} {}

  const fesa::ElementDofLayout& DofLayout() const noexcept override {
    return layout_;
  }

  fesa::Result<fesa::ElementStiffnessContribution> ComputeStiffness()
      const override {
    const std::size_t local_dof_count =
        layout_.node_indices.size() * layout_.components_per_node.size();
    return fesa::Result<fesa::ElementStiffnessContribution>::Success(
        {layout_, fesa::Matrix{local_dof_count, local_dof_count}});
  }

  fesa::Result<fesa::ElementResultBundle> Recover(
      const fesa::Vector&) const override {
    return fesa::Result<fesa::ElementResultBundle>::Success(
        {layout_.source_id, fesa::BeamElementResultRows{}});
  }

 private:
  fesa::ElementDofLayout layout_;
};

fesa::ModelDefinition MakeDefinition() {
  const std::filesystem::path source{"models/dof-manager.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "fnv1a64:0123456789abcdef";
  definition.nodes = {{{"Beam-1", 10, "10"}, {0.0, 0.0, 0.0}, {source, 10U}},
                      {{"Beam-1", 20, "20"}, {1.0, 0.0, 0.0}, {source, 11U}},
                      {{"Beam-1", 30, "30"}, {2.0, 0.0, 0.0}, {source, 12U}}};
  definition.materials = {{"Material", 1000.0, 0.25, {source, 20U}}};
  definition.sections = {
      {"Section", 1.0, 1.0, 0.0, 1.0, 1.0, {0.0, 1.0, 0.0}, {}, {source, 30U}}};
  definition.elements = {
      {{"Beam-1", 100, "100"}, {0U, 1U}, 0U, 0U, {source, 40U}},
      {{"Beam-1", 200, "200"}, {1U, 2U}, 0U, 0U, {source, 41U}}};
  definition.node_sets = {
      {"Root", std::optional<std::string>{"Beam-1"}, {0U}, {source, 50U}},
      {"Ends", std::optional<std::string>{"Beam-1"}, {0U, 2U}, {source, 51U}}};
  definition.steps = {{"Step-1",
                       {{"Root", 1, 2, 0.0, {source, 60U}},
                        {"ends", 3, 3, 0.25, {source, 61U}},
                        {"20", 6, 6, -0.5, {source, 62U}},
                        {"Root", 1, 1, 0.0, {source, 63U}}},
                       {},
                       0.1,
                       1.0,
                       0.01,
                       1.0,
                       {source, 59U}}};
  return definition;
}

fesa::ModelDefinition MakeShellDefinition() {
  const std::filesystem::path source{"models/shell-dof-manager.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "fnv1a64:1122334455667788";
  definition.nodes = {{{"Shell-1", 10, "10"}, {0.0, 0.0, 0.0}, {source, 10U}},
                      {{"Shell-1", 20, "20"}, {1.0, 0.0, 0.0}, {source, 11U}},
                      {{"Shell-1", 30, "30"}, {1.0, 1.0, 0.0}, {source, 12U}},
                      {{"Shell-1", 40, "40"}, {0.0, 1.0, 0.0}, {source, 13U}}};
  definition.materials = {{"Material", 1000.0, 0.25, {source, 20U}}};
  definition.shell_sections = {{"ShellSection", 0.1, 0U, {source, 30U}}};
  definition.shell_elements = {{{"Shell-1", 100, "100"},
                                fesa::ShellSourceElementType::kS4,
                                {2U, 0U, 3U, 1U},
                                0U,
                                0U,
                                {source, 40U}}};
  definition.steps = {{"Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0, {source, 50U}}};
  return definition;
}

struct DofFixture {
  fesa::DofManager dofs;
};

DofFixture MakeDofFixture(fesa::ModelDefinition definition = MakeDefinition()) {
  auto domain = fesa::Domain::Create(std::move(definition));
  EXPECT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  EXPECT_TRUE(model.HasValue());
  auto dofs = fesa::DofManager::Create(model.Value());
  EXPECT_TRUE(dofs.HasValue());
  return {std::move(dofs.Value())};
}

std::vector<std::size_t> RowColumns(const fesa::SparsePattern& pattern,
                                    std::size_t row) {
  return {pattern.column_indices.begin() + pattern.row_offsets[row],
          pattern.column_indices.begin() + pattern.row_offsets[row + 1U]};
}

void ExpectInvariantFailure(const fesa::DofManager& dofs,
                            const std::string& code) {
  const fesa::Status status = dofs.ValidateInvariants();
  ASSERT_FALSE(status.IsOk());
  EXPECT_EQ(status.Category(), fesa::FailureCategory::kModel);
  ASSERT_EQ(status.Diagnostics().size(), 1U);
  EXPECT_EQ(status.Diagnostics()[0U].code, code);
}

}  // namespace

// C-DOF-001
TEST(DofManager, BuildsGenericRuntimeLayoutsInDeclaredSourceOrder) {
  auto domain = fesa::Domain::Create(MakeDefinition());
  ASSERT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());

  GenericLayoutElement first{fesa::ElementDofLayout{
      {"Beam-1", 100, "100"},
      {2U, 0U},
      {fesa::DofComponent::kUz, fesa::DofComponent::kUx}}};
  GenericLayoutElement second{
      fesa::ElementDofLayout{{"Beam-1", 200, "200"},
                             {1U},
                             {fesa::DofComponent::kUrz, fesa::DofComponent::kUy,
                              fesa::DofComponent::kUrx}}};
  const fesa::ElementView elements{std::cref(first), std::cref(second)};

  fesa::DofManager dofs;
  ASSERT_TRUE(dofs.Build(model.Value(), elements).IsOk());
  ASSERT_TRUE(dofs.ValidateInvariants().IsOk());

  auto first_scatter = dofs.ElementScatter(first.DofLayout());
  auto second_scatter = dofs.ElementScatter(second.DofLayout());
  ASSERT_TRUE(first_scatter.HasValue());
  ASSERT_TRUE(second_scatter.HasValue());
  EXPECT_EQ(first_scatter.Value(),
            (std::vector<std::size_t>{14U, 12U, 2U, 0U}));
  EXPECT_EQ(second_scatter.Value(), (std::vector<std::size_t>{11U, 7U, 9U}));

  const auto& pattern = dofs.GetSparsePattern();
  EXPECT_EQ(RowColumns(pattern, 14U),
            (std::vector<std::size_t>{0U, 2U, 12U, 14U}));
  EXPECT_EQ(RowColumns(pattern, 7U), (std::vector<std::size_t>{7U, 9U, 11U}));

  const auto expected_free_dofs = dofs.FreeDofs();
  const auto expected_constrained_dofs = dofs.ConstrainedDofs();
  const auto expected_row_offsets = pattern.row_offsets;
  const auto expected_columns = pattern.column_indices;
  const fesa::ElementView reversed_elements{std::cref(second),
                                            std::cref(first)};
  const fesa::Status reversed_status =
      dofs.Build(model.Value(), reversed_elements);
  ASSERT_FALSE(reversed_status.IsOk());
  ASSERT_EQ(reversed_status.Diagnostics().size(), 1U);
  EXPECT_EQ(reversed_status.Diagnostics()[0U].code,
            "invalid-element-layout-order");
  EXPECT_EQ(dofs.FreeDofs(), expected_free_dofs);
  EXPECT_EQ(dofs.ConstrainedDofs(), expected_constrained_dofs);
  EXPECT_EQ(dofs.GetSparsePattern().row_offsets, expected_row_offsets);
  EXPECT_EQ(dofs.GetSparsePattern().column_indices, expected_columns);
}

TEST(DofManager, RejectsEveryCorruptedOwnerMappingInvariant) {
  const auto fixture = MakeDofFixture();

  auto duplicate = fixture.dofs;
  fesa::DofManagerTestAccess::DuplicateOwnership(duplicate);
  ExpectInvariantFailure(duplicate, "duplicate-dof-mapping");

  auto full = fixture.dofs;
  fesa::DofManagerTestAccess::RemoveFullMapping(full);
  ExpectInvariantFailure(full, "invalid-dof-dimensions");

  auto free = fixture.dofs;
  fesa::DofManagerTestAccess::ReverseFreeMapping(free);
  ExpectInvariantFailure(free, "invalid-free-dof-mapping");

  auto constrained = fixture.dofs;
  fesa::DofManagerTestAccess::ReverseConstrainedMapping(constrained);
  ExpectInvariantFailure(constrained, "invalid-constrained-dof-mapping");

  auto equation = fixture.dofs;
  fesa::DofManagerTestAccess::CorruptEquationMapping(equation);
  ExpectInvariantFailure(equation, "invalid-equation-mapping");

  auto missing_pattern_entry = fixture.dofs;
  fesa::DofManagerTestAccess::RemoveSparsePatternEntry(missing_pattern_entry);
  ExpectInvariantFailure(missing_pattern_entry, "invalid-dof-sparse-pattern");

  auto extra_pattern_entry = fixture.dofs;
  fesa::DofManagerTestAccess::AddSparsePatternEntry(extra_pattern_entry);
  ExpectInvariantFailure(extra_pattern_entry, "invalid-dof-sparse-pattern");
}

TEST(DofManager, NumbersSixDofsAndFreeEquationsStably) {
  const auto fixture = MakeDofFixture();
  const auto& dofs = fixture.dofs;

  EXPECT_EQ(dofs.FullDofCount(), 18U);
  EXPECT_EQ(dofs.FreeDofCount(), 13U);
  EXPECT_EQ(dofs.ConstrainedDofCount(), 5U);
  EXPECT_EQ(dofs.FullDof(0U, fesa::DofComponent::kUx), 0U);
  EXPECT_EQ(dofs.FullDof(0U, fesa::DofComponent::kUrz), 5U);
  EXPECT_EQ(dofs.FullDof(1U, fesa::DofComponent::kUx), 6U);
  EXPECT_EQ(dofs.FullDof(2U, fesa::DofComponent::kUrz), 17U);
  EXPECT_EQ(dofs.FreeDofs(),
            (std::vector<std::size_t>{3U, 4U, 5U, 6U, 7U, 8U, 9U, 10U, 12U, 13U,
                                      15U, 16U, 17U}));
  EXPECT_EQ(dofs.ConstrainedDofs(),
            (std::vector<std::size_t>{0U, 1U, 2U, 11U, 14U}));
  EXPECT_EQ(dofs.FreeEquation(0U), std::nullopt);
  EXPECT_EQ(dofs.FreeEquation(3U), std::optional<std::size_t>{0U});
  EXPECT_EQ(dofs.FreeEquation(10U), std::optional<std::size_t>{7U});
  EXPECT_EQ(dofs.FreeEquation(17U), std::optional<std::size_t>{12U});
}

TEST(DofManager, ExpandsAndValidatesPrescribedValues) {
  const auto fixture = MakeDofFixture();
  const auto& values = fixture.dofs.PrescribedValues();
  ASSERT_EQ(values.Size(), 5U);
  EXPECT_DOUBLE_EQ(values[0], 0.0);
  EXPECT_DOUBLE_EQ(values[1], 0.0);
  EXPECT_DOUBLE_EQ(values[2], 0.25);
  EXPECT_DOUBLE_EQ(values[3], -0.5);
  EXPECT_DOUBLE_EQ(values[4], 0.25);

  auto conflicting_definition = MakeDefinition();
  conflicting_definition.steps[0].boundaries.push_back(
      {"root", 1, 1, 1.0, {conflicting_definition.source_path, 77U}});
  auto domain = fesa::Domain::Create(std::move(conflicting_definition));
  ASSERT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());

  auto conflict = fesa::DofManager::Create(model.Value());
  ASSERT_FALSE(conflict.HasValue());
  EXPECT_EQ(conflict.GetStatus().Category(), fesa::FailureCategory::kInput);
  ASSERT_EQ(conflict.GetStatus().Diagnostics().size(), 1U);
  const auto& diagnostic = conflict.GetStatus().Diagnostics()[0];
  EXPECT_EQ(diagnostic.code, "conflicting-boundary-condition");
  EXPECT_EQ(diagnostic.keyword, "BOUNDARY");
  EXPECT_EQ(diagnostic.entity_identity, "root");
  EXPECT_EQ(diagnostic.location.file,
            std::filesystem::path{"models/dof-manager.inp"});
  EXPECT_EQ(diagnostic.location.line, 77U);
}

TEST(DofManager, BuildsTwelveDofScatterAndSortedUniquePattern) {
  const auto fixture = MakeDofFixture();
  const auto& dofs = fixture.dofs;

  EXPECT_EQ(dofs.ElementScatter(0U),
            (std::array<std::size_t, 12>{0U, 1U, 2U, 3U, 4U, 5U, 6U, 7U, 8U, 9U,
                                         10U, 11U}));
  EXPECT_EQ(dofs.ElementScatter(1U),
            (std::array<std::size_t, 12>{6U, 7U, 8U, 9U, 10U, 11U, 12U, 13U,
                                         14U, 15U, 16U, 17U}));

  const auto& pattern = dofs.GetSparsePattern();
  EXPECT_EQ(pattern.row_offsets,
            (std::vector<std::size_t>{0U, 12U, 24U, 36U, 48U, 60U, 72U, 90U,
                                      108U, 126U, 144U, 162U, 180U, 192U, 204U,
                                      216U, 228U, 240U, 252U}));
  const std::vector<std::size_t> first_block{0U, 1U, 2U, 3U, 4U,  5U,
                                             6U, 7U, 8U, 9U, 10U, 11U};
  const std::vector<std::size_t> shared_block{0U,  1U,  2U,  3U,  4U,  5U,
                                              6U,  7U,  8U,  9U,  10U, 11U,
                                              12U, 13U, 14U, 15U, 16U, 17U};
  const std::vector<std::size_t> last_block{6U,  7U,  8U,  9U,  10U, 11U,
                                            12U, 13U, 14U, 15U, 16U, 17U};
  EXPECT_EQ(RowColumns(pattern, 0U), first_block);
  EXPECT_EQ(RowColumns(pattern, 7U), shared_block);
  EXPECT_EQ(RowColumns(pattern, 17U), last_block);
  for (std::size_t row = 0U; row < dofs.FullDofCount(); ++row) {
    const auto columns = RowColumns(pattern, row);
    EXPECT_TRUE(std::is_sorted(columns.begin(), columns.end()));
    EXPECT_EQ(std::adjacent_find(columns.begin(), columns.end()),
              columns.end());
  }
}

// MITC4-DOF-001
TEST(DofManager, BuildsShellScatterInSourceNodeAndComponentOrder) {
  const auto fixture = MakeDofFixture(MakeShellDefinition());
  const auto& dofs = fixture.dofs;

  EXPECT_EQ(dofs.FullDofCount(), 24U);
  EXPECT_EQ(dofs.ShellElementScatter(0U),
            (std::array<std::size_t, 24>{
                12U, 13U, 14U, 15U, 16U, 17U, 0U, 1U, 2U, 3U, 4U,  5U,
                18U, 19U, 20U, 21U, 22U, 23U, 6U, 7U, 8U, 9U, 10U, 11U}));
}

// MITC4-DOF-002
TEST(DofManager, IncludesShellConnectivityInSortedUniqueSparsePattern) {
  const auto fixture = MakeDofFixture(MakeShellDefinition());
  const auto& dofs = fixture.dofs;
  const auto& pattern = dofs.GetSparsePattern();

  ASSERT_EQ(pattern.row_offsets.size(), 25U);
  ASSERT_EQ(pattern.column_indices.size(), 24U * 24U);
  for (std::size_t row = 0U; row < dofs.FullDofCount(); ++row) {
    EXPECT_EQ(pattern.row_offsets[row], row * 24U);
    EXPECT_EQ(pattern.row_offsets[row + 1U], (row + 1U) * 24U);
    const auto columns = RowColumns(pattern, row);
    ASSERT_EQ(columns.size(), 24U);
    for (std::size_t column = 0U; column < columns.size(); ++column) {
      EXPECT_EQ(columns[column], column);
    }
    EXPECT_TRUE(std::binary_search(columns.begin(), columns.end(), row));
    EXPECT_EQ(std::adjacent_find(columns.begin(), columns.end()),
              columns.end());
  }
}

TEST(DofManager, ReconstructsFullReducedRoundTrip) {
  const auto fixture = MakeDofFixture();
  const auto& dofs = fixture.dofs;
  fesa::Vector full{dofs.FullDofCount()};
  for (std::size_t index = 0U; index < full.Size(); ++index) {
    full[index] = static_cast<double>(index) + 0.5;
  }
  for (std::size_t index = 0U; index < dofs.ConstrainedDofCount(); ++index) {
    full[dofs.ConstrainedDofs()[index]] = dofs.PrescribedValues()[index];
  }

  fesa::Vector reduced{dofs.FreeDofCount()};
  for (std::size_t equation = 0U; equation < reduced.Size(); ++equation) {
    reduced[equation] = full[dofs.FreeDofs()[equation]];
  }
  fesa::Vector reconstructed{dofs.FullDofCount()};
  for (std::size_t equation = 0U; equation < reduced.Size(); ++equation) {
    reconstructed[dofs.FreeDofs()[equation]] = reduced[equation];
  }
  for (std::size_t index = 0U; index < dofs.ConstrainedDofCount(); ++index) {
    reconstructed[dofs.ConstrainedDofs()[index]] =
        dofs.PrescribedValues()[index];
  }

  ASSERT_EQ(reconstructed.Size(), full.Size());
  for (std::size_t index = 0U; index < full.Size(); ++index) {
    EXPECT_DOUBLE_EQ(reconstructed[index], full[index]);
  }
}