#include "fesa/analysis/nonlinear_analysis_state.h"

#include <gtest/gtest.h>

#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/model/domain.h"
#include "fesa/model/model_types.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}};

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

Fixture MakeFixture() {
  const std::filesystem::path source{"nonlinear-analysis-state.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "nonlinear-analysis-state-fixture";
  definition.nodes = {
      {{"Shell-A", 10, "10"}, {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.shell_sections = {{"Shell", 0.1, 0U, {source, 6U}}};
  definition.shell_elements = {
      {{"Shell-A", 1, "1"},
       fesa::ShellSourceElementType::kS4,
       {0U, 1U, 2U, 3U},
       0U,
       0U,
       {source, 7U}}};
  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",
                       {},
                       {},
                       fesa::StaticProcedureKind::kNonlinearFullNewton,
                       true,
                       0.1,
                       1.0,
                       0.001,
                       1.0,
                       {source, 8U}}};

  auto domain_result = fesa::Domain::Create(std::move(definition));
  if (!domain_result.HasValue()) {
    throw std::runtime_error{"State 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{"State fixture AnalysisModel construction failed."};
  }
  auto model =
      std::make_unique<fesa::AnalysisModel>(std::move(model_result.Value()));
  auto dofs_result = fesa::DofManager::Create(*model);
  if (!dofs_result.HasValue()) {
    throw std::runtime_error{"State fixture DofManager construction failed."};
  }
  auto dofs =
      std::make_unique<fesa::DofManager>(std::move(dofs_result.Value()));
  return {std::move(domain), std::move(model), std::move(dofs)};
}

fesa::NonlinearStateBaselineInput MakeBaseline(const fesa::Domain& domain) {
  fesa::NonlinearStateBaselineInput baseline{};
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    baseline.node_baselines.push_back(
        {static_cast<fesa::EntityIndex>(node), domain.Nodes()[node].source_id,
         kIdentityTriad, {0.0, 0.0, 1.0}});
  }
  return baseline;
}

fesa::NonlinearAnalysisState MakeState(const Fixture& fixture) {
  auto state = fesa::NonlinearAnalysisState::Create(
      *fixture.domain, *fixture.dofs, MakeBaseline(*fixture.domain));
  if (!state.HasValue()) {
    throw std::runtime_error{"State fixture construction failed."};
  }
  return std::move(state.Value());
}

fesa::ShellResultRow MakeShellRow(
    const fesa::ShellMidsurfaceLocation location,
    const std::array<double, 2> coordinates) {
  return {0U,
          location,
          coordinates,
          kIdentityTriad,
          {},
          {},
          {{{fesa::ShellSectionPosition::kBottom, -1.0, {}},
            {fesa::ShellSectionPosition::kMiddle, 0.0, {}},
            {fesa::ShellSectionPosition::kTop, 1.0, {}}}}};
}

std::vector<fesa::ShellResultRow> MakeShellRows() {
  const double gauss = 1.0 / std::sqrt(3.0);
  return {MakeShellRow(fesa::ShellMidsurfaceLocation::kGp1, {-gauss, -gauss}),
          MakeShellRow(fesa::ShellMidsurfaceLocation::kGp2, {gauss, -gauss}),
          MakeShellRow(fesa::ShellMidsurfaceLocation::kGp3, {gauss, gauss}),
          MakeShellRow(fesa::ShellMidsurfaceLocation::kGp4, {-gauss, gauss})};
}

fesa::AcceptedFrameState MakeFrame(
    const fesa::NonlinearTrialState& committed, const std::uint32_t frame_index,
    const std::uint32_t increment_index, const double load_factor,
    const std::uint32_t iteration_count) {
  return {frame_index,
          increment_index,
          load_factor,
          iteration_count,
          load_factor == 1.0,
          committed.nodal_internal_order,
          committed.nodal_source_ids,
          committed.generalized_coordinates,
          committed.displacement,
          committed.rotation_vector,
          committed.nodal_triad,
          committed.nodal_director,
          committed.internal_force,
          committed.external_force,
          committed.residual,
          committed.reaction,
          MakeShellRows(),
          0.0,
          {}};
}

template <class T, class = void>
struct HasVerificationMetrics : std::false_type {};

template <class T>
struct HasVerificationMetrics<
    T, std::void_t<decltype(std::declval<T&>().verification_metrics)>>
    : std::true_type {};

void ExpectSourceIdEqual(const fesa::SourceEntityId& actual,
                         const fesa::SourceEntityId& expected) {
  EXPECT_EQ(actual.instance_name, expected.instance_name);
  EXPECT_EQ(actual.source_label, expected.source_label);
  EXPECT_EQ(actual.source_label_text, expected.source_label_text);
}

void ExpectVectorEqual(const fesa::Vector& actual, const fesa::Vector& expected) {
  ASSERT_EQ(actual.Size(), expected.Size());
  for (std::size_t index = 0U; index < actual.Size(); ++index) {
    EXPECT_DOUBLE_EQ(actual[index], expected[index]);
  }
}

}  // namespace

// NSNR-STATE-000, NSNR-STATE-002, NSNR-STATE-003, and NSNR-STATE-014
TEST(NonlinearAnalysisState, CreatesFactorZeroBaselineInStableIdentityOrder) {
  const auto fixture = MakeFixture();
  const auto baseline = MakeBaseline(*fixture.domain);
  const auto result = fesa::NonlinearAnalysisState::Create(
      *fixture.domain, *fixture.dofs, baseline);

  ASSERT_TRUE(result.HasValue());
  const auto& committed = result.Value().CommittedState();
  EXPECT_EQ(committed.increment_index, 0U);
  EXPECT_DOUBLE_EQ(committed.load_factor, 0.0);
  EXPECT_TRUE(committed.is_factor_zero_baseline);
  EXPECT_TRUE(result.Value().AcceptedFrames().empty());
  ASSERT_EQ(committed.nodal_internal_order.size(), fixture.domain->Nodes().size());
  ASSERT_EQ(committed.nodal_source_ids.size(), fixture.domain->Nodes().size());
  for (std::size_t node = 0U; node < fixture.domain->Nodes().size(); ++node) {
    EXPECT_EQ(committed.nodal_internal_order[node], node);
    ExpectSourceIdEqual(committed.nodal_source_ids[node],
                         fixture.domain->Nodes()[node].source_id);
  }
  const fesa::Vector* const vectors[] = {
      &committed.generalized_coordinates, &committed.displacement,
      &committed.rotation_vector, &committed.internal_force,
      &committed.external_force, &committed.residual, &committed.reaction};
  for (const auto* vector : vectors) {
    ASSERT_EQ(vector->Size(), fixture.dofs->FullDofCount());
    for (std::size_t index = 0U; index < vector->Size(); ++index) {
      EXPECT_DOUBLE_EQ((*vector)[index], 0.0);
    }
  }
}

// NSNR-STATE-001
TEST(NonlinearAnalysisState, RejectsInvalidBaselineInventoryBeforeStateCreation) {
  const auto fixture = MakeFixture();
  const auto baseline = MakeBaseline(*fixture.domain);
  const std::array<fesa::NonlinearStateBaselineInput, 5> invalid{
      fesa::NonlinearStateBaselineInput{{baseline.node_baselines[0U],
                                         baseline.node_baselines[1U],
                                         baseline.node_baselines[2U]}},
      fesa::NonlinearStateBaselineInput{{baseline.node_baselines[0U],
                                         baseline.node_baselines[1U],
                                         baseline.node_baselines[2U],
                                         baseline.node_baselines[3U],
                                         baseline.node_baselines[3U]}},
      fesa::NonlinearStateBaselineInput{{baseline.node_baselines[1U],
                                         baseline.node_baselines[0U],
                                         baseline.node_baselines[2U],
                                         baseline.node_baselines[3U]}},
      fesa::NonlinearStateBaselineInput{{baseline.node_baselines[0U],
                                         baseline.node_baselines[0U],
                                         baseline.node_baselines[2U],
                                         baseline.node_baselines[3U]}},
      fesa::NonlinearStateBaselineInput{{baseline.node_baselines[0U],
                                         baseline.node_baselines[1U],
                                         baseline.node_baselines[2U],
                                         {3U,
                                          {"Other-Instance", 40, "40"},
                                          kIdentityTriad,
                                          {0.0, 0.0, 1.0}}}}};

  for (const auto& invalid_baseline : invalid) {
    const auto result = fesa::NonlinearAnalysisState::Create(
        *fixture.domain, *fixture.dofs, invalid_baseline);
    EXPECT_FALSE(result.HasValue());
  }
}

// NSNR-STATE-001 and NSNR-STATE-002
TEST(NonlinearAnalysisState, RejectsNonfiniteBaselineAndDofMismatch) {
  const auto fixture = MakeFixture();
  auto nonfinite_triad = MakeBaseline(*fixture.domain);
  nonfinite_triad.node_baselines[0U].triad[0U][0U] =
      std::numeric_limits<double>::quiet_NaN();
  EXPECT_FALSE(fesa::NonlinearAnalysisState::Create(
                   *fixture.domain, *fixture.dofs, nonfinite_triad)
                   .HasValue());

  auto nonfinite_director = MakeBaseline(*fixture.domain);
  nonfinite_director.node_baselines[0U].director[2U] =
      std::numeric_limits<double>::infinity();
  EXPECT_FALSE(fesa::NonlinearAnalysisState::Create(
                   *fixture.domain, *fixture.dofs, nonfinite_director)
                   .HasValue());

  auto empty_domain_result = fesa::Domain::Create(fesa::ModelDefinition{});
  ASSERT_TRUE(empty_domain_result.HasValue());
  const auto dof_mismatch = fesa::NonlinearAnalysisState::Create(
      empty_domain_result.Value(), *fixture.dofs,
      MakeBaseline(empty_domain_result.Value()));
  EXPECT_FALSE(dof_mismatch.HasValue());
}

// NSNR-STATE-001 and NSNR-STATE-003
TEST(NonlinearAnalysisState, PreservesDuplicateLabelsWhenInstancesDiffer) {
  const auto fixture = MakeFixture();
  auto baseline = MakeBaseline(*fixture.domain);
  baseline.node_baselines[1U].source_id = {"Shell-B", 10, "10"};
  auto changed_definition = fesa::ModelDefinition{};
  changed_definition.source_path = "duplicate-instance-labels.inp";
  changed_definition.source_content_identity = "duplicate-instance-labels";
  changed_definition.nodes = {
      {{"Shell-A", 10, "10"}, {0.0, 0.0, 0.0}, {changed_definition.source_path, 1U}},
      {{"Shell-B", 10, "10"}, {1.0, 0.0, 0.0}, {changed_definition.source_path, 2U}}};
  changed_definition.steps = {{"Step-1",
                               {},
                               {},
                               fesa::StaticProcedureKind::kNonlinearFullNewton,
                               true,
                               0.1,
                               1.0,
                               0.001,
                               1.0,
                               {changed_definition.source_path, 3U}}};
  auto domain_result = fesa::Domain::Create(std::move(changed_definition));
  ASSERT_TRUE(domain_result.HasValue());
  auto domain = std::move(domain_result.Value());
  auto model_result = fesa::AnalysisModel::Create(domain);
  ASSERT_TRUE(model_result.HasValue());
  auto dofs_result = fesa::DofManager::Create(model_result.Value());
  ASSERT_TRUE(dofs_result.HasValue());

  baseline.node_baselines.resize(domain.Nodes().size());
  baseline = MakeBaseline(domain);
  const auto result = fesa::NonlinearAnalysisState::Create(
      domain, dofs_result.Value(), baseline);

  ASSERT_TRUE(result.HasValue());
  ASSERT_EQ(result.Value().CommittedState().nodal_source_ids.size(), 2U);
  EXPECT_EQ(result.Value().CommittedState().nodal_source_ids[0U].source_label,
            result.Value().CommittedState().nodal_source_ids[1U].source_label);
  EXPECT_NE(result.Value().CommittedState().nodal_source_ids[0U].instance_name,
            result.Value().CommittedState().nodal_source_ids[1U].instance_name);
}

// NSNR-STATE-004 and NSNR-STATE-012
TEST(NonlinearAnalysisState, StartsIndependentTrialAndDiscardPreservesCommit) {
  const auto fixture = MakeFixture();
  auto state = MakeState(fixture);
  const auto before = state.CommittedState();
  auto trial = state.BeginTrialFromCommitted();

  ASSERT_TRUE(trial.HasValue());
  EXPECT_NE(trial.Value().displacement.Data(),
            state.CommittedState().displacement.Data());
  trial.Value().displacement[0U] = 2.0;
  trial.Value().nodal_director[0U][2U] = 0.5;
  state.DiscardTrial();

  ExpectVectorEqual(state.CommittedState().displacement, before.displacement);
  EXPECT_DOUBLE_EQ(state.CommittedState().nodal_director[0U][2U], 1.0);
}

// NSNR-STATE-005, NSNR-STATE-006, NSNR-STATE-007, NSNR-STATE-008,
// NSNR-STATE-009, and NSNR-STATE-010
TEST(NonlinearAnalysisState, CommitsOnlyStableCompleteAcceptedFrames) {
  const auto fixture = MakeFixture();
  auto state = MakeState(fixture);
  auto first = MakeFrame(state.CommittedState(), 0U, 1U, 0.5, 3U);
  first.internal_force[0U] = 7.0;
  first.external_force[0U] = 2.0;
  first.residual[0U] = 5.0;
  first.reaction[0U] = 5.0;

  ASSERT_TRUE(state.CommitAcceptedFrame(first).IsOk());
  ASSERT_EQ(state.AcceptedFrames().size(), 1U);
  EXPECT_EQ(state.AcceptedFrames()[0U].frame_index, 0U);
  EXPECT_EQ(state.AcceptedFrames()[0U].increment_index, 1U);
  EXPECT_FALSE(state.AcceptedFrames()[0U].is_factor_one);
  EXPECT_EQ(state.CommittedState().increment_index, 1U);
  EXPECT_FALSE(state.CommittedState().is_factor_zero_baseline);
  EXPECT_DOUBLE_EQ(state.CommittedState().residual[0U], 5.0);
  EXPECT_DOUBLE_EQ(state.CommittedState().reaction[0U], 5.0);
  ASSERT_EQ(state.CommittedState().shell_results.size(), 4U);
  EXPECT_EQ(state.CommittedState().shell_results[0U].location,
            fesa::ShellMidsurfaceLocation::kGp1);
  EXPECT_EQ(state.CommittedState().shell_results[3U].location,
            fesa::ShellMidsurfaceLocation::kGp4);

  auto final_frame = MakeFrame(state.CommittedState(), 1U, 2U, 1.0, 2U);
  ASSERT_TRUE(state.CommitAcceptedFrame(final_frame).IsOk());
  ASSERT_EQ(state.AcceptedFrames().size(), 2U);
  EXPECT_EQ(state.AcceptedFrames()[1U].frame_index, 1U);
  EXPECT_TRUE(state.AcceptedFrames()[1U].is_factor_one);
}

// NSNR-STATE-005, NSNR-STATE-009, NSNR-STATE-010, and NSNR-STATE-012
TEST(NonlinearAnalysisState, AtomicallyRejectsMalformedAcceptedFrames) {
  const auto fixture = MakeFixture();
  auto state = MakeState(fixture);
  const auto first = MakeFrame(state.CommittedState(), 0U, 1U, 0.5, 1U);
  ASSERT_TRUE(state.CommitAcceptedFrame(first).IsOk());

  const auto MakeNextFrame = [&state]() {
    return MakeFrame(state.CommittedState(), 1U, 2U, 1.0, 1U);
  };
  const auto ExpectRejected = [&state](const fesa::AcceptedFrameState& frame) {
    const auto before = state.CommittedState();
    const std::size_t frame_count = state.AcceptedFrames().size();
    const auto status = state.CommitAcceptedFrame(frame);

    EXPECT_FALSE(status.IsOk());
    ASSERT_EQ(status.Diagnostics().size(), 1U);
    EXPECT_EQ(status.Diagnostics()[0U].code, "invalid-accepted-state");
    EXPECT_EQ(state.AcceptedFrames().size(), frame_count);
    EXPECT_EQ(state.CommittedState().increment_index, before.increment_index);
    EXPECT_DOUBLE_EQ(state.CommittedState().load_factor, before.load_factor);
    ExpectVectorEqual(state.CommittedState().displacement, before.displacement);
  };

  auto wrong_frame_index = MakeNextFrame();
  wrong_frame_index.frame_index = 2U;
  ExpectRejected(wrong_frame_index);

  auto nonincreasing_factor = MakeFrame(state.CommittedState(), 1U, 2U, 0.5, 1U);
  ExpectRejected(nonincreasing_factor);

  auto mismatched_identity = MakeNextFrame();
  mismatched_identity.nodal_source_ids[0U].instance_name = "Wrong-Instance";
  ExpectRejected(mismatched_identity);

  auto wrong_vector_size = MakeNextFrame();
  wrong_vector_size.generalized_coordinates =
      fesa::Vector{fixture.dofs->FullDofCount() - 1U};
  ExpectRejected(wrong_vector_size);

  auto nonfinite_vector = MakeNextFrame();
  nonfinite_vector.rotation_vector[0U] =
      std::numeric_limits<double>::quiet_NaN();
  ExpectRejected(nonfinite_vector);

  auto residual_mismatch = MakeNextFrame();
  residual_mismatch.residual[0U] = 1.0;
  ExpectRejected(residual_mismatch);

  auto reaction_mismatch = MakeNextFrame();
  reaction_mismatch.reaction[0U] = 1.0;
  ExpectRejected(reaction_mismatch);

  auto invalid_shell_inventory = MakeNextFrame();
  invalid_shell_inventory.shell_results.pop_back();
  ExpectRejected(invalid_shell_inventory);

  auto invalid_factor_flag = MakeNextFrame();
  invalid_factor_flag.is_factor_one = false;
  ExpectRejected(invalid_factor_flag);
}

// NSNR-STATE-011 and NSNR-STATE-012
TEST(NonlinearAnalysisState, AtomicallyRejectsNonfiniteAcceptedFrame) {
  const auto fixture = MakeFixture();
  auto state = MakeState(fixture);
  const auto valid = MakeFrame(state.CommittedState(), 0U, 1U, 0.5, 1U);
  ASSERT_TRUE(state.CommitAcceptedFrame(valid).IsOk());
  const auto before = state.CommittedState();
  auto invalid = MakeFrame(state.CommittedState(), 1U, 2U, 1.0, 1U);
  invalid.physical_strain_energy = std::numeric_limits<double>::quiet_NaN();

  const auto status = state.CommitAcceptedFrame(invalid);

  EXPECT_FALSE(status.IsOk());
  ASSERT_EQ(status.Diagnostics().size(), 1U);
  EXPECT_EQ(status.Diagnostics()[0U].code, "invalid-accepted-state");
  ASSERT_EQ(state.AcceptedFrames().size(), 1U);
  ExpectVectorEqual(state.CommittedState().displacement, before.displacement);
  EXPECT_DOUBLE_EQ(state.CommittedState().load_factor, before.load_factor);
}

// NSNR-STATE-013
TEST(NonlinearAnalysisState, StoresNoUnapprovedVerificationMetrics) {
  EXPECT_FALSE((HasVerificationMetrics<fesa::NonlinearTrialState>::value));
  EXPECT_FALSE((HasVerificationMetrics<fesa::AcceptedFrameState>::value));
}