#include "fesa/analysis/nonlinear_analysis_state.h"

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

#include "fesa/elements/element_definition.h"

namespace fesa {
namespace {

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

/// @brief Creates one failure without exposing a partially built state.
Status StateFailure(const FailureCategory category, const std::string& code,
                    const std::string& identity, const std::string& message) {
  return Status::Failure(
      category,
      {{Severity::kError, code, {}, "NONLINEAR_ANALYSIS_STATE", identity,
        message}});
}

/// @brief Compares every preserved source-identity component.
bool SourceIdsMatch(const SourceEntityId& left, const SourceEntityId& right) {
  return left.instance_name == right.instance_name &&
         left.source_label == right.source_label &&
         left.source_label_text == right.source_label_text;
}

template <std::size_t Size>
/// @brief Tests one fixed-size numerical array for finite values.
bool IsFinite(const std::array<double, Size>& values) {
  return std::all_of(values.begin(), values.end(),
                     [](const double value) { return std::isfinite(value); });
}

/// @brief Tests each component of a stored trial vector for finiteness.
bool IsFinite(const Vector& values) {
  for (std::size_t index = 0U; index < values.Size(); ++index) {
    if (!std::isfinite(values[index])) {
      return false;
    }
  }
  return true;
}

/// @brief Tests all physical quantities attached to one shell recovery row.
bool IsFinite(const ShellResultRow& row) {
  if (!IsFinite(row.natural_coordinates) || !IsFinite(row.generalized_strain) ||
      !IsFinite(row.section_resultant)) {
    return false;
  }
  for (const auto& axis : row.local_frame) {
    if (!IsFinite(axis)) {
      return false;
    }
  }
  return std::all_of(row.stress.begin(), row.stress.end(),
                     [](const ShellSectionStressRow& stress) {
                       return std::isfinite(stress.zeta) &&
                              IsFinite(stress.components);
                     });
}

/// @brief Returns shell Domain indices in deterministic unified element order.
std::vector<EntityIndex> ShellElementOrder(const Domain& domain) {
  std::vector<EntityIndex> order;
  for (std::size_t index = 0U; index < domain.Elements().Size(); ++index) {
    if (domain.Elements()[index].Kind() == ElementDefinitionKind::kMitc4Shell) {
      order.push_back(static_cast<EntityIndex>(index));
    }
  }
  return order;
}

/// @brief Validates one accepted frame before any committed-state mutation.
Status ValidateAcceptedFrame(const AcceptedFrameState& frame,
                             const NonlinearTrialState& committed,
                             const std::size_t full_dof_count,
                             const std::vector<EntityIndex>& shell_order,
                             const std::size_t accepted_frame_count) {
  if (accepted_frame_count >
      static_cast<std::size_t>((std::numeric_limits<std::uint32_t>::max)())) {
    return StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                        "Accepted-frame inventory exceeds uint32 identity.");
  }
  if (frame.frame_index !=
      static_cast<std::uint32_t>(accepted_frame_count)) {
    return StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                        "Frame index must equal accepted-frame order.");
  }
  if (frame.increment_index <= committed.increment_index ||
      !std::isfinite(frame.load_factor) || frame.load_factor <= 0.0 ||
      frame.load_factor > 1.0 || frame.load_factor <= committed.load_factor ||
      frame.iteration_count == 0U ||
      frame.iteration_count > kMaximumNewtonIterations ||
      frame.is_factor_one != (frame.load_factor == 1.0)) {
    return StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                        "Accepted frame has an invalid increment identity.");
  }

  const std::size_t node_count = committed.nodal_internal_order.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 StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                        "Accepted frame node inventory does not match baseline.");
  }
  for (std::size_t node = 0U; node < node_count; ++node) {
    if (frame.nodal_internal_order[node] !=
            committed.nodal_internal_order[node] ||
        !SourceIdsMatch(frame.nodal_source_ids[node],
                         committed.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 StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                          "Accepted frame has invalid nodal identity or state.");
    }
  }

  const std::array<const Vector*, 7> vectors{
      &frame.generalized_coordinates, &frame.displacement,
      &frame.rotation_vector,         &frame.internal_force,
      &frame.external_force,          &frame.residual,
      &frame.reaction};
  for (const Vector* vector : vectors) {
    if (vector->Size() != full_dof_count || !IsFinite(*vector)) {
      return StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                          "Accepted frame vector inventory is invalid.");
    }
  }
  for (std::size_t dof = 0U; dof < full_dof_count; ++dof) {
    const double residual = frame.internal_force[dof] - frame.external_force[dof];
    if (frame.residual[dof] != residual || frame.reaction[dof] != residual) {
      return StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                          "Residual and reaction must retain full equilibrium.");
    }
  }

  if (shell_order.size() >
      (std::numeric_limits<std::size_t>::max)() / kShellLocationsPerElement ||
      frame.shell_results.size() !=
          shell_order.size() * kShellLocationsPerElement ||
      !std::isfinite(frame.physical_strain_energy) ||
      !IsFinite(frame.global_equilibrium)) {
    return StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                        "Accepted frame physical result inventory is invalid.");
  }

  const double gauss = 1.0 / std::sqrt(3.0);
  const std::array<ShellMidsurfaceLocation, kShellLocationsPerElement>
      expected_locations{ShellMidsurfaceLocation::kGp1,
                         ShellMidsurfaceLocation::kGp2,
                         ShellMidsurfaceLocation::kGp3,
                         ShellMidsurfaceLocation::kGp4};
  const std::array<std::array<double, 2>, kShellLocationsPerElement>
      expected_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}};
  const std::array<ShellSectionPosition, 3> expected_positions{
      ShellSectionPosition::kBottom, ShellSectionPosition::kMiddle,
      ShellSectionPosition::kTop};
  constexpr std::array<double, 3> expected_zeta{-1.0, 0.0, 1.0};
  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 != expected_locations[point] ||
          row.natural_coordinates != expected_coordinates[point] ||
          !IsFinite(row)) {
        return StateFailure(FailureCategory::kSolver, "invalid-accepted-state", {},
                            "Shell result rows do not preserve stable identity.");
      }
      for (std::size_t position = 0U; position < expected_positions.size();
           ++position) {
        if (row.stress[position].position != expected_positions[position] ||
            row.stress[position].zeta != expected_zeta[position]) {
          return StateFailure(
              FailureCategory::kSolver, "invalid-accepted-state", {},
              "Shell stress rows do not preserve stable section identity.");
        }
      }
    }
  }
  return Status::Ok();
}

}  // namespace

Result<NonlinearAnalysisState> NonlinearAnalysisState::Create(
    const Domain& domain, const DofManager& dofs,
    const NonlinearStateBaselineInput& baseline) {
  const std::size_t node_count = domain.Nodes().size();
  if (node_count > (std::numeric_limits<std::size_t>::max)() / kDofsPerNode ||
      dofs.FullDofCount() != node_count * kDofsPerNode ||
      !dofs.ValidateInvariants().IsOk()) {
    return Result<NonlinearAnalysisState>::Failure(StateFailure(
        FailureCategory::kModel, "invalid-nonlinear-state", {},
        "Domain node and full-DOF inventories must agree."));
  }
  if (baseline.node_baselines.size() != node_count) {
    return Result<NonlinearAnalysisState>::Failure(StateFailure(
        FailureCategory::kModel, "invalid-nonlinear-state", {},
        "Nonlinear baseline requires exactly one entry per Domain node."));
  }

  NonlinearTrialState committed{
      0U,
      0.0,
      true,
      {},
      {},
      Vector{dofs.FullDofCount()},
      Vector{dofs.FullDofCount()},
      Vector{dofs.FullDofCount()},
      {},
      {},
      Vector{dofs.FullDofCount()},
      Vector{dofs.FullDofCount()},
      Vector{dofs.FullDofCount()},
      Vector{dofs.FullDofCount()},
      {},
      0.0,
      {}};
  committed.nodal_internal_order.reserve(node_count);
  committed.nodal_source_ids.reserve(node_count);
  committed.nodal_triad.reserve(node_count);
  committed.nodal_director.reserve(node_count);
  for (std::size_t node = 0U; node < node_count; ++node) {
    const auto& entry = baseline.node_baselines[node];
    if (entry.node_index != static_cast<EntityIndex>(node) ||
        !SourceIdsMatch(entry.source_id, domain.Nodes()[node].source_id) ||
        !IsFinite(entry.triad[0U]) || !IsFinite(entry.triad[1U]) ||
        !IsFinite(entry.triad[2U]) || !IsFinite(entry.director)) {
      return Result<NonlinearAnalysisState>::Failure(StateFailure(
          FailureCategory::kModel, "invalid-nonlinear-state",
          std::to_string(node),
          "Baseline identity and finite node frame must match Domain order."));
    }
    committed.nodal_internal_order.push_back(entry.node_index);
    committed.nodal_source_ids.push_back(entry.source_id);
    committed.nodal_triad.push_back(entry.triad);
    committed.nodal_director.push_back(entry.director);
  }

  return Result<NonlinearAnalysisState>::Success(NonlinearAnalysisState{
      std::move(committed), dofs.FullDofCount(), ShellElementOrder(domain)});
}

Result<NonlinearTrialState> NonlinearAnalysisState::BeginTrialFromCommitted()
    const {
  return Result<NonlinearTrialState>::Success(committed_state_);
}

Status NonlinearAnalysisState::CommitAcceptedFrame(
    const AcceptedFrameState& frame) {
  const Status validation =
      ValidateAcceptedFrame(frame, committed_state_, full_dof_count_,
                            shell_element_order_, accepted_frames_.size());
  if (!validation.IsOk()) {
    return validation;
  }

  NonlinearTrialState candidate{
      frame.increment_index,
      frame.load_factor,
      false,
      frame.nodal_internal_order,
      frame.nodal_source_ids,
      frame.generalized_coordinates,
      frame.displacement,
      frame.rotation_vector,
      frame.nodal_triad,
      frame.nodal_director,
      frame.internal_force,
      frame.external_force,
      frame.residual,
      frame.reaction,
      frame.shell_results,
      frame.physical_strain_energy,
      frame.global_equilibrium};
  std::vector<AcceptedFrameState> accepted_candidate = accepted_frames_;
  accepted_candidate.push_back(frame);

  committed_state_ = std::move(candidate);
  accepted_frames_ = std::move(accepted_candidate);
  return Status::Ok();
}

void NonlinearAnalysisState::DiscardTrial() noexcept {}

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

const NonlinearTrialState& NonlinearAnalysisState::CommittedState() const noexcept {
  return committed_state_;
}

NonlinearAnalysisState::NonlinearAnalysisState(
    NonlinearTrialState committed_state, const std::size_t full_dof_count,
    std::vector<EntityIndex> shell_element_order)
    : committed_state_{std::move(committed_state)},
      full_dof_count_{full_dof_count},
      shell_element_order_{std::move(shell_element_order)} {}

}  // namespace fesa