#include "results/global_equilibrium_recovery.h"

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

#include "fesa/math/vector3.h"
#include "results/mitc4_drilling_evidence.h"

namespace fesa::results_internal {
namespace {

constexpr std::size_t kDofsPerNode = 6U;
constexpr double kFreeResidualTolerance = 1.0e-10;
constexpr double kGlobalEquilibriumTolerance = 1.0e-10;

double IndexedNorm(const Vector& values,
                   const std::vector<std::size_t>& indices) {
  double result = 0.0;
  for (const std::size_t index : indices) {
    result = std::hypot(result, values[index]);
  }
  return result;
}

/// @brief Computes separate Kff*df and Kfc*dc norms through sparse MKL.
std::array<double, 2> FreeEquationInternalTermNorms(
    const SparseMatrix& stiffness, const Vector& displacement,
    const DofManager& dofs) {
  Vector free_displacement{displacement.Size()};
  Vector constrained_displacement = displacement;
  for (const std::size_t full_dof : dofs.FreeDofs()) {
    free_displacement[full_dof] = displacement[full_dof];
    constrained_displacement[full_dof] = 0.0;
  }
  const auto free_term = stiffness.Multiply(free_displacement);
  const auto constrained_term = stiffness.Multiply(constrained_displacement);
  double free_term_norm = 0.0;
  double constrained_term_norm = 0.0;
  for (const std::size_t row : dofs.FreeDofs()) {
    // Only the original free-row terms enter these acceptance scales.
    free_term_norm = std::hypot(free_term_norm, free_term[row]);
    constrained_term_norm =
        std::hypot(constrained_term_norm, constrained_term[row]);
  }
  return {free_term_norm, constrained_term_norm};
}

bool AccumulateVectorAndScale(std::array<double, 3>& total, double& scale,
                              const std::array<double, 3>& contribution) {
  const Vector3 contribution_vector{contribution};
  const double magnitude = contribution_vector.Norm();
  const double accumulated_scale = scale + magnitude;
  if (!contribution_vector.IsFinite() || !std::isfinite(magnitude) ||
      !std::isfinite(accumulated_scale)) {
    return false;
  }
  for (std::size_t component = 0U; component < contribution.size();
       ++component) {
    const double accumulated = total[component] + contribution[component];
    if (!std::isfinite(accumulated)) {
      return false;
    }
    total[component] = accumulated;
  }
  scale = accumulated_scale;
  return true;
}

double NormalizedBalance(const std::array<double, 3>& balance,
                         const double scale) {
  const double balance_norm = Vector3{balance}.Norm();
  if (!std::isfinite(balance_norm) || !std::isfinite(scale)) {
    return (std::numeric_limits<double>::infinity)();
  }
  if (scale == 0.0) {
    return balance_norm == 0.0 ? 0.0
                               : (std::numeric_limits<double>::infinity)();
  }
  return balance_norm / scale;
}

}  // namespace

Result<GlobalEquilibriumResult> GlobalEquilibriumRecovery(
    const Domain& domain, const DofManager& dofs,
    const SparseMatrix& full_stiffness, const Vector& full_displacement,
    const Vector& full_external_force) try {
  const std::size_t full_count = dofs.FullDofCount();
  if (full_stiffness.Rows() != full_count ||
      full_stiffness.Columns() != full_count ||
      full_displacement.Size() != full_count ||
      full_external_force.Size() != full_count ||
      domain.Nodes().size() * kDofsPerNode != full_count) {
    return RecoveryResultFailure<GlobalEquilibriumResult>(
        "invalid-recovery-dimensions", {domain.SourcePath(), 0U},
        domain.SourceContentIdentity(),
        "Model, DOF, stiffness, and recovery vector dimensions must agree.");
  }
  const Status matrix_status = full_stiffness.Validate();
  if (!matrix_status.IsOk()) {
    return Result<GlobalEquilibriumResult>::Failure(matrix_status);
  }
  if (!IsFiniteVector(full_displacement) ||
      !IsFiniteVector(full_external_force)) {
    return RecoveryResultFailure<GlobalEquilibriumResult>(
        "nonfinite-recovery-value", {domain.SourcePath(), 0U},
        domain.SourceContentIdentity(),
        "Displacement and external-force inputs must be finite.");
  }

  GlobalEquilibriumResult result{full_stiffness.Multiply(full_displacement),
                                 Vector{full_count}, Vector{full_count}, 0.0};
  if (!IsFiniteVector(result.internal_force)) {
    return RecoveryResultFailure<GlobalEquilibriumResult>(
        "nonfinite-recovery-value", {domain.SourcePath(), 0U},
        domain.SourceContentIdentity(),
        "Full stiffness multiplication must produce finite internal force.");
  }
  for (std::size_t full_dof = 0U; full_dof < result.residual.Size();
       ++full_dof) {
    result.residual[full_dof] =
        result.internal_force[full_dof] - full_external_force[full_dof];
    if (!std::isfinite(result.residual[full_dof])) {
      return RecoveryResultFailure<GlobalEquilibriumResult>(
          "nonfinite-recovery-value", {domain.SourcePath(), 0U},
          std::to_string(full_dof),
          "Internal-minus-external residual must remain finite.");
    }
  }

  const double residual_norm = IndexedNorm(result.residual, dofs.FreeDofs());
  const auto internal_term_norms =
      FreeEquationInternalTermNorms(full_stiffness, full_displacement, dofs);
  const double external_norm =
      IndexedNorm(full_external_force, dofs.FreeDofs());
  // Normalize against Kff*df, Kfc*dc, and Ff separately; the already-cancelled
  // K*d term would hide prescribed-displacement scaling.
  const double denominator =
      (std::max)({internal_term_norms[0U], internal_term_norms[1U],
                  external_norm});
  if (!std::isfinite(residual_norm) || !std::isfinite(denominator)) {
    return RecoveryResultFailure<GlobalEquilibriumResult>(
        "nonfinite-recovery-value", {domain.SourcePath(), 0U}, "free-residual",
        "Free residual and its physical normalization scale must be finite.");
  }
  result.normalized_free_residual =
      denominator == 0.0
          ? (residual_norm == 0.0 ? 0.0
                                  : (std::numeric_limits<double>::infinity)())
          : residual_norm / denominator;
  if (!std::isfinite(result.normalized_free_residual) ||
      result.normalized_free_residual > kFreeResidualTolerance) {
    return RecoveryResultFailure<GlobalEquilibriumResult>(
        "free-residual-tolerance-failure", {domain.SourcePath(), 0U},
        "free-residual", "The normalized free residual exceeds 1e-10.");
  }

  // The state reaction vector intentionally preserves free residual evidence;
  // shell global equilibrium below separately filters constrained reactions.
  result.reaction = result.residual;
  return Result<GlobalEquilibriumResult>::Success(std::move(result));
} catch (const std::length_error&) {
  return RecoveryResultFailure<GlobalEquilibriumResult>(
      "recovery-index-overflow", {domain.SourcePath(), 0U},
      domain.SourceContentIdentity(),
      "Recovery multiplication exceeds the backend integer range.");
} catch (const std::runtime_error&) {
  return RecoveryResultFailure<GlobalEquilibriumResult>(
      "recovery-multiplication-failed", {domain.SourcePath(), 0U},
      domain.SourceContentIdentity(), "Recovery sparse multiplication failed.");
}

Status RecoverShellGlobalEvidence(const Domain& domain, const DofManager& dofs,
                                  const Vector& external_force,
                                  const GlobalEquilibriumResult& global,
                                  RecoveryCandidate& candidate) {
  std::vector<unsigned char> constrained(dofs.FullDofCount(), 0U);
  for (const std::size_t full_dof : dofs.ConstrainedDofs()) {
    constrained[full_dof] = 1U;
  }

  std::array<double, 3> applied_force{};
  std::array<double, 3> reaction_force{};
  std::array<double, 3> applied_moment{};
  std::array<double, 3> reaction_moment{};
  double applied_force_scale = 0.0;
  double reaction_force_scale = 0.0;
  double applied_moment_scale = 0.0;
  double reaction_moment_scale = 0.0;
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const std::size_t offset = node * kDofsPerNode;
    std::array<double, 3> nodal_applied_force{};
    std::array<double, 3> nodal_reaction_force{};
    std::array<double, 3> nodal_applied_moment{};
    std::array<double, 3> nodal_reaction_moment{};
    for (std::size_t component = 0U; component < 3U; ++component) {
      nodal_applied_force[component] = external_force[offset + component];
      nodal_applied_moment[component] = external_force[offset + 3U + component];
      if (constrained[offset + component] != 0U) {
        nodal_reaction_force[component] = global.residual[offset + component];
      }
      if (constrained[offset + 3U + component] != 0U) {
        nodal_reaction_moment[component] =
            global.residual[offset + 3U + component];
      }
    }

    const Vector3 position{domain.Nodes()[node].coordinates};
    const Vector3 applied_force_moment =
        position.Cross(Vector3{nodal_applied_force});
    const Vector3 reaction_force_moment =
        position.Cross(Vector3{nodal_reaction_force});
    for (std::size_t component = 0U; component < 3U; ++component) {
      nodal_applied_moment[component] += applied_force_moment[component];
      nodal_reaction_moment[component] += reaction_force_moment[component];
    }
    if (!AccumulateVectorAndScale(applied_force, applied_force_scale,
                                  nodal_applied_force) ||
        !AccumulateVectorAndScale(reaction_force, reaction_force_scale,
                                  nodal_reaction_force) ||
        !AccumulateVectorAndScale(applied_moment, applied_moment_scale,
                                  nodal_applied_moment) ||
        !AccumulateVectorAndScale(reaction_moment, reaction_moment_scale,
                                  nodal_reaction_moment)) {
      return RecoveryFailure("nonfinite-recovery-value",
                             domain.Nodes()[node].location,
                             domain.Nodes()[node].source_id.source_label_text,
                             "Global force and moment evidence must remain "
                             "finite in source-node order.");
    }
  }

  const Vector3 force_balance =
      Vector3{applied_force} + Vector3{reaction_force};
  const Vector3 moment_balance =
      Vector3{applied_moment} + Vector3{reaction_moment};
  for (std::size_t component = 0U; component < 3U; ++component) {
    candidate.shell.equilibrium[component] = force_balance[component];
    candidate.shell.equilibrium[3U + component] = moment_balance[component];
  }
  Vector3 acceptance_moment = moment_balance;
  if (candidate.shell.mitc4_drilling) {
    CompleteDrillingEvidence(candidate.shell.equilibrium,
                             &*candidate.shell.mitc4_drilling);
    if (!ValidDrillingEvidence(candidate.shell.mitc4_drilling,
                               candidate.shell.equilibrium))
      return RecoveryFailure("invalid-shell-drilling-evidence",
                             {domain.SourcePath(), 0U}, "global-equilibrium",
                             "Shell drilling diagnostics must be finite and consistent.");
    if (candidate.shell.mitc4_drilling->mode == Mitc4DrillingMode::kDirectSpring)
      acceptance_moment = moment_balance -
          Vector3(candidate.shell.mitc4_drilling->drilling_moment);
  }
  const double force_metric =
      NormalizedBalance(force_balance.Components(),
                        (std::max)(applied_force_scale, reaction_force_scale));
  const double moment_metric = NormalizedBalance(
      acceptance_moment.Components(),
      (std::max)(applied_moment_scale, reaction_moment_scale));
  candidate.shell.verification_metrics = {global.normalized_free_residual,
                                          force_metric, moment_metric};
  if (!IsFiniteArray(candidate.shell.equilibrium) ||
      !Vector3{candidate.shell.verification_metrics}.IsFinite()) {
    return RecoveryFailure("nonfinite-recovery-value",
                           {domain.SourcePath(), 0U}, "global-equilibrium",
                           "Global equilibrium values and their physical "
                           "normalization scales must be finite.");
  }
  if (force_metric > kGlobalEquilibriumTolerance ||
      moment_metric > kGlobalEquilibriumTolerance) {
    return RecoveryFailure(
        "global-equilibrium-tolerance-failure", {domain.SourcePath(), 0U},
        "global-equilibrium",
        "Normalized global force or moment balance exceeds 1e-10.");
  }
  return Status::Ok();
}

}  // namespace fesa::results_internal