#include "fesa/analysis/hht_dynamic_procedure.h"

#include <gtest/gtest.h>

#include <filesystem>
#include <optional>

#include "fesa/analysis/newton_iteration_log.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/solvers/linear/mkl_pardiso_solver.h"
#include "fesa/solvers/linear/mkl_pardiso_tangent_solver.h"
namespace {
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;
}

TEST(HhtDynamicProcedure, ClampedStepLoadHasPhysicalReactionAndZeroMotion) {
  auto domain = fesa::Domain::Create(MakeDefinition());
  ASSERT_TRUE(domain.HasValue());
  fesa::SerialParallelFor parallel;
  fesa::MklPardisoSolver mass;
  fesa::MklPardisoTangentSolver tangent;
  fesa::NewtonIterationLogWriter log;
  const auto path =
      std::filesystem::temp_directory_path() / "fesa-hht-clamped.log";
  ASSERT_TRUE(log.Open(path).IsOk());
  const fesa::HhtDynamicProcedureExecutor executor(parallel, mass, tangent);
  auto result = executor.Execute(domain.Value(), log);
  ASSERT_TRUE(log.Close().IsOk());
  std::filesystem::remove(path);
  ASSERT_TRUE(result.HasValue())
      << (result.HasValue() ? ""
                            : result.GetStatus().Diagnostics().front().message);
  ASSERT_EQ(result.Value().state.AcceptedFrames().size(), 3U);
  for (const auto& frame : result.Value().state.AcceptedFrames()) {
    EXPECT_EQ(frame.iteration_count, 0U);
    EXPECT_DOUBLE_EQ(frame.reaction[2U], -4.0);
    EXPECT_DOUBLE_EQ(frame.external_force[2U], 4.0);
    EXPECT_DOUBLE_EQ(frame.kinetic_energy, 0.0);
    EXPECT_DOUBLE_EQ(frame.energy_defect, 0.0);
    for (std::size_t i = 0; i < 24U; ++i) {
      EXPECT_DOUBLE_EQ(frame.generalized_coordinates[i], 0.0);
      EXPECT_DOUBLE_EQ(frame.acceleration[i], 0.0);
    }
  }
}

TEST(HhtDynamicProcedure, UnloadedFreeChannelsRemainExactlyAtRest) {
  for (const int fixed_first_dof : {0, 1, 4}) {
    SCOPED_TRACE(fixed_first_dof);
    auto definition = MakeDefinition(false);
    if (fixed_first_dof != 0) {
      definition.dynamic_step->boundaries = {{"All",
                                              fixed_first_dof,
                                              fixed_first_dof + 2,
                                              0.0,
                                              {"dynamic-state.inp", 11U}}};
    }
    auto domain = fesa::Domain::Create(std::move(definition));
    ASSERT_TRUE(domain.HasValue());
    fesa::SerialParallelFor parallel;
    fesa::MklPardisoSolver mass;
    fesa::MklPardisoTangentSolver tangent;
    fesa::NewtonIterationLogWriter log;
    const auto path =
        std::filesystem::temp_directory_path() / "fesa-hht-rest.log";
    ASSERT_TRUE(log.Open(path).IsOk());
    auto result = fesa::HhtDynamicProcedureExecutor(parallel, mass, tangent)
                      .Execute(domain.Value(), log);
    ASSERT_TRUE(log.Close().IsOk());
    std::filesystem::remove(path);
    ASSERT_TRUE(result.HasValue())
        << (result.HasValue()
                ? ""
                : result.GetStatus().Diagnostics().front().message);
    ASSERT_EQ(result.Value().state.AcceptedFrames().size(), 3U);
    for (const auto& frame : result.Value().state.AcceptedFrames()) {
      EXPECT_EQ(frame.iteration_count, 0U);
      for (const auto* values :
           {&frame.generalized_coordinates, &frame.velocity,
            &frame.acceleration, &frame.algorithmic_acceleration,
            &frame.reaction, &frame.residual}) {
        for (std::size_t i = 0; i < values->Size(); ++i)
          EXPECT_DOUBLE_EQ((*values)[i], 0.0);
      }
      EXPECT_DOUBLE_EQ(frame.kinetic_energy, 0.0);
      EXPECT_DOUBLE_EQ(frame.physical_strain_energy, 0.0);
      EXPECT_DOUBLE_EQ(frame.drilling_energy, 0.0);
      EXPECT_DOUBLE_EQ(frame.energy_defect, 0.0);
      for (const double metric : frame.free_residual_metrics)
        EXPECT_DOUBLE_EQ(metric, 0.0);
    }
  }
}

TEST(HhtDynamicProcedure, LargeDirectStepStartsFromAcceptedGeometry) {
  auto definition = MakeDefinition(false);
  definition.materials[0U].youngs_modulus = 1e9;
  definition.dynamic_step->increment_count = 1U;
  definition.dynamic_step->time_period =
      definition.dynamic_step->initial_increment;
  definition.node_sets[0U].node_indices = {0U, 3U};
  definition.node_sets[1U].node_indices = {2U};
  definition.dynamic_step->boundaries = {
      {"All", 1, 6, 0.0, {"dynamic-state.inp", 11U}}};
  definition.dynamic_step->loads = {
      {"Loaded", 3, -40000.0, {"dynamic-state.inp", 12U}}};
  auto domain = fesa::Domain::Create(std::move(definition));
  ASSERT_TRUE(domain.HasValue());
  fesa::SerialParallelFor parallel;
  fesa::MklPardisoSolver mass;
  fesa::MklPardisoTangentSolver tangent;
  fesa::NewtonIterationLogWriter log;
  const auto path =
      std::filesystem::temp_directory_path() / "fesa-hht-large-direct.log";
  ASSERT_TRUE(log.Open(path).IsOk());
  auto result = fesa::HhtDynamicProcedureExecutor(parallel, mass, tangent)
                    .Execute(domain.Value(), log);
  ASSERT_TRUE(log.Close().IsOk());
  ASSERT_TRUE(result.HasValue())
      << (result.HasValue() ? ""
                            : result.GetStatus().Diagnostics().front().message);
  EXPECT_GT(
      std::abs(result.Value().state.AcceptedFrames()[0U].acceleration[14U]),
      1e5);
  EXPECT_LT(result.Value().state.CommittedState().generalized_coordinates[14U],
            0.0);
  for (double metric :
       result.Value().state.CommittedState().free_residual_metrics)
    EXPECT_LE(metric, 1e-10);
  EXPECT_GT(result.Value().state.CommittedState().iteration_count, 0U);
  std::filesystem::remove(path);
}

}  // namespace