#include "fesa/analysis/dynamic_analysis_state.h"

#include <gtest/gtest.h>

#include <array>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/fem/dof_manager.h"

namespace {

constexpr std::array<std::array<double, 3>, 3> kIdentityTriad{
    std::array<double, 3>{1.0, 0.0, 0.0}, std::array<double, 3>{0.0, 1.0, 0.0},
    std::array<double, 3>{0.0, 0.0, 1.0}};

fesa::ModelDefinition MakeDefinition(const bool fixed = true) {
  const std::filesystem::path source{"dynamic-state.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "hht-output-unit-fixture";
  definition.nodes = {{{"Shell-A", 10, "0010"}, {0.0, 0.0, 0.0}, {source, 1U}},
                      {{"Shell-A", 20, "20"}, {1.0, 0.0, 0.0}, {source, 2U}},
                      {{"Shell-A", 30, "30"}, {1.0, 1.0, 0.0}, {source, 3U}},
                      {{"Shell-A", 40, "40"}, {0.0, 1.0, 0.0}, {source, 4U}}};
  definition.materials = {{"Material", 1000.0, 0.25, {source, 5U}}};
  definition.densities = {{0U, 2.0, {source, 6U}}};
  definition.shell_sections = {{"Shell", 0.1, 0U, {source, 7U}}};
  definition.shell_elements = {{{"Shell-A", 9, "0009"},
                                fesa::ShellSourceElementType::kS4,
                                {0U, 1U, 2U, 3U},
                                0U,
                                0U,
                                {source, 8U}}};
  for (std::size_t node = 0U; node < 4U; ++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.node_sets = {{"All", std::nullopt, {0U, 1U, 2U, 3U}, {source, 9U}},
                          {"Loaded", std::nullopt, {0U}, {source, 10U}}};
  fesa::DynamicStepDefinition step{};
  step.source_step_name_raw = "Transient Load";
  step.source_step_name_canonical = "TRANSIENT LOAD";
  step.initial_increment = 0.25;
  step.time_period = 0.5;
  step.increment_count = 2U;
  step.maximum_increments = 3U;
  if (fixed) {
    step.boundaries = {{"All", 1, 6, 0.0, {source, 11U}}};
    step.loads = {{"Loaded", 3, 4.0, {source, 12U}}};
  }
  definition.dynamic_step = step;
  return definition;
}

struct Fixture {
  std::unique_ptr<fesa::Domain> domain;
  fesa::DynamicAnalysisState state;
};

Fixture MakeFixture(fesa::ModelDefinition definition = MakeDefinition()) {
  auto domain_result = fesa::Domain::Create(std::move(definition));
  if (!domain_result.HasValue())
    throw std::runtime_error{"Invalid fixture Domain."};
  auto domain =
      std::make_unique<fesa::Domain>(std::move(domain_result.Value()));
  auto model = fesa::AnalysisModel::Create(*domain);
  if (!model.HasValue())
    throw std::runtime_error{"Invalid fixture AnalysisModel."};
  auto dofs = fesa::DofManager::Create(model.Value());
  if (!dofs.HasValue()) throw std::runtime_error{"Invalid fixture DofManager."};
  auto state = fesa::DynamicAnalysisState::Create(*domain, dofs.Value());
  if (!state.HasValue()) throw std::runtime_error{"Invalid fixture state."};
  return {std::move(domain), std::move(state.Value())};
}

// A fully restrained, initially undeformed shell under a dead nodal force has
// zero motion and strain, and nonzero support force at every accepted endpoint.
fesa::DynamicAcceptedFrameState MakeFrame(const fesa::Domain& domain,
                                          const std::uint32_t index) {
  fesa::DynamicAcceptedFrameState frame{};
  frame.frame_index = index;
  frame.increment_index = index;
  frame.time = static_cast<double>(index) * domain.Dynamic()->initial_increment;
  frame.is_initial = index == 0U;
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    frame.nodal_internal_order.push_back(static_cast<fesa::EntityIndex>(node));
    frame.nodal_source_ids.push_back(domain.Nodes()[node].source_id);
    frame.nodal_triad.push_back(kIdentityTriad);
    frame.nodal_director.push_back({0.0, 0.0, 1.0});
  }
  for (auto* vector :
       {&frame.generalized_coordinates, &frame.velocity, &frame.acceleration,
        &frame.algorithmic_acceleration, &frame.inertia_force,
        &frame.internal_force, &frame.external_force, &frame.residual,
        &frame.reaction, &frame.residual_scale}) {
    *vector = fesa::Vector{24U};
  }
  frame.angular_velocity = fesa::Vector{12U};
  frame.angular_acceleration = fesa::Vector{12U};
  if (!domain.Dynamic()->loads.empty()) {
    frame.external_force[2U] = 4.0;
    frame.residual[2U] = -4.0;
    frame.reaction[2U] = -4.0;
    frame.residual_scale[2U] = 4.0;
  }
  const double gauss = 1.0 / std::sqrt(3.0);
  const std::array<std::array<double, 2>, 4> natural{
      std::array<double, 2>{-gauss, -gauss},
      {gauss, -gauss},
      {gauss, gauss},
      {-gauss, gauss}};
  const std::array<fesa::ShellMidsurfaceLocation, 4> locations{
      fesa::ShellMidsurfaceLocation::kGp1, fesa::ShellMidsurfaceLocation::kGp2,
      fesa::ShellMidsurfaceLocation::kGp3, fesa::ShellMidsurfaceLocation::kGp4};
  for (std::size_t point = 0U; point < 4U; ++point) {
    frame.shell_results.push_back(
        {0U,
         locations[point],
         natural[point],
         kIdentityTriad,
         {},
         {},
         {{{fesa::ShellSectionPosition::kBottom, -1.0, {}},
           {fesa::ShellSectionPosition::kMiddle, 0.0, {}},
           {fesa::ShellSectionPosition::kTop, 1.0, {}}}}});
  }
  return frame;
}

TEST(DynamicAnalysisState, RejectsUnknownOrChangingIntegrationFormulation) {
  auto fixture = MakeFixture();
  auto initial = MakeFrame(*fixture.domain, 0U);
  initial.time_integration_formulation =
      static_cast<fesa::DynamicTimeIntegrationFormulation>(99);
  EXPECT_FALSE(fixture.state.CommitAcceptedFrame(initial).IsOk());
  EXPECT_TRUE(fixture.state.AcceptedFrames().empty());
  initial.time_integration_formulation =
      fesa::DynamicTimeIntegrationFormulation::kReducedCoordinateHht;
  ASSERT_TRUE(fixture.state.CommitAcceptedFrame(initial).IsOk());
  auto next = MakeFrame(*fixture.domain, 1U);
  EXPECT_FALSE(fixture.state.CommitAcceptedFrame(next).IsOk());
  EXPECT_EQ(fixture.state.AcceptedFrames().size(), 1U);
  next.time_integration_formulation = initial.time_integration_formulation;
  EXPECT_TRUE(fixture.state.CommitAcceptedFrame(next).IsOk());
}

TEST(DynamicAnalysisState, DefaultStateRejectsTrialAndCommit) {
  fesa::DynamicAnalysisState state;
  EXPECT_FALSE(state.IsInitialized());
  EXPECT_TRUE(state.AcceptedFrames().empty());
  EXPECT_FALSE(state.BeginTrialFromCommitted().HasValue());
  EXPECT_THROW((void)state.CommittedState(), std::logic_error);
  const auto status = state.CommitAcceptedFrame({});
  EXPECT_FALSE(status.IsOk());
  ASSERT_EQ(status.Diagnostics().size(), 1U);
  EXPECT_EQ(status.Diagnostics()[0U].code, "dynamic-invalid-state");
}

TEST(DynamicAnalysisState,
     CommitsCompleteZeroCorrectionHistoryAndCopiesTrials) {
  auto fixture = MakeFixture();
  EXPECT_TRUE(fixture.state.IsInitialized());
  EXPECT_FALSE(fixture.state.BeginTrialFromCommitted().HasValue());
  for (std::uint32_t index = 0U; index <= 2U; ++index) {
    auto frame = MakeFrame(*fixture.domain, index);
    ASSERT_TRUE(fixture.state.CommitAcceptedFrame(frame).IsOk());
    frame.reaction[2U] = 123.0;
    EXPECT_DOUBLE_EQ(fixture.state.CommittedState().reaction[2U], -4.0);
    EXPECT_EQ(fixture.state.AcceptedFrames().size(), index + 1U);
  }
  auto trial = fixture.state.BeginTrialFromCommitted();
  ASSERT_TRUE(trial.HasValue());
  trial.Value().generalized_coordinates[0U] = 1.0;
  trial.Value().shell_results[0U].stress[0U].components[0U] = 1.0;
  fixture.state.DiscardTrial();
  EXPECT_DOUBLE_EQ(fixture.state.CommittedState().generalized_coordinates[0U],
                   0.0);
  EXPECT_DOUBLE_EQ(fixture.state.CommittedState()
                       .shell_results[0U]
                       .stress[0U]
                       .components[0U],
                   0.0);
  EXPECT_DOUBLE_EQ(fixture.state.CommittedState().time, 0.5);
  EXPECT_FALSE(
      fixture.state.CommitAcceptedFrame(MakeFrame(*fixture.domain, 3U)).IsOk());
  EXPECT_EQ(fixture.state.AcceptedFrames().size(), 3U);
}

TEST(DynamicAnalysisState,
     RejectsMalformedInventoryScheduleAndPreservesEndpoint) {
  auto fixture = MakeFixture();
  ASSERT_TRUE(
      fixture.state.CommitAcceptedFrame(MakeFrame(*fixture.domain, 0U)).IsOk());
  using Frame = fesa::DynamicAcceptedFrameState;
  const std::vector<std::function<void(Frame&)>> mutations{
      [](Frame& frame) { frame.frame_index = 2U; },
      [](Frame& frame) { frame.increment_index = 2U; },
      [](Frame& frame) { frame.time = 0.26; },
      [](Frame& frame) {
        frame.time = std::numeric_limits<double>::quiet_NaN();
      },
      [](Frame& frame) { frame.is_initial = true; },
      [](Frame& frame) { frame.iteration_count = 101U; },
      [](Frame& frame) { frame.nodal_internal_order[1U] = 0U; },
      [](Frame& frame) {
        frame.nodal_source_ids[1U].source_label_text = "wrong";
      },
      [](Frame& frame) { frame.nodal_triad.pop_back(); },
      [](Frame& frame) {
        frame.nodal_director[0U][2U] = std::numeric_limits<double>::infinity();
      },
      [](Frame& frame) { frame.acceleration = fesa::Vector{23U}; },
      [](Frame& frame) { frame.velocity[0U] = 1.0; },
      [](Frame& frame) { frame.acceleration[0U] = 1.0; },
      [](Frame& frame) { frame.algorithmic_acceleration[0U] = 1.0; },
      [](Frame& frame) { frame.generalized_coordinates[0U] = 1.0; },
      [](Frame& frame) { frame.reaction[2U] = 0.0; },
      [](Frame& frame) { frame.residual[2U] = 0.0; },
      [](Frame& frame) { frame.residual_scale[0U] = -1.0; },
      [](Frame& frame) {
        frame.external_force[0U] = 1.0;
        frame.residual[0U] = -1.0;
        frame.reaction[0U] = -1.0;
      },
      [](Frame& frame) { frame.shell_results.pop_back(); },
      [](Frame& frame) { frame.shell_results[0U].element = 1U; },
      [](Frame& frame) {
        frame.shell_results[0U].natural_coordinates[0U] = 0.0;
      },
      [](Frame& frame) { frame.shell_results[0U].stress[0U].zeta = 0.0; },
      [](Frame& frame) {
        frame.shell_results[0U].stress[0U].components[0U] =
            std::numeric_limits<double>::quiet_NaN();
      },
      [](Frame& frame) {
        frame.kinetic_energy = std::numeric_limits<double>::infinity();
      },
      [](Frame& frame) {
        frame.angular_momentum_rate[0U] =
            std::numeric_limits<double>::quiet_NaN();
      },
      [](Frame& frame) { frame.hidden_constraint = fesa::Vector{1U}; },
      [](Frame& frame) { frame.free_residual_metrics[0U] = -1.0; }};
  for (std::size_t index = 0U; index < mutations.size(); ++index) {
    SCOPED_TRACE(index);
    auto frame = MakeFrame(*fixture.domain, 1U);
    mutations[index](frame);
    const auto status = fixture.state.CommitAcceptedFrame(frame);
    ASSERT_FALSE(status.IsOk());
    ASSERT_EQ(status.Diagnostics().size(), 1U);
    EXPECT_EQ(status.Diagnostics()[0U].code, "dynamic-invalid-state");
    ASSERT_EQ(fixture.state.AcceptedFrames().size(), 1U);
    EXPECT_DOUBLE_EQ(fixture.state.CommittedState().time, 0.0);
    EXPECT_DOUBLE_EQ(fixture.state.CommittedState().reaction[2U], -4.0);
  }
  EXPECT_TRUE(
      fixture.state.CommitAcceptedFrame(MakeFrame(*fixture.domain, 1U)).IsOk());
}

TEST(DynamicAnalysisState, RejectsInvalidInitialStateAndFreeReaction) {
  auto fixture = MakeFixture(MakeDefinition(false));
  auto initial = MakeFrame(*fixture.domain, 0U);
  initial.generalized_coordinates[0U] = 1.0;
  EXPECT_FALSE(fixture.state.CommitAcceptedFrame(initial).IsOk());
  initial = MakeFrame(*fixture.domain, 0U);
  initial.acceleration[0U] = 1.0;
  EXPECT_FALSE(fixture.state.CommitAcceptedFrame(initial).IsOk());
  initial = MakeFrame(*fixture.domain, 0U);
  initial.time = -0.0;
  EXPECT_FALSE(fixture.state.CommitAcceptedFrame(initial).IsOk());
  initial = MakeFrame(*fixture.domain, 0U);
  initial.reaction[0U] = 1.0;
  EXPECT_FALSE(fixture.state.CommitAcceptedFrame(initial).IsOk());
  EXPECT_TRUE(fixture.state.AcceptedFrames().empty());
  EXPECT_TRUE(
      fixture.state.CommitAcceptedFrame(MakeFrame(*fixture.domain, 0U)).IsOk());
}

TEST(DynamicAnalysisState, CreateRejectsInvalidScheduleAndNonzeroBoundary) {
  const std::vector<std::function<void(fesa::ModelDefinition&)>> mutations{
      [](auto& definition) { definition.dynamic_step->increment_count = 0U; },
      [](auto& definition) { definition.dynamic_step->time_period = 0.6; },
      [](auto& definition) {
        definition.dynamic_step->maximum_increments = 1U;
      },
      [](auto& definition) {
        definition.dynamic_step->initial_increment =
            std::numeric_limits<double>::quiet_NaN();
      },
      [](auto& definition) {
        definition.dynamic_step->boundaries[0U].value = 1.0;
      }};
  for (std::size_t index = 0U; index < mutations.size(); ++index) {
    SCOPED_TRACE(index);
    auto definition = MakeDefinition();
    mutations[index](definition);
    auto domain = fesa::Domain::Create(std::move(definition));
    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());
    EXPECT_FALSE(
        fesa::DynamicAnalysisState::Create(domain.Value(), dofs.Value())
            .HasValue());
  }
}

}  // namespace

TEST(DynamicAnalysisState, DrillingEvidenceRejectsPartialCommitAndModeChanges) {
  auto fixture = MakeFixture();
  auto initial = MakeFrame(*fixture.domain, 0U);
  initial.mitc4_drilling = fesa::Mitc4DrillingEvidence{fesa::Mitc4DrillingMode::kDirectSpring};
  ASSERT_TRUE(fixture.state.CommitAcceptedFrame(initial).IsOk());
  for (const int kind : {0, 1, 2, 3, 4}) {
    auto next = MakeFrame(*fixture.domain, 1U);
    next.mitc4_drilling = initial.mitc4_drilling;
    if (kind == 0) next.mitc4_drilling.reset();
    if (kind == 1) next.mitc4_drilling->mode = fesa::Mitc4DrillingMode::kLegacyCorotatedTwist;
    if (kind == 2) next.mitc4_drilling->mode = static_cast<fesa::Mitc4DrillingMode>(99);
    if (kind == 3) next.mitc4_drilling->equilibrium_with_drilling[3U] = 1.0;
    if (kind == 4) next.mitc4_drilling->drilling_moment[0U] = std::numeric_limits<double>::quiet_NaN();
    EXPECT_FALSE(fixture.state.CommitAcceptedFrame(next).IsOk()) << kind;
    ASSERT_EQ(fixture.state.AcceptedFrames().size(), 1U);
    EXPECT_EQ(fixture.state.CommittedState().frame_index, 0U);
    EXPECT_EQ(fixture.state.CommittedState().mitc4_drilling->mode,
              fesa::Mitc4DrillingMode::kDirectSpring);
  }
  auto valid = MakeFrame(*fixture.domain, 1U);
  valid.mitc4_drilling = initial.mitc4_drilling;
  EXPECT_TRUE(fixture.state.CommitAcceptedFrame(valid).IsOk());
}