#include <gtest/gtest.h>

#include <array>
#include <cmath>
#include <filesystem>
#include <limits>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_state.h"

namespace {

fesa::DofManager MakeEmptyDofs() {
  fesa::ModelDefinition definition{};
  definition.source_path = "models/result-records.inp";
  definition.source_content_identity = "fnv1a64:0123456789abcdef";
  definition.steps = {
      {"Step-1", {}, {}, 0.1, 1.0, 0.01, 1.0, {definition.source_path, 10U}}};

  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());
}

fesa::ShellResultRow MakeShellRow(fesa::EntityIndex element,
                                  fesa::ShellMidsurfaceLocation location,
                                  std::array<double, 2> natural_coordinates,
                                  double seed) {
  return {element,
          location,
          natural_coordinates,
          {{{1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}, {0.0, 0.0, 1.0}}},
          {seed + 1.0, seed + 2.0, seed + 3.0, seed + 4.0, seed + 5.0,
           seed + 6.0, seed + 7.0, seed + 8.0},
          {seed + 11.0, seed + 12.0, seed + 13.0, seed + 14.0, seed + 15.0,
           seed + 16.0, seed + 17.0, seed + 18.0},
          {{{fesa::ShellSectionPosition::kBottom,
             -1.0,
             {seed + 21.0, seed + 22.0, seed + 23.0}},
            {fesa::ShellSectionPosition::kMiddle,
             0.0,
             {seed + 24.0, seed + 25.0, seed + 26.0}},
            {fesa::ShellSectionPosition::kTop,
             1.0,
             {seed + 27.0, seed + 28.0, seed + 29.0}}}}};
}

fesa::ShellStateCandidate MakeShellCandidate(
    const std::vector<fesa::EntityIndex>& elements) {
  const double gauss = 1.0 / std::sqrt(3.0);
  const std::array<fesa::ShellMidsurfaceLocation, 4> locations{
      fesa::ShellMidsurfaceLocation::kGp1, fesa::ShellMidsurfaceLocation::kGp2,
      fesa::ShellMidsurfaceLocation::kGp3, fesa::ShellMidsurfaceLocation::kGp4};
  const std::array<std::array<double, 2>, 4> coordinates{
      std::array<double, 2>{-gauss, -gauss},
      std::array<double, 2>{gauss, -gauss}, std::array<double, 2>{gauss, gauss},
      std::array<double, 2>{-gauss, gauss}};

  fesa::ShellStateCandidate candidate{};
  for (const auto element : elements) {
    for (std::size_t point = 0U; point < locations.size(); ++point) {
      candidate.rows.push_back(
          MakeShellRow(element, locations[point], coordinates[point],
                       100.0 * static_cast<double>(element) +
                           10.0 * static_cast<double>(point)));
    }
  }
  candidate.physical_strain_energy = 35.5;
  candidate.equilibrium = {1.0, -2.0, 3.0, -4.0, 5.0, -6.0};
  candidate.verification_metrics = {1.0e-11, 2.0e-11, 3.0e-11};
  return candidate;
}

void ExpectShellStateEquals(const fesa::AnalysisState& state,
                            const std::vector<fesa::ShellResultRow>& rows,
                            double physical_strain_energy,
                            const std::array<double, 6>& equilibrium,
                            const std::array<double, 3>& verification_metrics) {
  ASSERT_EQ(state.ShellResults().size(), rows.size());
  for (std::size_t row = 0U; row < rows.size(); ++row) {
    EXPECT_EQ(state.ShellResults()[row].element, rows[row].element);
    EXPECT_EQ(state.ShellResults()[row].location, rows[row].location);
    EXPECT_EQ(state.ShellResults()[row].natural_coordinates,
              rows[row].natural_coordinates);
    EXPECT_EQ(state.ShellResults()[row].local_frame, rows[row].local_frame);
    EXPECT_EQ(state.ShellResults()[row].generalized_strain,
              rows[row].generalized_strain);
    EXPECT_EQ(state.ShellResults()[row].section_resultant,
              rows[row].section_resultant);
    for (std::size_t position = 0U; position < rows[row].stress.size();
         ++position) {
      EXPECT_EQ(state.ShellResults()[row].stress[position].position,
                rows[row].stress[position].position);
      EXPECT_DOUBLE_EQ(state.ShellResults()[row].stress[position].zeta,
                       rows[row].stress[position].zeta);
      EXPECT_EQ(state.ShellResults()[row].stress[position].components,
                rows[row].stress[position].components);
    }
  }
  EXPECT_DOUBLE_EQ(state.PhysicalStrainEnergy(), physical_strain_energy);
  EXPECT_EQ(state.Equilibrium(), equilibrium);
  EXPECT_EQ(state.VerificationMetrics(), verification_metrics);
}

void ExpectShellCandidateRejectedWithoutMutation(
    fesa::AnalysisState& state,
    const std::vector<fesa::EntityIndex>& expected_elements,
    const fesa::ShellStateCandidate& candidate,
    const fesa::ShellStateCandidate& committed) {
  const auto status = state.CommitShellResults(expected_elements, candidate);
  EXPECT_FALSE(status.IsOk());
  EXPECT_EQ(status.Category(), fesa::FailureCategory::kModel);
  ExpectShellStateEquals(state, committed.rows,
                         committed.physical_strain_energy,
                         committed.equilibrium, committed.verification_metrics);
  ASSERT_TRUE(state.CommitShellResults(expected_elements, committed).IsOk());
}

}  // namespace

TEST(AnalysisState, PreservesStepFrameAndStableRowOrder) {
  const auto dofs = MakeEmptyDofs();
  auto state = fesa::AnalysisState::Create(dofs, {"Step-1", 0U});

  const fesa::EndpointResultRow first_endpoint{2U,
                                               -1,
                                               {"Beam-1", 10, "010"},
                                               {1.0, 2.0, 3.0, 4.0, 5.0, 6.0},
                                               {7.0, 8.0, 9.0, 10.0}};
  const fesa::EndpointResultRow second_endpoint{
      2U,
      1,
      {"Beam-1", 20, "020"},
      {11.0, 12.0, 13.0, 14.0, 15.0, 16.0},
      {17.0, 18.0, 19.0, 20.0}};
  const fesa::GaussResultRow first_gauss{
      2U, 1, {0.1, 0.2, 0.3, 0.4}, {1.1, 1.2, 1.3, 1.4}};
  const fesa::GaussResultRow second_gauss{
      2U, 2, {0.5, 0.6, 0.7, 0.8}, {1.5, 1.6, 1.7, 1.8}};
  const fesa::StressS11Row first_stress{2U, 1, 1U, -0.25, 0.5, 12.5, "input"};
  const fesa::StressS11Row second_stress{2U, 1, 2U, 0.25, -0.5, -7.5, "input"};

  state.EndpointResults().push_back(first_endpoint);
  state.EndpointResults().push_back(second_endpoint);
  state.GaussResults().push_back(first_gauss);
  state.GaussResults().push_back(second_gauss);
  state.StressResults().push_back(first_stress);
  state.StressResults().push_back(second_stress);

  const auto* const endpoint_storage = state.EndpointResults().data();
  const auto* const gauss_storage = state.GaussResults().data();
  const auto* const stress_storage = state.StressResults().data();
  const fesa::AnalysisState& const_state = state;

  EXPECT_EQ(const_state.Identity().step_name, "Step-1");
  EXPECT_EQ(const_state.Identity().frame_index, 0U);
  ASSERT_EQ(const_state.EndpointResults().size(), 2U);
  EXPECT_EQ(const_state.EndpointResults().data(), endpoint_storage);
  EXPECT_EQ(const_state.EndpointResults()[0].element, 2U);
  EXPECT_EQ(const_state.EndpointResults()[0].endpoint, -1);
  EXPECT_EQ(const_state.EndpointResults()[0].node.instance_name, "Beam-1");
  EXPECT_EQ(const_state.EndpointResults()[0].node.source_label, 10);
  EXPECT_EQ(const_state.EndpointResults()[0].node.source_label_text, "010");
  EXPECT_EQ(const_state.EndpointResults()[0].end_action,
            (std::array<double, 6>{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}));
  EXPECT_EQ(const_state.EndpointResults()[1].section_resultant,
            (std::array<double, 4>{17.0, 18.0, 19.0, 20.0}));

  ASSERT_EQ(const_state.GaussResults().size(), 2U);
  EXPECT_EQ(const_state.GaussResults().data(), gauss_storage);
  EXPECT_EQ(const_state.GaussResults()[0].gauss_point, 1);
  EXPECT_EQ(const_state.GaussResults()[1].gauss_point, 2);
  EXPECT_EQ(const_state.GaussResults()[0].generalized_strain,
            (std::array<double, 4>{0.1, 0.2, 0.3, 0.4}));
  EXPECT_EQ(const_state.GaussResults()[1].generalized_resultant,
            (std::array<double, 4>{1.5, 1.6, 1.7, 1.8}));

  ASSERT_EQ(const_state.StressResults().size(), 2U);
  EXPECT_EQ(const_state.StressResults().data(), stress_storage);
  EXPECT_EQ(const_state.StressResults()[0].section_point, 1U);
  EXPECT_DOUBLE_EQ(const_state.StressResults()[0].x1, -0.25);
  EXPECT_DOUBLE_EQ(const_state.StressResults()[0].x2, 0.5);
  EXPECT_DOUBLE_EQ(const_state.StressResults()[0].s11, 12.5);
  EXPECT_EQ(const_state.StressResults()[0].source, "input");
  EXPECT_EQ(const_state.StressResults()[1].section_point, 2U);
}

// MITC4-STATE-001
TEST(AnalysisState, OwnsExactShellRowsInStableElementAndLocationOrder) {
  const auto dofs = MakeEmptyDofs();
  auto state = fesa::AnalysisState::Create(dofs, {"Step-1", 0U});
  const std::vector<fesa::EntityIndex> expected_elements{3U, 7U};
  auto candidate = MakeShellCandidate(expected_elements);

  const auto status = state.CommitShellResults(expected_elements, candidate);

  ASSERT_TRUE(status.IsOk());
  const fesa::AnalysisState& const_state = state;
  ASSERT_EQ(const_state.ShellResults().size(), 8U);
  EXPECT_EQ(const_state.ShellResults()[0].element, 3U);
  EXPECT_EQ(const_state.ShellResults()[0].location,
            fesa::ShellMidsurfaceLocation::kGp1);
  EXPECT_EQ(const_state.ShellResults()[3].location,
            fesa::ShellMidsurfaceLocation::kGp4);
  EXPECT_EQ(const_state.ShellResults()[4].element, 7U);
  EXPECT_EQ(const_state.ShellResults()[4].location,
            fesa::ShellMidsurfaceLocation::kGp1);
  EXPECT_EQ(
      const_state.ShellResults()[0].natural_coordinates,
      (std::array<double, 2>{-1.0 / std::sqrt(3.0), -1.0 / std::sqrt(3.0)}));
  EXPECT_EQ(const_state.ShellResults()[2].generalized_strain,
            (std::array<double, 8>{321.0, 322.0, 323.0, 324.0, 325.0, 326.0,
                                   327.0, 328.0}));
  EXPECT_EQ(const_state.ShellResults()[7].section_resultant,
            (std::array<double, 8>{741.0, 742.0, 743.0, 744.0, 745.0, 746.0,
                                   747.0, 748.0}));
  EXPECT_EQ(const_state.ShellResults()[7].stress[0].position,
            fesa::ShellSectionPosition::kBottom);
  EXPECT_DOUBLE_EQ(const_state.ShellResults()[7].stress[0].zeta, -1.0);
  EXPECT_EQ(const_state.ShellResults()[7].stress[2].components,
            (std::array<double, 3>{757.0, 758.0, 759.0}));
}

// MITC4-STATE-002
TEST(AnalysisState, CommitsFiniteShellGlobalEvidence) {
  const auto dofs = MakeEmptyDofs();
  auto state = fesa::AnalysisState::Create(dofs, {"Step-1", 0U});
  const std::vector<fesa::EntityIndex> expected_elements{5U};
  const auto candidate = MakeShellCandidate(expected_elements);

  const auto status = state.CommitShellResults(expected_elements, candidate);

  ASSERT_TRUE(status.IsOk());
  EXPECT_DOUBLE_EQ(state.PhysicalStrainEnergy(), 35.5);
  EXPECT_EQ(state.Equilibrium(),
            (std::array<double, 6>{1.0, -2.0, 3.0, -4.0, 5.0, -6.0}));
  EXPECT_EQ(state.VerificationMetrics(),
            (std::array<double, 3>{1.0e-11, 2.0e-11, 3.0e-11}));
}

// MITC4-STATE-003
TEST(AnalysisState, InvalidShellCandidatesLeavePriorStateUnchanged) {
  const auto dofs = MakeEmptyDofs();
  auto state = fesa::AnalysisState::Create(dofs, {"Step-1", 0U});
  const std::vector<fesa::EntityIndex> expected_elements{5U};
  const auto committed = MakeShellCandidate(expected_elements);
  ASSERT_TRUE(state.CommitShellResults(expected_elements, committed).IsOk());

  auto invalid_location = MakeShellCandidate(expected_elements);
  invalid_location.rows[0].location = fesa::ShellMidsurfaceLocation::kGp2;
  ExpectShellCandidateRejectedWithoutMutation(state, expected_elements,
                                              invalid_location, committed);

  auto nonfinite = MakeShellCandidate(expected_elements);
  nonfinite.rows[2].generalized_strain[6] =
      (std::numeric_limits<double>::quiet_NaN)();
  ExpectShellCandidateRejectedWithoutMutation(state, expected_elements,
                                              nonfinite, committed);

  auto nonfinite_frame = MakeShellCandidate(expected_elements);
  nonfinite_frame.rows[1].local_frame[2][0] =
      (std::numeric_limits<double>::infinity)();
  ExpectShellCandidateRejectedWithoutMutation(state, expected_elements,
                                              nonfinite_frame, committed);

  auto invalid_section_position = MakeShellCandidate(expected_elements);
  invalid_section_position.rows[3].stress[0].position =
      fesa::ShellSectionPosition::kTop;
  ExpectShellCandidateRejectedWithoutMutation(
      state, expected_elements, invalid_section_position, committed);

  auto nonfinite_global_evidence = MakeShellCandidate(expected_elements);
  nonfinite_global_evidence.verification_metrics[1] =
      (std::numeric_limits<double>::infinity)();
  ExpectShellCandidateRejectedWithoutMutation(
      state, expected_elements, nonfinite_global_evidence, committed);

  auto incomplete_inventory = MakeShellCandidate(expected_elements);
  incomplete_inventory.rows.pop_back();
  ExpectShellCandidateRejectedWithoutMutation(state, expected_elements,
                                              incomplete_inventory, committed);
}