#include "fesa/assembly/load_assembler.h"

#include <gtest/gtest.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/loads/load.h"
#include "fesa/model/domain.h"

namespace {

struct LoadFixture {
  std::unique_ptr<fesa::Domain> domain;
  std::unique_ptr<fesa::AnalysisModel> model;
  std::unique_ptr<fesa::DofManager> dofs;
};

LoadFixture MakeFixture(const std::size_t node_count,
                        std::vector<fesa::NodeSet> node_sets,
                        std::vector<fesa::BoundaryCondition> boundaries,
                        std::vector<fesa::NodalLoad> loads) {
  const std::filesystem::path source{"models/load-assembly.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "fnv1a64:abcdef0123456789";
  for (std::size_t index = 0U; index < node_count; ++index) {
    const auto label = static_cast<std::int64_t>((index + 1U) * 10U);
    definition.nodes.push_back({{"Beam-1", label, std::to_string(label)},
                                {static_cast<double>(index), 0.0, 0.0},
                                {source, index + 2U}});
  }
  definition.node_sets = std::move(node_sets);
  definition.steps = {{"Step-1",
                       std::move(boundaries),
                       std::move(loads),
                       0.1,
                       1.0,
                       0.01,
                       1.0,
                       {source, 20U}}};

  auto domain_result = fesa::Domain::Create(std::move(definition));
  if (!domain_result.HasValue()) {
    throw std::runtime_error{"Load fixture Domain construction failed."};
  }
  auto domain =
      std::make_unique<fesa::Domain>(std::move(domain_result.Value()));

  auto model_result = fesa::AnalysisModel::Create(*domain);
  if (!model_result.HasValue()) {
    throw std::runtime_error{"Load fixture AnalysisModel construction failed."};
  }
  auto model =
      std::make_unique<fesa::AnalysisModel>(std::move(model_result.Value()));

  auto dof_result = fesa::DofManager::Create(*model);
  if (!dof_result.HasValue()) {
    throw std::runtime_error{"Load fixture DofManager construction failed."};
  }
  auto dofs = std::make_unique<fesa::DofManager>(std::move(dof_result.Value()));
  return {std::move(domain), std::move(model), std::move(dofs)};
}

LoadFixture MakeShellFixture(std::vector<fesa::BoundaryCondition> boundaries,
                             std::vector<fesa::NodalLoad> loads) {
  const std::filesystem::path source{"models/shell-load-assembly.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, 2U}},
                      {{"Shell-1", 20, "20"}, {1.0, 0.0, 0.0}, {source, 3U}},
                      {{"Shell-1", 30, "30"}, {1.0, 1.0, 0.0}, {source, 4U}},
                      {{"Shell-1", 40, "40"}, {0.0, 1.0, 0.0}, {source, 5U}}};
  definition.materials = {{"Material", 1000.0, 0.25, {source, 6U}}};
  definition.shell_sections = {{"ShellSection", 0.1, 0U, {source, 7U}}};
  definition.shell_elements = {{{"Shell-1", 1, "1"},
                                fesa::ShellSourceElementType::kS4,
                                {0U, 1U, 2U, 3U},
                                0U,
                                0U,
                                {source, 8U}}};
  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.steps = {{"Step-1",
                       std::move(boundaries),
                       std::move(loads),
                       0.1,
                       1.0,
                       0.01,
                       1.0,
                       {source, 20U}}};

  auto domain_result = fesa::Domain::Create(std::move(definition));
  if (!domain_result.HasValue()) {
    throw std::runtime_error{"Shell load fixture Domain construction failed."};
  }
  auto domain =
      std::make_unique<fesa::Domain>(std::move(domain_result.Value()));

  auto model_result = fesa::AnalysisModel::Create(*domain);
  if (!model_result.HasValue()) {
    throw std::runtime_error{
        "Shell load fixture AnalysisModel construction failed."};
  }
  auto model =
      std::make_unique<fesa::AnalysisModel>(std::move(model_result.Value()));

  auto dof_result = fesa::DofManager::Create(*model);
  if (!dof_result.HasValue()) {
    throw std::runtime_error{
        "Shell load fixture DofManager construction failed."};
  }
  auto dofs = std::make_unique<fesa::DofManager>(std::move(dof_result.Value()));
  return {std::move(domain), std::move(model), std::move(dofs)};
}

fesa::SparseMatrix MakeDenseSparse(const std::size_t rows,
                                   const std::size_t columns,
                                   const std::vector<double>& values) {
  if (values.size() != rows * columns) {
    throw std::invalid_argument{
        "Dense sparse fixture has the wrong value count."};
  }

  fesa::SparsePattern pattern;
  std::vector<fesa::CooContribution> contributions;
  pattern.row_offsets.reserve(rows + 1U);
  pattern.row_offsets.push_back(0U);
  for (std::size_t row = 0U; row < rows; ++row) {
    for (std::size_t column = 0U; column < columns; ++column) {
      pattern.column_indices.push_back(column);
      contributions.push_back(
          {row, column, values[row * columns + column], row, column});
    }
    pattern.row_offsets.push_back(pattern.column_indices.size());
  }

  auto result = fesa::SparseMatrix::FromCoo(rows, columns,
                                            std::move(contributions), pattern);
  if (!result.HasValue()) {
    throw std::runtime_error{"Sparse fixture construction failed."};
  }
  return std::move(result.Value());
}

void ExpectFailureCode(const fesa::Result<fesa::Vector>& result,
                       const std::string& code) {
  ASSERT_FALSE(result.HasValue());
  EXPECT_EQ(result.GetStatus().Category(), fesa::FailureCategory::kModel);
  ASSERT_EQ(result.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(result.GetStatus().Diagnostics()[0U].code, code);
}

class FakeLoad final : public fesa::Load {
 public:
  explicit FakeLoad(std::vector<fesa::LoadContribution> contributions)
      : contributions_{std::move(contributions)} {}

  fesa::Result<std::vector<fesa::LoadContribution>> ComputeContributions(
      const fesa::LoadContext&) const override {
    return fesa::Result<std::vector<fesa::LoadContribution>>::Success(
        contributions_);
  }

 private:
  std::vector<fesa::LoadContribution> contributions_;
};

}  // namespace

// C-LOAD-001
TEST(LoadAssembly, AccumulatesGenericContributionsInSuppliedSourceOrder) {
  auto fixture = MakeFixture(1U, {}, {}, {});
  const FakeLoad first(
      {{0U, 0U, 1.0e16}, {0U, 0U, -1.0e16}, {0U, 0U, 1.0}, {0U, 2U, -4.0}});
  const FakeLoad second({{1U, 2U, 1.5}, {1U, 0U, -2.0}});
  const fesa::LoadView loads{std::cref(first), std::cref(second)};

  const auto result = fesa::LoadAssembler::AssembleFullNodalLoad(
      *fixture.model, *fixture.dofs, loads);

  ASSERT_TRUE(result.HasValue());
  EXPECT_DOUBLE_EQ(result.Value()[0U], -1.0);
  EXPECT_DOUBLE_EQ(result.Value()[2U], -2.5);
}

TEST(LoadAssembly, RejectsGenericContributionErrorsBeforeCandidateCommit) {
  auto fixture = MakeFixture(1U, {}, {}, {});
  const FakeLoad valid({{0U, 1U, 3.0}});
  const FakeLoad nonfinite(
      {{1U, 2U, std::numeric_limits<double>::quiet_NaN()}});
  const FakeLoad out_of_range({{1U, 6U, 4.0}});
  const FakeLoad wrong_order({{3U, 2U, 4.0}});

  ExpectFailureCode(fesa::LoadAssembler::AssembleFullNodalLoad(
                        *fixture.model, *fixture.dofs,
                        {std::cref(valid), std::cref(nonfinite)}),
                    "nonfinite-load-value");
  ExpectFailureCode(fesa::LoadAssembler::AssembleFullNodalLoad(
                        *fixture.model, *fixture.dofs,
                        {std::cref(valid), std::cref(out_of_range)}),
                    "invalid-load-index");
  ExpectFailureCode(fesa::LoadAssembler::AssembleFullNodalLoad(
                        *fixture.model, *fixture.dofs,
                        {std::cref(valid), std::cref(wrong_order)}),
                    "invalid-load-order");

  const auto repeated = fesa::LoadAssembler::AssembleFullNodalLoad(
      *fixture.model, *fixture.dofs, {std::cref(valid)});
  ASSERT_TRUE(repeated.HasValue());
  EXPECT_DOUBLE_EQ(repeated.Value()[1U], 3.0);
}

TEST(LoadAssembly, AssemblesNodeSetAndSixComponentLoads) {
  const std::filesystem::path source{"models/load-assembly.inp"};
  auto fixture =
      MakeFixture(2U, {{"Pair", std::nullopt, {0U, 1U}, {source, 10U}}},
                  {{"10", 1, 1, 0.0, {source, 21U}}},
                  {{"pair", 1, 1.0, {source, 30U}},
                   {"10", 2, 2.0, {source, 31U}},
                   {"20", 3, 3.0, {source, 32U}},
                   {"10", 4, -4.0, {source, 33U}},
                   {"PAIR", 5, 5.0, {source, 34U}},
                   {"20", 6, 6.0, {source, 35U}}});

  auto result =
      fesa::LoadAssembler::AssembleFullNodalLoad(*fixture.model, *fixture.dofs);
  ASSERT_TRUE(result.HasValue());
  ASSERT_EQ(result.Value().Size(), 12U);
  EXPECT_EQ(
      std::vector<double>(result.Value().Data(), result.Value().Data() + 12U),
      (std::vector<double>{1.0, 2.0, 0.0, -4.0, 5.0, 0.0, 1.0, 0.0, 3.0, 0.0,
                           5.0, 6.0}));
  EXPECT_EQ(fixture.dofs->ConstrainedDofs(), (std::vector<std::size_t>{0U}));
  EXPECT_DOUBLE_EQ(result.Value()[0U], 1.0);
}

TEST(LoadAssembly, AccumulatesSignedLoadsInSourceOrder) {
  const std::filesystem::path source{"models/load-assembly.inp"};
  auto first_order = MakeFixture(1U, {}, {},
                                 {{"10", 1, 1.0e16, {source, 30U}},
                                  {"10", 1, -1.0e16, {source, 31U}},
                                  {"10", 1, 1.0, {source, 32U}}});
  auto second_order = MakeFixture(1U, {}, {},
                                  {{"10", 1, 1.0e16, {source, 30U}},
                                   {"10", 1, 1.0, {source, 31U}},
                                   {"10", 1, -1.0e16, {source, 32U}}});

  auto first = fesa::LoadAssembler::AssembleFullNodalLoad(*first_order.model,
                                                          *first_order.dofs);
  auto second = fesa::LoadAssembler::AssembleFullNodalLoad(*second_order.model,
                                                           *second_order.dofs);
  ASSERT_TRUE(first.HasValue());
  ASSERT_TRUE(second.HasValue());
  EXPECT_DOUBLE_EQ(first.Value()[0U], 1.0);
  EXPECT_DOUBLE_EQ(second.Value()[0U], 0.0);
}

// MITC4-LOAD-001
TEST(LoadAssembly, AggregatesAllSixGlobalShellLoadComponentsInSourceOrder) {
  const std::filesystem::path source{"models/shell-load-assembly.inp"};
  auto fixture = MakeShellFixture({}, {{"10", 1, 1.0e16, {source, 30U}},
                                       {"10", 1, -1.0e16, {source, 31U}},
                                       {"10", 1, 1.0, {source, 32U}},
                                       {"10", 2, 2.0, {source, 33U}},
                                       {"10", 3, 3.0, {source, 34U}},
                                       {"10", 4, 1.0e16, {source, 35U}},
                                       {"10", 4, -1.0e16, {source, 36U}},
                                       {"10", 4, 4.0, {source, 37U}},
                                       {"10", 5, 5.0, {source, 38U}},
                                       {"10", 6, 6.0, {source, 39U}},
                                       {"10", 6, -6.0, {source, 40U}}});

  const auto result =
      fesa::LoadAssembler::AssembleFullNodalLoad(*fixture.model, *fixture.dofs);

  ASSERT_TRUE(result.HasValue());
  ASSERT_EQ(result.Value().Size(), 24U);
  EXPECT_EQ(
      std::vector<double>(result.Value().Data(), result.Value().Data() + 6U),
      (std::vector<double>{1.0, 2.0, 3.0, 4.0, 5.0, 0.0}));
}

// MITC4-LOAD-002
TEST(LoadAssembly, AcceptsExactlyZeroAggregateShellMoment) {
  const std::filesystem::path source{"models/shell-load-assembly.inp"};
  auto fixture = MakeShellFixture({}, {{"10", 4, 3.0, {source, 30U}},
                                       {"10", 4, -3.0, {source, 31U}},
                                       {"10", 5, 4.0, {source, 32U}},
                                       {"10", 5, -4.0, {source, 33U}},
                                       {"10", 6, 5.0, {source, 34U}},
                                       {"10", 6, -5.0, {source, 35U}}});

  const auto result =
      fesa::LoadAssembler::AssembleFullNodalLoad(*fixture.model, *fixture.dofs);

  ASSERT_TRUE(result.HasValue());
  EXPECT_DOUBLE_EQ(result.Value()[3U], 0.0);
  EXPECT_DOUBLE_EQ(result.Value()[4U], 0.0);
  EXPECT_DOUBLE_EQ(result.Value()[5U], 0.0);
}

// MITC4-LOAD-003
TEST(LoadAssembly, EnforcesAggregateShellMomentDirectorProjectionThreshold) {
  const std::filesystem::path source{"models/shell-load-assembly.inp"};
  auto accepted_fixture = MakeShellFixture(
      {}, {{"10", 4, 1.0, {source, 30U}}, {"10", 6, 1.0e-12, {source, 31U}}});
  auto rejected_fixture = MakeShellFixture(
      {}, {{"10", 4, 1.0, {source, 30U}}, {"10", 6, 2.0e-12, {source, 31U}}});

  const auto accepted = fesa::LoadAssembler::AssembleFullNodalLoad(
      *accepted_fixture.model, *accepted_fixture.dofs);
  const auto rejected = fesa::LoadAssembler::AssembleFullNodalLoad(
      *rejected_fixture.model, *rejected_fixture.dofs);

  ASSERT_TRUE(accepted.HasValue());
  ASSERT_FALSE(rejected.HasValue());
  EXPECT_EQ(rejected.GetStatus().Category(), fesa::FailureCategory::kModel);
  ASSERT_EQ(rejected.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(rejected.GetStatus().Diagnostics()[0U].code,
            "unsupported-drilling-load");
  EXPECT_EQ(rejected.GetStatus().Diagnostics()[0U].keyword, "CLOAD");
  EXPECT_EQ(rejected.GetStatus().Diagnostics()[0U].entity_identity, "10");
}

// MITC4-LOAD-004
TEST(LoadAssembly, RejectsDrillingMomentBeforeEffectiveRhsCanBeFormed) {
  const std::filesystem::path source{"models/shell-load-assembly.inp"};
  auto fixture = MakeShellFixture({{"10", 1, 1, 2.0, {source, 21U}}},
                                  {{"10", 6, 1.0, {source, 30U}}});

  const auto rejected =
      fesa::LoadAssembler::AssembleFullNodalLoad(*fixture.model, *fixture.dofs);

  ASSERT_FALSE(rejected.HasValue());
  EXPECT_EQ(rejected.GetStatus().Category(), fesa::FailureCategory::kModel);
  ASSERT_EQ(rejected.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(rejected.GetStatus().Diagnostics()[0U].code,
            "unsupported-drilling-load");
}

TEST(LoadAssembly, FormsNonzeroPrescribedEffectiveRhs) {
  const std::filesystem::path source{"models/load-assembly.inp"};
  auto fixture = MakeFixture(
      1U, {},
      {{"10", 2, 2, 2.0, {source, 21U}}, {"10", 5, 5, -1.0, {source, 22U}}},
      {{"10", 1, 10.0, {source, 30U}},
       {"10", 2, 900.0, {source, 31U}},
       {"10", 3, 20.0, {source, 32U}},
       {"10", 4, 30.0, {source, 33U}},
       {"10", 5, 800.0, {source, 34U}},
       {"10", 6, 40.0, {source, 35U}}});
  auto full =
      fesa::LoadAssembler::AssembleFullNodalLoad(*fixture.model, *fixture.dofs);
  ASSERT_TRUE(full.HasValue());
  const auto kfc =
      MakeDenseSparse(4U, 2U, {1.0, 2.0, 3.0, 4.0, -2.0, 5.0, 0.5, -1.0});

  auto rhs = fesa::LoadAssembler::EffectiveFreeRhs(
      full.Value(), kfc, fixture.dofs->PrescribedValues(), *fixture.dofs);
  ASSERT_TRUE(rhs.HasValue());
  ASSERT_EQ(rhs.Value().Size(), 4U);
  EXPECT_EQ(std::vector<double>(rhs.Value().Data(), rhs.Value().Data() + 4U),
            (std::vector<double>{10.0, 18.0, 39.0, 38.0}));
}

TEST(LoadAssembly, RejectsNonfiniteOrDimensionMismatch) {
  const std::filesystem::path source{"models/load-assembly.inp"};
  const double maximum = (std::numeric_limits<double>::max)();
  auto nonfinite = MakeFixture(
      1U, {}, {},
      {{"10", 1, std::numeric_limits<double>::quiet_NaN(), {source, 30U}}});
  ExpectFailureCode(fesa::LoadAssembler::AssembleFullNodalLoad(*nonfinite.model,
                                                               *nonfinite.dofs),
                    "nonfinite-load-value");

  auto overflow = MakeFixture(
      1U, {}, {},
      {{"10", 1, maximum, {source, 30U}}, {"10", 1, maximum, {source, 31U}}});
  ExpectFailureCode(fesa::LoadAssembler::AssembleFullNodalLoad(*overflow.model,
                                                               *overflow.dofs),
                    "nonfinite-load-accumulation");

  auto one_node = MakeFixture(
      1U, {},
      {{"10", 2, 2, 2.0, {source, 21U}}, {"10", 5, 5, -1.0, {source, 22U}}},
      {});
  auto two_nodes = MakeFixture(2U, {}, {}, {});
  ExpectFailureCode(fesa::LoadAssembler::AssembleFullNodalLoad(*two_nodes.model,
                                                               *one_node.dofs),
                    "invalid-load-dimensions");

  const auto valid_kfc = MakeDenseSparse(4U, 2U, std::vector<double>(8U, 0.0));
  ExpectFailureCode(fesa::LoadAssembler::EffectiveFreeRhs(
                        fesa::Vector{5U}, valid_kfc,
                        one_node.dofs->PrescribedValues(), *one_node.dofs),
                    "invalid-load-dimensions");
  ExpectFailureCode(fesa::LoadAssembler::EffectiveFreeRhs(
                        fesa::Vector{6U},
                        MakeDenseSparse(3U, 2U, std::vector<double>(6U, 0.0)),
                        one_node.dofs->PrescribedValues(), *one_node.dofs),
                    "invalid-load-dimensions");
  ExpectFailureCode(fesa::LoadAssembler::EffectiveFreeRhs(
                        fesa::Vector{6U},
                        MakeDenseSparse(4U, 1U, std::vector<double>(4U, 0.0)),
                        one_node.dofs->PrescribedValues(), *one_node.dofs),
                    "invalid-load-dimensions");
  ExpectFailureCode(
      fesa::LoadAssembler::EffectiveFreeRhs(fesa::Vector{6U}, valid_kfc,
                                            fesa::Vector{1U}, *one_node.dofs),
      "invalid-load-dimensions");

  fesa::Vector nonfinite_full{6U};
  nonfinite_full[0U] = std::numeric_limits<double>::infinity();
  ExpectFailureCode(fesa::LoadAssembler::EffectiveFreeRhs(
                        nonfinite_full, valid_kfc,
                        one_node.dofs->PrescribedValues(), *one_node.dofs),
                    "nonfinite-load-value");

  fesa::Vector nonfinite_prescribed{2U};
  nonfinite_prescribed[0U] = std::numeric_limits<double>::quiet_NaN();
  ExpectFailureCode(
      fesa::LoadAssembler::EffectiveFreeRhs(
          fesa::Vector{6U}, valid_kfc, nonfinite_prescribed, *one_node.dofs),
      "nonfinite-load-value");

  const auto overflowing_kfc =
      MakeDenseSparse(4U, 2U, {maximum, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0});
  ExpectFailureCode(fesa::LoadAssembler::EffectiveFreeRhs(
                        fesa::Vector{6U}, overflowing_kfc,
                        one_node.dofs->PrescribedValues(), *one_node.dofs),
                    "nonfinite-load-accumulation");
}

TEST(LoadAssembly, ZeroLoadsRemainZero) {
  const std::filesystem::path source{"models/load-assembly.inp"};
  auto free_fixture = MakeFixture(1U, {}, {}, {{"10", 3, 0.0, {source, 30U}}});
  auto full = fesa::LoadAssembler::AssembleFullNodalLoad(*free_fixture.model,
                                                         *free_fixture.dofs);
  ASSERT_TRUE(full.HasValue());
  EXPECT_TRUE(std::all_of(full.Value().Data(),
                          full.Value().Data() + full.Value().Size(),
                          [](const double value) { return value == 0.0; }));
  const auto no_constrained_columns = MakeDenseSparse(6U, 0U, {});
  auto free_rhs = fesa::LoadAssembler::EffectiveFreeRhs(
      full.Value(), no_constrained_columns,
      free_fixture.dofs->PrescribedValues(), *free_fixture.dofs);
  ASSERT_TRUE(free_rhs.HasValue());
  EXPECT_EQ(free_rhs.Value().Size(), 6U);
  EXPECT_TRUE(std::all_of(free_rhs.Value().Data(),
                          free_rhs.Value().Data() + free_rhs.Value().Size(),
                          [](const double value) { return value == 0.0; }));

  auto constrained_fixture =
      MakeFixture(1U, {}, {{"10", 1, 6, 0.0, {source, 21U}}}, {});
  auto constrained_full = fesa::LoadAssembler::AssembleFullNodalLoad(
      *constrained_fixture.model, *constrained_fixture.dofs);
  ASSERT_TRUE(constrained_full.HasValue());
  const auto no_free_rows = MakeDenseSparse(0U, 6U, {});
  auto constrained_rhs = fesa::LoadAssembler::EffectiveFreeRhs(
      constrained_full.Value(), no_free_rows,
      constrained_fixture.dofs->PrescribedValues(), *constrained_fixture.dofs);
  ASSERT_TRUE(constrained_rhs.HasValue());
  EXPECT_EQ(constrained_rhs.Value().Size(), 0U);
}