#include "fesa/analysis/nonlinear_static_procedure.h"

#include "equilibrium_algorithm.h"
#include "analysis/nonlinear_static_core.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/analysis/newton_iteration_log.h"
#include "fesa/analysis/nonlinear_static_control.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/assembly/surface_pressure_assembler.h"
#include "fesa/constraints/essential_constraint_policy.h"
#include "fesa/elements/element_capabilities.h"
#include "fesa/elements/finite_rotation.h"
#include "fesa/elements/mitc4_shell.h"
#include "fesa/elements/nonlinear_mitc4_shell.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/math/sparse_matrix.h"
#include "fesa/properties/shell_section.h"
#include "fesa/solvers/linear/tangent_solver.h"

#include "assembly/ordered_contribution_internal.h"
#include "math/finite_values.h"
#include "elements/shell_kernel_values.h"
#include "results/mitc4_drilling_evidence.h"

namespace fesa {
namespace {

using shell_kernel_internal::FrameMatrix;
using math_internal::IsFinite;

constexpr std::size_t kDofsPerNode = 6U;
constexpr std::size_t kTranslationsPerNode = 3U;
constexpr std::size_t kShellDofCount = 24U;
constexpr std::uint32_t kMaximumNewtonIterations = 100U;

struct RuntimeShell {
  EntityIndex element_index;
  std::array<std::size_t, kShellDofCount> scatter;
  std::unique_ptr<NonlinearMitc4Shell> kernel;
  std::unique_ptr<ShellTrialElement> element;
};

struct InternalAssembly {
  Vector force;
  SparseMatrix tangent;
  std::vector<ShellResultRow> shell_rows;
  double physical_strain_energy;
  Mitc4DrillingEvidence mitc4_drilling;
};

Status ProcedureFailure(const FailureCategory category, const std::string& code,
                        const SourceLocation& location,
                        const std::string& keyword,
                        const std::string& identity,
                        const std::string& message) {
  return Status::Failure(category, {{Severity::kError, code, location, keyword,
                                     identity, message}});
}

std::string DiagnosticCode(const Status& status) {
  if (status.Diagnostics().empty()) {
    return "nonlinear-procedure-failure";
  }
  return status.Diagnostics().front().code;
}

std::string DiagnosticReason(const Status& status) {
  if (status.Diagnostics().empty()) {
    return "The nonlinear procedure returned no diagnostic.";
  }
  return status.Diagnostics().front().message;
}

NewtonChannelDecision UnevaluatedDecision(const std::string& identity,
                                          const NewtonChannelUnit unit) {
  return {identity,
          0.0,
          0.0,
          NewtonConvergenceBranch::kZeroScale,
          kNewtonZeroScaleThreshold,
          unit,
          false};
}

NewtonIterationLogRecord NewLogRecord(const StepDefinition& step,
                                      const LoadFactorIncrement& increment,
                                      const std::uint32_t iteration) {
  return {step.Name(),
          increment.increment_index,
          increment.load_factor,
          iteration,
          NewtonIterationStage::kEntered,
          {false, NewtonIterationStage::kEntered, {}, {}},
          "entered",
          "pending",
          {},
          UnevaluatedDecision("residual_force", NewtonChannelUnit::kForce),
          UnevaluatedDecision("residual_moment",
                              NewtonChannelUnit::kForceLength),
          UnevaluatedDecision("correction_translation",
                              NewtonChannelUnit::kLength),
          UnevaluatedDecision("correction_rotation",
                              NewtonChannelUnit::kDimensionless),
          false};
}

Status WriteAndFlush(NewtonIterationLogWriter& log,
                     const NewtonIterationLogRecord& record) {
  const Status write_status = log.Write(record);
  if (!write_status.IsOk()) {
    return write_status;
  }
  return log.Flush();
}

Status RecordFailure(NewtonIterationLogWriter& log,
                     NewtonIterationLogRecord* record, const Status& failure,
                     const NewtonIterationStage stage) {
  record->last_completed_stage = stage;
  record->failure = {true, stage, DiagnosticCode(failure),
                     DiagnosticReason(failure)};
  record->transition = "failed";
  record->outcome = "failed";
  record->diagnostic_code = DiagnosticCode(failure);
  const Status log_status = WriteAndFlush(log, *record);
  return log_status.IsOk() ? failure : log_status;
}

Result<DofManager> BuildDofManager(const Domain& domain) {
  Result<AnalysisModel> model = AnalysisModel::Create(domain);
  if (!model.HasValue()) {
    return Result<DofManager>::Failure(model.GetStatus());
  }
  return DofManager::Create(model.Value());
}

Result<NonlinearStateBaselineInput> BuildBaseline(const Domain& domain) {
  if (domain.ShellNodeInitialFrames().size() != domain.Nodes().size()) {
    return Result<NonlinearStateBaselineInput>::Failure(ProcedureFailure(
        FailureCategory::kModel, "invalid-nonlinear-state",
        {domain.SourcePath(), 0U},
        "NODE", domain.SourceContentIdentity(),
        "Nonlinear analysis requires one initial shell frame per node."));
  }

  NonlinearStateBaselineInput baseline;
  baseline.node_baselines.reserve(domain.Nodes().size());
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const ShellNodeInitialFrame& frame = domain.ShellNodeInitialFrames()[node];
    if (frame.node_index != static_cast<EntityIndex>(node)) {
      return Result<NonlinearStateBaselineInput>::Failure(ProcedureFailure(
          FailureCategory::kModel, "invalid-nonlinear-state",
          domain.Nodes()[node].location, "NODE",
          domain.Nodes()[node].source_id.source_label_text,
          "Initial shell frames must follow stable Domain node order."));
    }
    baseline.node_baselines.push_back(
        {static_cast<EntityIndex>(node), domain.Nodes()[node].source_id,
         {frame.tangent_a, frame.tangent_b, frame.director}, frame.director});
  }
  return Result<NonlinearStateBaselineInput>::Success(std::move(baseline));
}

Result<std::vector<RuntimeShell>> BuildRuntimeShells(const Domain& domain,
                                                     const DofManager& dofs) {
  std::vector<RuntimeShell> shells;
  if (domain.Elements().Empty())
    return Result<std::vector<RuntimeShell>>::Failure(ProcedureFailure(
        FailureCategory::kModel, "unsupported-nonlinear-element", {}, "ELEMENT",
        {}, "The nonlinear static procedure requires a MITC4 shell element."));
  shells.reserve(domain.Elements().Size());
  for (std::size_t index = 0U; index < domain.Elements().Size(); ++index) {
    const auto* definition =
        dynamic_cast<const Mitc4ShellDefinition*>(&domain.Elements()[index]);
    if (definition == nullptr) {
      return Result<std::vector<RuntimeShell>>::Failure(ProcedureFailure(
          FailureCategory::kModel, "unsupported-nonlinear-element",
          {}, "ELEMENT",
          domain.Elements()[index].SourceId().source_label_text,
           "The nonlinear static procedure supports MITC4 shell elements "
           "only."));
    }
    if (definition->section_index >= domain.Properties().Size() ||
        definition->material_index >= domain.Materials().Size()) {
      return Result<std::vector<RuntimeShell>>::Failure(ProcedureFailure(
          FailureCategory::kModel, "invalid-nonlinear-static-data",
           definition->location, "ELEMENT",
           definition->source_id.source_label_text,
           "The nonlinear shell element property or material index is "
           "invalid."));
    }

    const auto* section =
        dynamic_cast<const ShellSection*>(
            &domain.Properties()[definition->section_index]);
    const auto* material = dynamic_cast<const LinearElasticMaterial*>(
        &domain.Materials()[definition->material_index]);
    if (section == nullptr || material == nullptr) {
      return Result<std::vector<RuntimeShell>>::Failure(ProcedureFailure(
          FailureCategory::kModel, "incompatible-nonlinear-shell-property",
           definition->location, "ELEMENT",
           definition->source_id.source_label_text,
           "A nonlinear MITC4 shell requires a shell section and linear "
           "elastic material."));
    }

    std::array<const Node*, 4U> nodes{};
    std::array<ShellNodeInitialFrame, 4U> frames{};
    for (std::size_t node = 0U; node < definition->node_indices.size();
         ++node) {
      const EntityIndex node_index = definition->node_indices[node];
      if (node_index >= domain.Nodes().size() ||
          node_index >= domain.ShellNodeInitialFrames().size()) {
        return Result<std::vector<RuntimeShell>>::Failure(ProcedureFailure(
            FailureCategory::kModel, "invalid-shell-director",
            definition->location, "ELEMENT",
            definition->source_id.source_label_text,
            "A nonlinear shell node does not have a stable initial frame."));
      }
      nodes[node] = &domain.Nodes()[node_index];
      frames[node] = domain.ShellNodeInitialFrames()[node_index];
      if (frames[node].node_index != node_index) {
        return Result<std::vector<RuntimeShell>>::Failure(ProcedureFailure(
            FailureCategory::kModel, "invalid-shell-director",
            definition->location, "ELEMENT",
            definition->source_id.source_label_text,
            "A nonlinear shell initial frame does not match its node "
            "identity."));
      }
    }

    Result<NonlinearMitc4Shell> shell =
        NonlinearMitc4Shell::Create(nodes, frames, *section, *material);
    if (!shell.HasValue()) {
      return Result<std::vector<RuntimeShell>>::Failure(shell.GetStatus());
    }
    const EntityIndex element_index = static_cast<EntityIndex>(index);
    auto kernel =
        std::make_unique<NonlinearMitc4Shell>(std::move(shell.Value()));
    auto element = std::make_unique<Mitc4TrialAdapter>(
        ElementDofLayout{definition->SourceId(), definition->NodeIndices(),
                         {DofComponent::kUx, DofComponent::kUy, DofComponent::kUz,
                          DofComponent::kUrx, DofComponent::kUry, DofComponent::kUrz}},
        *kernel);
    auto mapped = dofs.ElementScatter(element->DofLayout());
    if (!mapped.HasValue())
      return Result<std::vector<RuntimeShell>>::Failure(mapped.GetStatus());
    std::array<std::size_t, kShellDofCount> scatter{};
    std::copy(mapped.Value().begin(), mapped.Value().end(), scatter.begin());
    shells.push_back(
        {element_index, scatter, std::move(kernel), std::move(element)});
  }
  return Result<std::vector<RuntimeShell>>::Success(std::move(shells));
}

Status PopulateNodalKinematics(const Domain& domain,
                               NonlinearTrialState* trial) {
  if (trial->generalized_coordinates.Size() !=
      domain.Nodes().size() * kDofsPerNode) {
    return ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data",
        {domain.SourcePath(), 0U},
        "NONLINEAR_STATIC_PROCEDURE", domain.SourceContentIdentity(),
        "Nonlinear trial coordinates do not match the Domain node inventory.");
  }

  trial->displacement = trial->generalized_coordinates;
  trial->rotation_vector = Vector{trial->generalized_coordinates.Size()};
  trial->nodal_triad.clear();
  trial->nodal_director.clear();
  trial->nodal_triad.reserve(domain.Nodes().size());
  trial->nodal_director.reserve(domain.Nodes().size());
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const std::size_t first = node * kDofsPerNode;
    const Vector3 phi{trial->generalized_coordinates[first + 3U],
                      trial->generalized_coordinates[first + 4U],
                      trial->generalized_coordinates[first + 5U]};
    Result<FiniteRotationState> rotation = MakeFiniteRotationState(
        phi, FrameMatrix(domain.ShellNodeInitialFrames()[node]));
    if (!rotation.HasValue()) {
      return rotation.GetStatus();
    }
    for (std::size_t component = 0U; component < kTranslationsPerNode;
         ++component) {
      trial->rotation_vector[first + 3U + component] = phi[component];
    }
    const Matrix& q = rotation.Value().q;
    trial->nodal_triad.push_back(
        {{{q(0U, 0U), q(1U, 0U), q(2U, 0U)},
          {q(0U, 1U), q(1U, 1U), q(2U, 1U)},
          {q(0U, 2U), q(1U, 2U), q(2U, 2U)}}});
    trial->nodal_director.push_back(rotation.Value().director.Components());
  }
  return Status::Ok();
}

NonlinearMitc4TrialDofs ExtractElementTrial(
    const RuntimeShell& shell, const Vector& coordinates) {
  NonlinearMitc4TrialDofs trial{};
  for (std::size_t node = 0U; node < 4U; ++node) {
    const std::size_t first = node * kDofsPerNode;
    trial.displacement[node] =
        {coordinates[shell.scatter[first]],
         coordinates[shell.scatter[first + 1U]],
         coordinates[shell.scatter[first + 2U]]};
    trial.rotation_vector[node] =
        {coordinates[shell.scatter[first + 3U]],
         coordinates[shell.scatter[first + 4U]],
         coordinates[shell.scatter[first + 5U]]};
  }
  return trial;
}

Result<InternalAssembly> AssembleInternal(
    const std::vector<RuntimeShell>& shells, const DofManager& dofs,
    const Domain& domain, const Vector& coordinates, const ParallelFor& parallel_for) {
  std::vector<std::optional<NonlinearMitc4InternalEquilibrium>> local(
      shells.size());
  std::vector<std::optional<Status>> failures(shells.size());
  parallel_for.Execute(shells.size(), [&](const std::size_t index) {
    Result<NonlinearMitc4InternalEquilibrium> result =
        shells[index].element->EvaluateTrial(
            ExtractElementTrial(shells[index], coordinates));
    if (!result.HasValue()) {
      failures[index] = result.GetStatus();
      return;
    }
    local[index] = std::move(result.Value());
  });

  Vector force{dofs.FullDofCount()};
  std::vector<CooContribution> contributions;
  contributions.reserve(shells.size() * kShellDofCount * kShellDofCount);
  std::vector<ShellResultRow> rows;
  rows.reserve(shells.size() * 4U);
  double energy = 0.0;
  Mitc4DrillingEvidence drilling{shells.front().kernel->DrillingMode()};
  for (std::size_t index = 0U; index < shells.size(); ++index) {
    if (failures[index].has_value()) {
      return Result<InternalAssembly>::Failure(*failures[index]);
    }
    if (!local[index].has_value() ||
        !IsFinite(local[index]->internal_residual) ||
        !IsFinite(local[index]->internal_tangent) ||
        !std::isfinite(local[index]->physical_strain_energy)) {
      return Result<InternalAssembly>::Failure(ProcedureFailure(
          FailureCategory::kSolver, "nonfinite-newton-data", {},
          "NONLINEAR_STATIC_PROCEDURE", std::to_string(index),
          "A nonlinear shell internal residual or tangent is nonfinite."));
    }
    if (shells[index].kernel->DrillingMode() != drilling.mode ||
        local[index]->drilling_force.Size() != kShellDofCount ||
        !IsFinite(local[index]->drilling_force))
      return Result<InternalAssembly>::Failure(ProcedureFailure(
          FailureCategory::kSolver, "invalid-shell-drilling-evidence", {},
          "NONLINEAR_STATIC_PROCEDURE", {}, "Drilling mode or force is invalid."));
    const auto element_trial = ExtractElementTrial(shells[index], coordinates);
    for (std::size_t node = 0U; node < 4U; ++node) {
      const auto first = 6U * node;
      const auto global_node = shells[index].scatter[first] / 6U;
      const auto& values = local[index]->drilling_force;
      const Vector3 drilling_force{values[first], values[first + 1U], values[first + 2U]};
      const auto spatial = SpatialMomentSO3(element_trial.rotation_vector[node],
          {values[first + 3U], values[first + 4U], values[first + 5U]});
      if (!spatial.HasValue())
        return Result<InternalAssembly>::Failure(spatial.GetStatus());
      const Vector3 position = Vector3(domain.Nodes()[global_node].coordinates) +
                               element_trial.displacement[node];
      const Vector3 moment = position.Cross(drilling_force) + spatial.Value();
      for (std::size_t c = 0U; c < 3U; ++c)
        drilling.drilling_moment[c] += moment[c];
    }
    for (std::size_t row = 0U; row < kShellDofCount; ++row) {
      force[shells[index].scatter[row]] += local[index]->internal_residual[row];
      assembly_internal::AppendLocalMatrixRow(
          local[index]->internal_tangent, shells[index].scatter, row, index,
          contributions);
    }
    for (const NonlinearMitc4RecoveryPoint& point :
         local[index]->recovery.gauss_points) {
      rows.push_back({shells[index].element_index,
                      point.location,
                      point.natural_coordinates,
                      {point.local_frame.e1, point.local_frame.e2,
                       point.local_frame.e3},
                      point.generalized_strain,
                      point.section_resultant,
                      point.stress});
    }
    energy += local[index]->physical_strain_energy;
  }
  if (!IsFinite(force) || !std::isfinite(energy) ||
      !IsFinite(drilling.drilling_moment)) {
    return Result<InternalAssembly>::Failure(ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {},
        "NONLINEAR_STATIC_PROCEDURE", {},
        "Global nonlinear internal force or energy is nonfinite."));
  }

  Result<SparseMatrix> tangent = SparseMatrix::FromCoo(
      dofs.FullDofCount(), dofs.FullDofCount(), std::move(contributions),
      dofs.GetSparsePattern());
  if (!tangent.HasValue()) {
    return Result<InternalAssembly>::Failure(tangent.GetStatus());
  }
  return Result<InternalAssembly>::Success(
      {std::move(force), std::move(tangent.Value()), std::move(rows), energy, drilling});
}

Result<Vector> Difference(const Vector& left, const Vector& right) {
  if (left.Size() != right.Size()) {
    return Result<Vector>::Failure(ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {},
        "NONLINEAR_STATIC_PROCEDURE", {},
        "Nonlinear internal and external vectors must have equal dimensions."));
  }
  Vector result{left.Size()};
  for (std::size_t index = 0U; index < result.Size(); ++index) {
    result[index] = left[index] - right[index];
  }
  if (!IsFinite(result)) {
    return Result<Vector>::Failure(ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {},
        "NONLINEAR_STATIC_PROCEDURE", {},
        "Nonlinear residual must remain finite."));
  }
  return Result<Vector>::Success(std::move(result));
}

Result<SparseMatrix> Difference(const SparseMatrix& internal,
                                const SparseMatrix& external,
                                const DofManager& dofs) {
  if (internal.Rows() != external.Rows() ||
      internal.Columns() != external.Columns() ||
      internal.RowOffsets() != external.RowOffsets() ||
      internal.ColumnIndices() != external.ColumnIndices()) {
    return Result<SparseMatrix>::Failure(ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {},
        "NONLINEAR_STATIC_PROCEDURE", {},
        "Nonlinear internal and external tangent patterns must match."));
  }
  std::vector<CooContribution> contributions;
  contributions.reserve(internal.Values().size());
  for (std::size_t row = 0U; row < internal.Rows(); ++row) {
    for (std::size_t offset = internal.RowOffsets()[row];
         offset < internal.RowOffsets()[row + 1U]; ++offset) {
      const double value =
          internal.Values()[offset] - external.Values()[offset];
      if (!std::isfinite(value)) {
        return Result<SparseMatrix>::Failure(ProcedureFailure(
            FailureCategory::kSolver, "nonfinite-newton-data", {},
            "NONLINEAR_STATIC_PROCEDURE", {},
            "Nonlinear tangent must remain finite."));
      }
      contributions.push_back(
          {row, internal.ColumnIndices()[offset], value, row, offset});
    }
  }
  return SparseMatrix::FromCoo(internal.Rows(), internal.Columns(),
                               std::move(contributions),
                               dofs.GetSparsePattern());
}

/// @brief Adds fixed and follower external candidates in stable CSR order.
Result<GlobalFixedLoadContribution> CombineExternalContributions(
    const GlobalFixedLoadContribution& fixed,
    const SurfacePressureAssemblyCandidate& pressure, const DofManager& dofs) {
  if (fixed.generalized_external_force.Size() !=
          pressure.full_external_force.Size() ||
      fixed.generalized_external_tangent.Rows() !=
          pressure.full_external_tangent.Rows() ||
      fixed.generalized_external_tangent.Columns() !=
          pressure.full_external_tangent.Columns() ||
      fixed.generalized_external_tangent.RowOffsets() !=
          pressure.full_external_tangent.RowOffsets() ||
      fixed.generalized_external_tangent.ColumnIndices() !=
          pressure.full_external_tangent.ColumnIndices()) {
    return Result<GlobalFixedLoadContribution>::Failure(ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {},
        "NONLINEAR_STATIC_PROCEDURE", {},
        "Nonlinear fixed and pressure external candidates must share the "
        "stable full-space dimension and pattern."));
  }

  Vector force{fixed.generalized_external_force.Size()};
  for (std::size_t index = 0U; index < force.Size(); ++index) {
    force[index] = fixed.generalized_external_force[index] +
                   pressure.full_external_force[index];
    if (!std::isfinite(force[index])) {
      return Result<GlobalFixedLoadContribution>::Failure(ProcedureFailure(
          FailureCategory::kSolver, "nonfinite-newton-data", {},
          "NONLINEAR_STATIC_PROCEDURE", std::to_string(index),
          "Nonlinear fixed and pressure force composition is nonfinite."));
    }
  }

  std::vector<CooContribution> contributions;
  contributions.reserve(fixed.generalized_external_tangent.Values().size());
  for (std::size_t row = 0U; row < fixed.generalized_external_tangent.Rows();
       ++row) {
    for (std::size_t offset =
             fixed.generalized_external_tangent.RowOffsets()[row];
         offset < fixed.generalized_external_tangent.RowOffsets()[row + 1U];
         ++offset) {
      const double tangent =
          fixed.generalized_external_tangent.Values()[offset] +
          pressure.full_external_tangent.Values()[offset];
      if (!std::isfinite(tangent)) {
        return Result<GlobalFixedLoadContribution>::Failure(ProcedureFailure(
            FailureCategory::kSolver, "nonfinite-newton-data", {},
            "NONLINEAR_STATIC_PROCEDURE", std::to_string(row),
            "Nonlinear fixed and pressure tangent composition is nonfinite."));
      }
      contributions.push_back(
          {row, fixed.generalized_external_tangent.ColumnIndices()[offset],
           tangent, row, offset});
    }
  }
  Result<SparseMatrix> tangent =
      SparseMatrix::FromCoo(fixed.generalized_external_tangent.Rows(),
                            fixed.generalized_external_tangent.Columns(),
                            std::move(contributions), dofs.GetSparsePattern());
  if (!tangent.HasValue()) {
    return Result<GlobalFixedLoadContribution>::Failure(tangent.GetStatus());
  }

  std::vector<Diagnostic> diagnostics = fixed.diagnostics;
  diagnostics.insert(diagnostics.end(), pressure.diagnostics.begin(),
                     pressure.diagnostics.end());
  return Result<GlobalFixedLoadContribution>::Success(
      {std::move(force), std::move(tangent.Value()), std::move(diagnostics)});
}

Vector SelectFreeComponents(const Vector& full, const DofManager& dofs,
                            const bool rotations) {
  std::size_t count = 0U;
  for (const std::size_t full_dof : dofs.FreeDofs()) {
    if ((full_dof % kDofsPerNode >= kTranslationsPerNode) == rotations) {
      ++count;
    }
  }
  Vector selected{count};
  std::size_t output = 0U;
  for (const std::size_t full_dof : dofs.FreeDofs()) {
    if ((full_dof % kDofsPerNode >= kTranslationsPerNode) == rotations) {
      selected[output++] = full[full_dof];
    }
  }
  return selected;
}

Vector SelectCorrectionComponents(const Vector& correction,
                                  const DofManager& dofs,
                                  const bool rotations) {
  std::size_t count = 0U;
  for (const std::size_t full_dof : dofs.FreeDofs()) {
    if ((full_dof % kDofsPerNode >= kTranslationsPerNode) == rotations) {
      ++count;
    }
  }
  Vector selected{count};
  std::size_t output = 0U;
  for (std::size_t equation = 0U; equation < dofs.FreeDofs().size();
       ++equation) {
    if ((dofs.FreeDofs()[equation] % kDofsPerNode >=
         kTranslationsPerNode) == rotations) {
      selected[output++] = correction[equation];
    }
  }
  return selected;
}

Vector CurrentIncrementComponents(const Vector& current, const Vector& start,
                                  const DofManager& dofs,
                                  const bool rotations) {
  Vector difference{current.Size()};
  for (std::size_t index = 0U; index < difference.Size(); ++index) {
    difference[index] = current[index] - start[index];
  }
  return SelectFreeComponents(difference, dofs, rotations);
}

/// @brief Computes accepted-state applied-plus-reaction balance evidence.
Result<std::array<double, 6>> GlobalEquilibrium(
    const Domain& domain, const DofManager& dofs,
    const NonlinearTrialState& trial) {
  const std::size_t full_dof_count = domain.Nodes().size() * kDofsPerNode;
  if (trial.displacement.Size() != full_dof_count ||
      trial.rotation_vector.Size() != full_dof_count ||
      trial.external_force.Size() != full_dof_count ||
      trial.residual.Size() != full_dof_count) {
    return Result<std::array<double, 6>>::Failure(ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data",
        {domain.SourcePath(), 0U}, "NONLINEAR_STATIC_PROCEDURE",
        domain.SourceContentIdentity(),
        "Nonlinear equilibrium vectors do not match the Domain node inventory."));
  }

  std::vector<bool> constrained(full_dof_count, false);
  for (const std::size_t full_dof : dofs.ConstrainedDofs()) {
    if (full_dof >= constrained.size()) {
      return Result<std::array<double, 6>>::Failure(ProcedureFailure(
          FailureCategory::kSolver, "nonfinite-newton-data",
          {domain.SourcePath(), 0U}, "NONLINEAR_STATIC_PROCEDURE",
          domain.SourceContentIdentity(),
          "Nonlinear constraint identity is outside the full DOF range."));
    }
    constrained[full_dof] = true;
  }

  Vector3 applied_force_resultant;
  Vector3 applied_moment_resultant;
  Vector3 reaction_force_resultant;
  Vector3 reaction_moment_resultant;
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const std::size_t first = node * kDofsPerNode;
    const Vector3 position = Vector3{domain.Nodes()[node].coordinates} +
                             Vector3{trial.displacement[first],
                                     trial.displacement[first + 1U],
                                     trial.displacement[first + 2U]};
    const Vector3 phi{trial.rotation_vector[first + 3U],
                      trial.rotation_vector[first + 4U],
                      trial.rotation_vector[first + 5U]};
    const Vector3 applied_force{trial.external_force[first],
                                trial.external_force[first + 1U],
                                trial.external_force[first + 2U]};
    const Vector3 reaction_force{
        constrained[first] ? trial.residual[first] : 0.0,
        constrained[first + 1U] ? trial.residual[first + 1U] : 0.0,
        constrained[first + 2U] ? trial.residual[first + 2U] : 0.0};
    const std::array<double, 3> applied_generalized{
        trial.external_force[first + 3U], trial.external_force[first + 4U],
        trial.external_force[first + 5U]};
    const std::array<double, 3> reaction_generalized{
        constrained[first + 3U] ? trial.residual[first + 3U] : 0.0,
        constrained[first + 4U] ? trial.residual[first + 4U] : 0.0,
        constrained[first + 5U] ? trial.residual[first + 5U] : 0.0};
    const auto applied = SpatialMomentSO3(phi, Vector3(applied_generalized));
    const auto reaction = SpatialMomentSO3(phi, Vector3(reaction_generalized));
    if (!applied.HasValue() || !reaction.HasValue())
      return Result<std::array<double, 6>>::Failure(ProcedureFailure(
          FailureCategory::kSolver, "nonfinite-newton-data",
          {domain.SourcePath(), 0U}, "NONLINEAR_STATIC_PROCEDURE",
          domain.SourceContentIdentity(),
          "Nonlinear global equilibrium balance is nonfinite."));
    const Vector3 applied_moment = applied.Value();
    const Vector3 reaction_moment = reaction.Value();
    applied_force_resultant = applied_force_resultant + applied_force;
    applied_moment_resultant =
        applied_moment_resultant + position.Cross(applied_force) + applied_moment;
    reaction_force_resultant = reaction_force_resultant + reaction_force;
    reaction_moment_resultant = reaction_moment_resultant +
                                position.Cross(reaction_force) + reaction_moment;
  }
  const Vector3 force_balance =
      applied_force_resultant + reaction_force_resultant;
  const Vector3 moment_balance =
      applied_moment_resultant + reaction_moment_resultant;
  const std::array<double, 6> equilibrium{
      force_balance.X(), force_balance.Y(), force_balance.Z(),
      moment_balance.X(), moment_balance.Y(), moment_balance.Z()};
  if (!std::all_of(equilibrium.begin(), equilibrium.end(),
                   [](const double value) { return std::isfinite(value); })) {
    return Result<std::array<double, 6>>::Failure(ProcedureFailure(
        FailureCategory::kSolver, "nonfinite-newton-data",
        {domain.SourcePath(), 0U}, "NONLINEAR_STATIC_PROCEDURE",
        domain.SourceContentIdentity(),
        "Nonlinear global equilibrium balance is nonfinite."));
  }
  return Result<std::array<double, 6>>::Success(equilibrium);
}

Status ApplyIncrementConstraints(const StepDefinition& step,
                                 const LoadFactorIncrement& increment,
                                 const DofManager& dofs,
                                 NonlinearTrialState* trial) {
  Result<Vector> source_targets =
      ApplyAbsolutePrescribedDisplacementTarget(step, increment.load_factor);
  if (!source_targets.HasValue()) {
    return source_targets.GetStatus();
  }
  for (std::size_t index = 0U; index < dofs.ConstrainedDofs().size();
       ++index) {
    const double target =
        increment.load_factor * dofs.PrescribedValues()[index];
    if (!std::isfinite(target)) {
      return ProcedureFailure(
          FailureCategory::kSolver, "nonfinite-newton-data", step.Location(),
          "BOUNDARY", step.Name(),
          "A nonlinear prescribed displacement target is nonfinite.");
    }
    trial->generalized_coordinates[dofs.ConstrainedDofs()[index]] = target;
  }
  return Status::Ok();
}

AcceptedFrameState AcceptedFrame(const std::uint32_t frame_index,
                                 const NonlinearTrialState& trial,
                                 const LoadFactorIncrement& increment,
                                 const std::uint32_t iteration) {
  return {frame_index,
          increment.increment_index,
          increment.load_factor,
          iteration,
          increment.is_factor_one,
          trial.nodal_internal_order,
          trial.nodal_source_ids,
          trial.generalized_coordinates,
          trial.displacement,
          trial.rotation_vector,
          trial.nodal_triad,
          trial.nodal_director,
          trial.internal_force,
          trial.external_force,
          trial.residual,
          trial.reaction,
          trial.shell_results,
          trial.physical_strain_energy,
          trial.global_equilibrium, trial.mitc4_drilling};
}

}  // namespace

NonlinearStaticProcedureExecutor::NonlinearStaticProcedureExecutor(
    const ParallelFor& parallel_for, TangentSolver& tangent_solver)
    : parallel_for_(parallel_for), tangent_solver_(tangent_solver) {}

Result<NonlinearStaticExecutionResult>
NonlinearStaticProcedureExecutor::Execute(
    const Domain& domain, NewtonIterationLogWriter& log) const {
  NewtonStepAlgorithm algorithm;
  return ExecuteNonlinearStaticCore(domain, log, parallel_for_, tangent_solver_,
                                     algorithm);
}

Result<NonlinearStaticExecutionResult> ExecuteNonlinearStaticCore(
    const Domain& domain, NewtonIterationLogWriter& log,
    const ParallelFor& parallel_for, TangentSolver& tangent_solver,
    EquilibriumAlgorithm& algorithm) {
  if (domain.Steps().Size() != 1U || !domain.Steps()[0U].IsNonlinearStatic()) {
    return Result<NonlinearStaticExecutionResult>::Failure(ProcedureFailure(
        FailureCategory::kInput, "invalid-nonlinear-static-data",
        {domain.SourcePath(), 0U}, "STATIC", domain.SourceContentIdentity(),
        "Nonlinear procedure execution requires exactly one NLGEOM static "
        "step."));
  }
  const StepDefinition& step = domain.Steps()[0U];

  Result<DofManager> dofs = BuildDofManager(domain);
  if (!dofs.HasValue()) {
    return Result<NonlinearStaticExecutionResult>::Failure(dofs.GetStatus());
  }
  Result<AnalysisModel> analysis_model = AnalysisModel::Create(domain);
  if (!analysis_model.HasValue()) {
    return Result<NonlinearStaticExecutionResult>::Failure(
        analysis_model.GetStatus());
  }
  Result<NonlinearStateBaselineInput> baseline = BuildBaseline(domain);
  if (!baseline.HasValue()) {
    return Result<NonlinearStaticExecutionResult>::Failure(
        baseline.GetStatus());
  }
  Result<NonlinearAnalysisState> state =
      NonlinearAnalysisState::Create(domain, dofs.Value(), baseline.Value());
  if (!state.HasValue()) {
    return Result<NonlinearStaticExecutionResult>::Failure(state.GetStatus());
  }
  Result<std::vector<RuntimeShell>> shells =
      BuildRuntimeShells(domain, dofs.Value());
  if (!shells.HasValue()) {
    return Result<NonlinearStaticExecutionResult>::Failure(shells.GetStatus());
  }
  Result<std::vector<LoadFactorIncrement>> increments =
      BuildFixedLoadFactorSequence(step);
  if (!increments.HasValue()) {
    return Result<NonlinearStaticExecutionResult>::Failure(
        increments.GetStatus());
  }

  const EssentialConstraintPolicy constraints;
  for (const LoadFactorIncrement& increment : increments.Value()) {
    Result<NonlinearTrialState> trial = state.Value().BeginTrialFromCommitted();
    if (!trial.HasValue()) {
      return Result<NonlinearStaticExecutionResult>::Failure(trial.GetStatus());
    }
    const Vector increment_start = trial.Value().generalized_coordinates;
    Status constraint_status =
        ApplyIncrementConstraints(step, increment, dofs.Value(),
                                  &trial.Value());
    if (!constraint_status.IsOk()) {
      return Result<NonlinearStaticExecutionResult>::Failure(constraint_status);
    }

    struct ShellStepCandidate {
      AcceptedFrameState frame;
      NewtonIterationLogRecord record;
    };
    const auto next_frame_index =
        static_cast<std::uint32_t>(state.Value().AcceptedFrames().size());
    const auto report_failure = [&](NewtonIterationLogRecord* record,
                                    const Status& failure,
                                    NewtonIterationStage stage) {
      return RecordFailure(log, record, failure, stage);
    };
    const auto write_iteration = [&](const NewtonIterationLogRecord& record) {
      return WriteAndFlush(log, record);
    };
    ShellFullNewtonStep<ShellStepCandidate> iteration_step(
        [&](std::uint32_t iteration)
            -> Result<std::optional<ShellStepCandidate>> {
          NewtonIterationLogRecord record =
              NewLogRecord(step, increment, iteration);
          trial.Value().increment_index = increment.increment_index;
          trial.Value().load_factor = increment.load_factor;
          trial.Value().is_factor_zero_baseline = false;

          Status kinematics = PopulateNodalKinematics(domain, &trial.Value());
          if (!kinematics.IsOk()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, kinematics,
                               NewtonIterationStage::kKinematicsEvaluated));
          }
          record.last_completed_stage =
              NewtonIterationStage::kKinematicsEvaluated;

          Result<InternalAssembly> internal = AssembleInternal(
              shells.Value(), dofs.Value(), domain,
              trial.Value().generalized_coordinates, parallel_for);
          if (!internal.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, internal.GetStatus(),
                               NewtonIterationStage::kInternalAssembled));
          }
          record.last_completed_stage =
              NewtonIterationStage::kInternalAssembled;

          Result<GlobalFixedLoadContribution> fixed =
              AssembleGlobalFixedLoadContribution(
                  domain, dofs.Value(), trial.Value().generalized_coordinates,
                  increment.load_factor);
          if (!fixed.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, fixed.GetStatus(),
                               NewtonIterationStage::kExternalAssembled));
          }
          Result<SurfacePressureAssemblyCandidate> pressure =
              AssembleSurfacePressure(
                  analysis_model.Value(), dofs.Value(),
                  trial.Value().generalized_coordinates, increment.load_factor,
                  SurfacePressureConfiguration::kCurrent, parallel_for);
          if (!pressure.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, pressure.GetStatus(),
                               NewtonIterationStage::kExternalAssembled));
          }
          Result<GlobalFixedLoadContribution> external =
              CombineExternalContributions(fixed.Value(), pressure.Value(),
                                           dofs.Value());
          if (!external.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, external.GetStatus(),
                               NewtonIterationStage::kExternalAssembled));
          }
          record.last_completed_stage =
              NewtonIterationStage::kExternalAssembled;

          Result<Vector> residual =
              Difference(internal.Value().force,
                         external.Value().generalized_external_force);
          if (!residual.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, residual.GetStatus(),
                               NewtonIterationStage::kExternalAssembled));
          }
          Result<SparseMatrix> tangent = Difference(
              internal.Value().tangent,
              external.Value().generalized_external_tangent, dofs.Value());
          if (!tangent.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, tangent.GetStatus(),
                               NewtonIterationStage::kExternalAssembled));
          }
          trial.Value().internal_force = std::move(internal.Value().force);
          trial.Value().external_force =
              std::move(external.Value().generalized_external_force);
          trial.Value().residual = std::move(residual.Value());
          trial.Value().reaction = trial.Value().residual;
          trial.Value().shell_results = std::move(internal.Value().shell_rows);
          trial.Value().physical_strain_energy =
              internal.Value().physical_strain_energy;
          Result<std::array<double, 6>> global_equilibrium =
              GlobalEquilibrium(domain, dofs.Value(), trial.Value());
          if (!global_equilibrium.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, global_equilibrium.GetStatus(),
                               NewtonIterationStage::kExternalAssembled));
          }
          trial.Value().global_equilibrium = global_equilibrium.Value();
          trial.Value().mitc4_drilling = internal.Value().mitc4_drilling;
          results_internal::CompleteDrillingEvidence(
              trial.Value().global_equilibrium, &*trial.Value().mitc4_drilling);

          Result<PartitionedStiffness> partition =
              constraints.Partition(tangent.Value(), dofs.Value());
          if (!partition.HasValue()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, partition.GetStatus(),
                               NewtonIterationStage::kExternalAssembled));
          }
          const Status factorization =
              tangent_solver.FactorizeFresh(partition.Value().k_ff);
          if (!factorization.IsOk()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, factorization,
                               NewtonIterationStage::kSolvedCorrection));
          }

          Vector correction_rhs =
              constraints.GatherFree(trial.Value().residual, dofs.Value());
          for (std::size_t index = 0U; index < correction_rhs.Size(); ++index) {
            correction_rhs[index] = -correction_rhs[index];
          }
          Vector correction{dofs.Value().FreeDofCount()};
          const Status solve =
              tangent_solver.SolveCorrection(correction_rhs, correction);
          if (!solve.IsOk()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, solve,
                               NewtonIterationStage::kSolvedCorrection));
          }
          record.last_completed_stage = NewtonIterationStage::kSolvedCorrection;

          const NewtonConvergenceResult convergence = EvaluateNewtonConvergence(
              {{"residual_force",
                SelectFreeComponents(trial.Value().residual, dofs.Value(),
                                     false),
                SelectFreeComponents(trial.Value().internal_force, dofs.Value(),
                                     false),
                SelectFreeComponents(trial.Value().external_force, dofs.Value(),
                                     false)},
               {"residual_moment",
                SelectFreeComponents(trial.Value().residual, dofs.Value(),
                                     true),
                SelectFreeComponents(trial.Value().internal_force, dofs.Value(),
                                     true),
                SelectFreeComponents(trial.Value().external_force, dofs.Value(),
                                     true)},
               {"correction_translation",
                SelectCorrectionComponents(correction, dofs.Value(), false),
                CurrentIncrementComponents(
                    trial.Value().generalized_coordinates, increment_start,
                    dofs.Value(), false)},
               {"correction_rotation",
                SelectCorrectionComponents(correction, dofs.Value(), true),
                CurrentIncrementComponents(
                    trial.Value().generalized_coordinates, increment_start,
                    dofs.Value(), true)}});
          record.force_residual = convergence.force_residual;
          record.moment_residual = convergence.moment_residual;
          record.translation_correction = convergence.translation_correction;
          record.rotation_correction = convergence.rotation_correction;
          record.combined_passed = convergence.combined_passed;
          record.last_completed_stage =
              NewtonIterationStage::kConvergenceEvaluated;

          if (convergence.combined_passed) {
            return Result<std::optional<ShellStepCandidate>>::Success(
                ShellStepCandidate{AcceptedFrame(next_frame_index,
                                                 trial.Value(), increment,
                                                 iteration),
                                   record});
          }

          if (iteration == kMaximumNewtonIterations) {
            const Status failure = ProcedureFailure(
                FailureCategory::kSolver, "nonlinear-nonconvergence",
                step.Location(), "STATIC", step.Name(),
                "The nonlinear procedure did not converge within 100 "
                "iterations.");
            return Result<std::optional<ShellStepCandidate>>::Failure(
                report_failure(&record, failure,
                               NewtonIterationStage::kConvergenceEvaluated));
          }

          for (std::size_t equation = 0U;
               equation < dofs.Value().FreeDofs().size(); ++equation) {
            const std::size_t full_dof = dofs.Value().FreeDofs()[equation];
            trial.Value().generalized_coordinates[full_dof] +=
                correction[equation];
            if (!std::isfinite(
                    trial.Value().generalized_coordinates[full_dof])) {
              const Status failure = ProcedureFailure(
                  FailureCategory::kSolver, "nonfinite-newton-data", {},
                  "NONLINEAR_STATIC_PROCEDURE", std::to_string(full_dof),
                  "A nonlinear free-coordinate correction is nonfinite.");
              return Result<std::optional<ShellStepCandidate>>::Failure(
                  report_failure(&record, failure,
                                 NewtonIterationStage::kSolvedCorrection));
            }
          }
          record.transition = "continue";
          record.outcome = "not-converged";
          const Status log_status = write_iteration(record);
          if (!log_status.IsOk()) {
            return Result<std::optional<ShellStepCandidate>>::Failure(
                log_status);
          }
          return Result<std::optional<ShellStepCandidate>>::Success(
              std::nullopt);
        },
        kMaximumNewtonIterations);
    const Status solved = iteration_step.Solve(algorithm);
    if (!solved.IsOk())
      return Result<NonlinearStaticExecutionResult>::Failure(solved);
    auto candidate = iteration_step.TakeCandidate();
    auto& record = candidate.record;
    const Status committed = state.Value().CommitAcceptedFrame(candidate.frame);
    if (!committed.IsOk()) {
      return Result<NonlinearStaticExecutionResult>::Failure(RecordFailure(
          log, &record, committed, NewtonIterationStage::kAccepted));
    }
    record.last_completed_stage = NewtonIterationStage::kAccepted;
    record.transition = "accepted";
    record.outcome = "accepted";
    const Status log_status = WriteAndFlush(log, record);
    if (!log_status.IsOk())
      return Result<NonlinearStaticExecutionResult>::Failure(log_status);
  }

  std::vector<Diagnostic> diagnostics = domain.Warnings();
  SortDiagnostics(diagnostics);
  return Result<NonlinearStaticExecutionResult>::Success(
      {std::move(state.Value()), std::move(diagnostics)});
}

} // namespace fesa