#include "fesa/analysis/dynamic_analysis_state.h"

#include "results/mitc4_drilling_evidence.h"

#include <cmath>
#include <limits>
#include <stdexcept>
#include <string>
#include <utility>

#include "fesa/core/source_identity.h"
#include "fesa/elements/element_definition.h"
#include "fesa/elements/finite_rotation.h"
#include "math/finite_values.h"
#include "results/shell_result_recovery.h"
#include "analysis/shell_state_inventory.h"

namespace fesa {
namespace {

using analysis_internal::ShellElementOrder;
using math_internal::IsFinite;

constexpr std::size_t kDofsPerNode = 6U;
constexpr std::size_t kAngularDofsPerNode = 3U;
constexpr std::size_t kShellLocationsPerElement = 4U;
constexpr std::uint32_t kMaximumNewtonIterations = 100U;

/// @brief Reports a rejected state without changing committed history.
Status Failure(const FailureCategory category, const std::string& message) {
  return Status::Failure(category, {{Severity::kError,
                                     "dynamic-invalid-state",
                                     {},
                                     "DYNAMIC_ANALYSIS_STATE",
                                     {},
                                     message}});
}

/// @brief Requires complete ordered Gauss and section-position rows.
Status ValidateShellRows(const DynamicAcceptedFrameState& frame,
                         const std::vector<EntityIndex>& shell_order) {
  if (shell_order.size() > (std::numeric_limits<std::size_t>::max)() /
                               kShellLocationsPerElement ||
      frame.shell_results.size() !=
          shell_order.size() * kShellLocationsPerElement) {
    return Failure(FailureCategory::kSolver,
                   "Shell recovery inventory does not match the Domain.");
  }
  constexpr std::array<ShellMidsurfaceLocation, kShellLocationsPerElement>
      kLocations{ShellMidsurfaceLocation::kGp1, ShellMidsurfaceLocation::kGp2,
                 ShellMidsurfaceLocation::kGp3, ShellMidsurfaceLocation::kGp4};
  constexpr std::array<ShellSectionPosition, 3U> kPositions{
      ShellSectionPosition::kBottom, ShellSectionPosition::kMiddle,
      ShellSectionPosition::kTop};
  constexpr std::array<double, 3U> kZeta{-1.0, 0.0, 1.0};
  const double gauss = 1.0 / std::sqrt(3.0);
  const std::array<std::array<double, 2U>, 4U> natural{
      std::array<double, 2U>{-gauss, -gauss},
      {gauss, -gauss},
      {gauss, gauss},
      {-gauss, gauss}};
  for (std::size_t element = 0U; element < shell_order.size(); ++element) {
    for (std::size_t point = 0U; point < kShellLocationsPerElement; ++point) {
      const auto& row =
          frame.shell_results[element * kShellLocationsPerElement + point];
      if (row.element != shell_order[element] ||
          row.location != kLocations[point] ||
          row.natural_coordinates != natural[point] || !results_internal::IsFiniteShellRow(row)) {
        return Failure(FailureCategory::kSolver,
                       "Shell recovery rows have invalid identity or values.");
      }
      for (std::size_t position = 0U; position < kPositions.size();
           ++position) {
        if (row.stress[position].position != kPositions[position] ||
            row.stress[position].zeta != kZeta[position]) {
          return Failure(FailureCategory::kSolver,
                         "Shell stress rows have invalid section identity.");
        }
      }
    }
  }
  return Status::Ok();
}

}  // namespace

Result<DynamicAnalysisState> DynamicAnalysisState::Create(
    const Domain& domain, const DofManager& dofs) {
  const DynamicStepDefinition* const dynamic = domain.Dynamic();
  const std::size_t node_count = domain.Nodes().size();
  if (dynamic == nullptr || dynamic->increment_count == 0U ||
      !std::isfinite(dynamic->initial_increment) ||
      !std::isfinite(dynamic->time_period) ||
      dynamic->initial_increment <= 0.0 ||
      dynamic->increment_count > dynamic->maximum_increments ||
      !std::isfinite(dynamic->alpha) || dynamic->alpha < -1.0 / 3.0 ||
      dynamic->alpha > 0.0 ||
      dynamic->beta != 0.25 * (1.0 - dynamic->alpha) * (1.0 - dynamic->alpha) ||
      dynamic->gamma != 0.5 - dynamic->alpha || dynamic->time_period <= 0.0 ||
      static_cast<double>(dynamic->increment_count) *
              dynamic->initial_increment !=
          dynamic->time_period ||
      node_count == 0U || domain.ShellElements().Empty() ||
      domain.Elements().Size() != domain.ShellElements().Size() ||
      domain.ShellNodeInitialFrames().size() != node_count ||
      node_count > (std::numeric_limits<std::size_t>::max)() / kDofsPerNode ||
      dofs.FullDofCount() != node_count * kDofsPerNode ||
      !dofs.ValidateInvariants().IsOk()) {
    return Result<DynamicAnalysisState>::Failure(
        Failure(FailureCategory::kModel,
                "Dynamic schedule and full-DOF inventory must be valid."));
  }
  for (std::size_t index = 0U; index < dofs.PrescribedValues().Size();
       ++index) {
    if (dofs.PrescribedValues()[index] != 0.0) {
      return Result<DynamicAnalysisState>::Failure(
          Failure(FailureCategory::kModel,
                  "Dynamic prescribed components must be exactly zero."));
    }
  }
  std::vector<SourceEntityId> source_ids;
  source_ids.reserve(node_count);
  for (const Node& node : domain.Nodes()) {
    source_ids.push_back(node.source_id);
  }
  return Result<DynamicAnalysisState>::Success(DynamicAnalysisState{
      dofs.FullDofCount(), dynamic->increment_count, dynamic->initial_increment,
      ShellElementOrder(domain), dofs.ConstrainedDofs(), std::move(source_ids),
      domain.ShellNodeInitialFrames()});
}

Result<DynamicAcceptedFrameState>
DynamicAnalysisState::BeginTrialFromCommitted() const {
  if (accepted_frames_.empty()) {
    return Result<DynamicAcceptedFrameState>::Failure(
        Failure(FailureCategory::kSolver,
                "A dynamic trial requires an accepted frame zero."));
  }
  return Result<DynamicAcceptedFrameState>::Success(accepted_frames_.back());
}

Status DynamicAnalysisState::CommitAcceptedFrame(
    const DynamicAcceptedFrameState& frame) {
  if (!initialized_) {
    return Failure(FailureCategory::kSolver,
                   "A default dynamic state cannot accept frames.");
  }
  if ((frame.time_integration_formulation !=
           DynamicTimeIntegrationFormulation::kFullCoordinateHht &&
       frame.time_integration_formulation !=
           DynamicTimeIntegrationFormulation::kReducedCoordinateHht) ||
      (!accepted_frames_.empty() &&
       frame.time_integration_formulation !=
           accepted_frames_.front().time_integration_formulation)) {
    return Failure(FailureCategory::kSolver,
                   "Dynamic integration formulation is invalid or changed.");
  }
  const std::size_t expected_frame = accepted_frames_.size();
  if (expected_frame > (std::numeric_limits<std::uint32_t>::max)() ||
      frame.frame_index != expected_frame ||
      frame.increment_index != expected_frame ||
      frame.increment_index > increment_count_ || !std::isfinite(frame.time) ||
      frame.time !=
          static_cast<double>(frame.increment_index) * initial_increment_ ||
      frame.is_initial != (frame.frame_index == 0U) ||
      (frame.is_initial && frame.iteration_count != 0U) ||
      frame.iteration_count > kMaximumNewtonIterations ||
      (frame.is_initial && std::signbit(frame.time))) {
    return Failure(FailureCategory::kSolver,
                   "Frame schedule or iteration identity is invalid.");
  }
  if (!accepted_frames_.empty() && frame.time <= accepted_frames_.back().time) {
    return Failure(FailureCategory::kSolver,
                   "Accepted dynamic frame times must strictly increase.");
  }
  const std::size_t node_count = nodal_source_ids_.size();
  if (frame.nodal_internal_order.size() != node_count ||
      frame.nodal_source_ids.size() != node_count ||
      frame.nodal_triad.size() != node_count ||
      frame.nodal_director.size() != node_count) {
    return Failure(FailureCategory::kSolver,
                   "Frame node inventory does not match the Domain.");
  }
  for (std::size_t node = 0U; node < node_count; ++node) {
    if (frame.nodal_internal_order[node] != node ||
        !SourceIdentityEquals(frame.nodal_source_ids[node], nodal_source_ids_[node]) ||
        !IsFinite(frame.nodal_triad[node][0U]) ||
        !IsFinite(frame.nodal_triad[node][1U]) ||
        !IsFinite(frame.nodal_triad[node][2U]) ||
        !IsFinite(frame.nodal_director[node])) {
      return Failure(FailureCategory::kSolver,
                     "Frame nodal identity or orientation is invalid.");
    }
  }
  const std::array<const Vector*, 10U> full_vectors{
      &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};
  for (const Vector* values : full_vectors) {
    if (values->Size() != full_dof_count_ || !IsFinite(*values)) {
      return Failure(FailureCategory::kSolver,
                     "Frame generalized vector inventory is invalid.");
    }
  }
  // N5 preserves the existing chart and 1e-12 orientation consistency checks.
  for (std::size_t node = 0U; node < node_count; ++node) {
    Matrix initial{3U, 3U};
    for (std::size_t component = 0U; component < 3U; ++component) {
      initial(component, 0U) = initial_frames_[node].tangent_a[component];
      initial(component, 1U) = initial_frames_[node].tangent_b[component];
      initial(component, 2U) = initial_frames_[node].director[component];
    }
    const std::size_t offset = node * kDofsPerNode;
    const Vector3 phi{frame.generalized_coordinates[offset + 3U],
                      frame.generalized_coordinates[offset + 4U],
                      frame.generalized_coordinates[offset + 5U]};
    const auto rotation = MakeFiniteRotationState(phi, initial);
    if (!rotation.HasValue()) {
      return Failure(FailureCategory::kSolver,
                     "Frame rotation is outside the valid initial chart.");
    }
    for (std::size_t axis = 0U; axis < 3U; ++axis) {
      for (std::size_t component = 0U; component < 3U; ++component) {
        if (std::abs(frame.nodal_triad[node][axis][component] -
                     rotation.Value().q(component, axis)) > 1.0e-12) {
          return Failure(FailureCategory::kSolver,
                         "Frame triad does not match its rotation coordinate.");
        }
      }
      if (std::abs(frame.nodal_director[node][axis] -
                   rotation.Value().q(axis, 2U)) > 1.0e-12) {
        return Failure(
            FailureCategory::kSolver,
            "Frame director does not match its rotation coordinate.");
      }
    }
  }
  const std::array<const Vector*, 6U> hidden_vectors{
      &frame.hidden_constraint,
      &frame.hidden_velocity_defect,
      &frame.hidden_acceleration_defect,
      &frame.hidden_constraint_scale,
      &frame.hidden_velocity_defect_scale,
      &frame.hidden_acceleration_defect_scale};
  for (const Vector* values : hidden_vectors) {
    if (values->Size() != frame.hidden_constraint.Size() ||
        !IsFinite(*values)) {
      return Failure(FailureCategory::kSolver,
                     "Frame hidden-constraint evidence has invalid size.");
    }
  }
  if ((!accepted_frames_.empty() &&
       frame.hidden_constraint.Size() != hidden_constraint_count_) ||
      frame.angular_velocity.Size() != node_count * kAngularDofsPerNode ||
      frame.angular_acceleration.Size() != node_count * kAngularDofsPerNode ||
      !IsFinite(frame.angular_velocity) ||
      !IsFinite(frame.angular_acceleration) ||
      !std::isfinite(frame.kinetic_energy) ||
      !std::isfinite(frame.physical_strain_energy) ||
      !std::isfinite(frame.drilling_energy) ||
      !std::isfinite(frame.external_work) ||
      !std::isfinite(frame.energy_defect) ||
      !IsFinite(frame.global_equilibrium) || !IsFinite(frame.linear_momentum) ||
      !IsFinite(frame.angular_momentum) ||
      !IsFinite(frame.linear_momentum_rate) ||
      !IsFinite(frame.angular_momentum_rate) ||
      !IsFinite(frame.free_residual_metrics)) {
    return Failure(FailureCategory::kSolver,
                   "Frame physical evidence is nonfinite or has invalid size.");
  }
  if (!results_internal::ValidDrillingEvidence(frame.mitc4_drilling,
                                               frame.global_equilibrium) ||
      (!accepted_frames_.empty() && !results_internal::SameDrillingInventory(
          frame.mitc4_drilling, accepted_frames_.back().mitc4_drilling)))
    return Failure(FailureCategory::kSolver,
                   "Accepted drilling evidence must be finite, consistent, and retain its mode.");
  for (const std::size_t dof : constrained_dofs_) {
    if (frame.generalized_coordinates[dof] != 0.0 ||
        frame.velocity[dof] != 0.0 || frame.acceleration[dof] != 0.0 ||
        frame.algorithmic_acceleration[dof] != 0.0) {
      return Failure(FailureCategory::kSolver,
                     "Fixed boundary components must remain exactly zero.");
    }
  }
  for (const Vector* scales :
       {&frame.hidden_constraint_scale, &frame.hidden_velocity_defect_scale,
        &frame.hidden_acceleration_defect_scale}) {
    for (std::size_t index = 0U; index < scales->Size(); ++index) {
      if ((*scales)[index] < 0.0) {
        return Failure(FailureCategory::kSolver,
                       "Hidden evidence scales must be nonnegative.");
      }
    }
  }
  if (frame.free_residual_metrics[0U] < 0.0 ||
      frame.free_residual_metrics[1U] < 0.0 ||
      (frame.is_initial &&
       (frame.external_work != 0.0 || frame.energy_defect != 0.0))) {
    return Failure(FailureCategory::kSolver,
                   "Frame metric or initial energy evidence is invalid.");
  }
  std::size_t constraint = 0U;
  for (std::size_t dof = 0U; dof < full_dof_count_; ++dof) {
    const bool constrained = constraint < constrained_dofs_.size() &&
                             constrained_dofs_[constraint] == dof;
    if (constrained) ++constraint;
    if (frame.residual[dof] != frame.inertia_force[dof] +
                                   frame.internal_force[dof] -
                                   frame.external_force[dof] ||
        frame.reaction[dof] != (constrained ? frame.residual[dof] : 0.0) ||
        frame.residual_scale[dof] < 0.0 ||
        (!accepted_frames_.empty() &&
         frame.external_force[dof] !=
             accepted_frames_.front().external_force[dof]) ||
        (frame.is_initial &&
         (frame.generalized_coordinates[dof] != 0.0 ||
          frame.velocity[dof] != 0.0 ||
          frame.acceleration[dof] != frame.algorithmic_acceleration[dof]))) {
      return Failure(
          FailureCategory::kSolver,
          "Frame force, reaction, or initial-state identity is invalid.");
    }
  }
  const Status shell_status = ValidateShellRows(frame, shell_element_order_);
  if (!shell_status.IsOk()) {
    return shell_status;
  }
  const std::size_t hidden_constraint_count = frame.hidden_constraint.Size();
  // Copy insertion preserves the previous vector when allocation/copy fails.
  accepted_frames_.push_back(frame);
  hidden_constraint_count_ = hidden_constraint_count;
  return Status::Ok();
}

void DynamicAnalysisState::DiscardTrial() noexcept {}

const std::vector<DynamicAcceptedFrameState>&
DynamicAnalysisState::AcceptedFrames() const noexcept {
  return accepted_frames_;
}

const DynamicAcceptedFrameState& DynamicAnalysisState::CommittedState() const {
  if (accepted_frames_.empty()) {
    throw std::logic_error{
        "Committed dynamic state is unavailable before frame zero."};
  }
  return accepted_frames_.back();
}

bool DynamicAnalysisState::IsInitialized() const noexcept {
  return initialized_;
}

DynamicAnalysisState::DynamicAnalysisState(
    const std::size_t full_dof_count, const std::uint32_t increment_count,
    const double initial_increment,
    std::vector<EntityIndex> shell_element_order,
    std::vector<std::size_t> constrained_dofs,
    std::vector<SourceEntityId> nodal_source_ids,
    std::vector<ShellNodeInitialFrame> initial_frames)
    : full_dof_count_{full_dof_count},
      increment_count_{increment_count},
      initial_increment_{initial_increment},
      shell_element_order_{std::move(shell_element_order)},
      constrained_dofs_{std::move(constrained_dofs)},
      nodal_source_ids_{std::move(nodal_source_ids)},
      initial_frames_{std::move(initial_frames)},
      initialized_{true} {}

}  // namespace fesa