#include "fesa/analysis/solid_static_procedure.h"

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

#include "fesa/analysis/analysis_model.h"
#include "fesa/assembly/load_assembler.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/assembly/sparse_assembler.h"
#include "fesa/constraints/essential_constraint_policy.h"
#include "fesa/elements/element_factory.h"
#include "fesa/elements/hex8_element.h"
#include "fesa/math/sparse_matrix.h"
#include "fesa/solvers/linear/linear_solver.h"
#include "fesa/solvers/linear/tangent_solver.h"

namespace fesa {
namespace {
constexpr std::uint32_t kMaximumIterations = 100U;
constexpr double kRoundoffBudget =
    512.0 * std::numeric_limits<double>::epsilon();

struct RuntimeSolid {
  const Hex8Element* element;
  std::vector<std::size_t> scatter;
};
struct SolidAssembly {
  Vector force;
  SparseMatrix tangent;
  std::vector<SolidIntegrationPointRow> points;
  double energy;
};

/// @brief Attaches immutable source-step context to a procedure failure.
Status Failure(const Domain& domain, const char* code, const char* message,
               FailureCategory category = FailureCategory::kSolver) {
  return Status::Failure(category, {{Severity::kError,
                                     code,
                                     {domain.SourcePath(), 0U},
                                     "STATIC",
                                     domain.SourceContentIdentity(),
                                     message}});
}
/// @brief Preserves kernel source/entity evidence for failed trial recovery.
Status TrialFailure(const Status& status) {
  return Status::Failure(FailureCategory::kSolver, status.Diagnostics());
}
/// @brief Returns component products with uncancelled arithmetic scales.
std::array<double, 3> Cross(const std::array<double, 3>& a,
                            const std::array<double, 3>& b,
                            std::array<double, 3>* scale = nullptr) {
  std::array<double, 3> result{};
  for (std::size_t i = 0U; i < 3U; ++i) {
    const auto j = (i + 1U) % 3U;
    const auto k = (i + 2U) % 3U;
    result[i] = a[j] * b[k] - a[k] * b[j];
    if (scale != nullptr)
      (*scale)[i] += std::abs(a[j] * b[k]) + std::abs(a[k] * b[j]);
  }
  return result;
}
/// @brief Tests a zero identity at its uncancelled same-unit operation scale.
bool RoundoffZero(double value, double scale) {
  return std::isfinite(value) && std::isfinite(scale) &&
         (scale == 0.0 ? value == 0.0
                       : std::abs(value) <= kRoundoffBudget * scale);
}
/// @brief Derives source DIRECT factors or the existing ordinary STATIC path.
Result<std::vector<LoadFactorIncrement>> Factors(const Domain& domain) {
  const auto& step = domain.Steps()[0];
  if (!step.IsDirectStatic()) return BuildFixedLoadFactorSequence(step);
  const auto count = step.IncrementCount();
  if (count == 0U || count > step.MaximumIncrements() ||
      !std::isfinite(step.InitialIncrement()) ||
      !(step.InitialIncrement() > 0.0) || !std::isfinite(step.TimePeriod()) ||
      !(step.TimePeriod() > 0.0) ||
      step.InitialIncrement() * static_cast<double>(count) != step.TimePeriod())
    return Result<std::vector<LoadFactorIncrement>>::Failure(Failure(
        domain, "invalid-solid-static-schedule",
        "DIRECT static requires a finite positive integer fixed schedule.",
        FailureCategory::kInput));
  std::vector<LoadFactorIncrement> factors;
  factors.reserve(count);
  double previous = 0.0;
  for (std::uint32_t n = 1U; n <= count; ++n) {
    const double factor = n == count
                              ? 1.0
                              : static_cast<double>(n) *
                                    step.InitialIncrement() / step.TimePeriod();
    if (!std::isfinite(factor) || !(factor > previous) || factor > 1.0)
      return Result<std::vector<LoadFactorIncrement>>::Failure(
          Failure(domain, "invalid-solid-static-schedule",
                  "DIRECT static factors must strictly increase to one.",
                  FailureCategory::kInput));
    factors.push_back({n, factor, n == count});
    previous = factor;
    if (n == count) break;  // Avoid unsigned wrap at the source limit.
  }
  return Result<std::vector<LoadFactorIncrement>>::Success(std::move(factors));
}
/// @brief Builds index-owned trials, then reduces all contributions stably.
Result<SolidAssembly> Assemble(const std::vector<RuntimeSolid>& solids,
                               const DofManager& dofs,
                               const Vector& displacement, bool nonlinear,
                               const ParallelFor& parallel) {
  std::vector<std::optional<Result<Hex8Response>>> responses(solids.size());
  parallel.Execute(solids.size(), [&](std::size_t e) {
    Vector local{24U};
    for (std::size_t i = 0U; i < 24U; ++i)
      local[i] = displacement[solids[e].scatter[i]];
    responses[e] = nonlinear ? solids[e].element->EvaluateNonlinear(local)
                             : solids[e].element->EvaluateLinear(local);
  });
  Vector force{dofs.FullDofCount()};
  std::vector<CooContribution> contributions;
  contributions.reserve(solids.size() * 24U * 24U);
  std::vector<SolidIntegrationPointRow> points;
  points.reserve(solids.size() * 8U);
  double energy = 0.0;
  for (std::size_t e = 0U; e < solids.size(); ++e) {
    if (!responses[e]->HasValue())
      return Result<SolidAssembly>::Failure(
          TrialFailure(responses[e]->GetStatus()));
    const auto& response = responses[e]->Value();
    for (std::size_t i = 0U; i < 24U; ++i) {
      force[solids[e].scatter[i]] += response.internal_force[i];
      for (std::size_t j = 0U; j < 24U; ++j)
        contributions.push_back({solids[e].scatter[i], solids[e].scatter[j],
                                 response.tangent(i, j), e, 24U * i + j});
    }
    energy += response.strain_energy;
    points.insert(points.end(), response.points.begin(), response.points.end());
  }
  auto tangent =
      SparseMatrix::FromCoo(dofs.FullDofCount(), dofs.FullDofCount(),
                            std::move(contributions), dofs.GetSparsePattern());
  if (!tangent.HasValue())
    return Result<SolidAssembly>::Failure(tangent.GetStatus());
  return Result<SolidAssembly>::Success({std::move(force),
                                         std::move(tangent.Value()),
                                         std::move(points), energy});
}
/// @brief Computes physical residual, projection and both moment-arm
/// identities.
Status CompleteFrame(const Domain& domain, const DofManager& dofs,
                     bool nonlinear, const SparseMatrix& stiffness,
                     SolidStaticFrame& frame) {
  const std::size_t size = dofs.FullDofCount();
  std::vector<bool> constrained(size, false);
  for (auto i : dofs.ConstrainedDofs()) constrained[i] = true;
  frame.residual = Vector{size};
  frame.reaction = Vector{size};
  for (std::size_t i = 0U; i < size; ++i) {
    frame.residual[i] = frame.internal_force[i] - frame.external_force[i];
    frame.reaction[i] = constrained[i] ? frame.residual[i] : 0.0;
    if (!std::isfinite(frame.displacement[i]) ||
        !std::isfinite(frame.internal_force[i]) ||
        !std::isfinite(frame.external_force[i]) ||
        !std::isfinite(frame.residual[i]))
      return Failure(domain, "nonfinite-solid-static-recovery",
                     "Static full-space quantities must remain finite.");
  }
  if (!std::isfinite(frame.physical_strain_energy))
    return Failure(domain, "nonfinite-solid-static-recovery",
                   "Recovered static strain energy must be finite.");
  std::array<double, 3> force_scale{};
  std::array<double, 3> reference_moment_scale{};
  std::array<double, 3> current_moment_scale{};
  std::array<double, 3> identity_moment_scale{};
  std::array<double, 3> free_force_sum{};
  std::array<double, 6> balance_scale{};
  Vector arithmetic_scale{size};
  // Linear K0 multiplication can cancel large translations within each row.
  // Use those original products, not the small resulting internal force.
  if (!nonlinear) {
    for (std::size_t i = 0U; i < size; ++i)
      for (std::size_t k = stiffness.RowOffsets()[i];
           k < stiffness.RowOffsets()[i + 1U]; ++k)
        arithmetic_scale[i] +=
            std::abs(stiffness.Values()[k] *
                     frame.displacement[stiffness.ColumnIndices()[k]]);
    for (auto i : dofs.FreeDofs())
      if (!RoundoffZero(
              frame.residual[i],
              arithmetic_scale[i] + std::abs(frame.external_force[i])))
        return Failure(
            domain, "solid-linear-residual-failed",
            "Linear free residual failed its uncancelled roundoff check.");
  }
  for (std::size_t a = 0U; a < domain.Nodes().size(); ++a) {
    const auto& reference = domain.Nodes()[a].coordinates;
    std::array<double, 3> current{};
    std::array<double, 3> u{};
    std::array<double, 3> f{};
    std::array<double, 3> e{};
    std::array<double, 3> applied{};
    for (std::size_t i = 0U; i < 3U; ++i) {
      const auto full = 3U * a + i;
      u[i] = frame.displacement[full];
      current[i] = reference[i] + u[i];
      f[i] = frame.internal_force[full];
      e[i] = constrained[full] ? 0.0 : frame.residual[full];
      applied[i] = frame.external_force[full] + frame.reaction[full];
      frame.balance.internal_force_sum[i] += f[i];
      free_force_sum[i] += e[i];
      force_scale[i] += nonlinear ? std::abs(f[i]) : arithmetic_scale[full];
      frame.global_equilibrium[i] += applied[i];
      balance_scale[i] += std::abs(frame.external_force[full]) +
                          std::abs(frame.reaction[full]) + std::abs(f[i]) +
                          std::abs(e[i]);
    }
    const auto reference_moment = Cross(reference, f, &reference_moment_scale);
    const auto current_moment = Cross(current, f, &current_moment_scale);
    const auto free_moment = Cross(current, e, &identity_moment_scale);
    const auto linear_moment = Cross(u, f, &identity_moment_scale);
    std::array<double, 3> applied_scale{};
    const auto applied_moment = Cross(current, applied, &applied_scale);
    for (std::size_t i = 0U; i < 3U; ++i) {
      frame.balance.reference_internal_moment[i] += reference_moment[i];
      frame.balance.current_internal_moment[i] += current_moment[i];
      frame.balance.free_residual_moment[i] += free_moment[i];
      frame.balance.linearization_moment[i] += linear_moment[i];
      frame.global_equilibrium[3U + i] += applied_moment[i];
      balance_scale[3U + i] += applied_scale[i];
      if (!nonlinear) {
        const auto j = (i + 1U) % 3U;
        const auto k = (i + 2U) % 3U;
        reference_moment_scale[i] +=
            std::abs(reference[j]) * arithmetic_scale[3U * a + k] +
            std::abs(reference[k]) * arithmetic_scale[3U * a + j];
      }
    }
  }
  frame.balance.internal_force_scale = force_scale;
  frame.balance.reference_moment_scale = reference_moment_scale;
  frame.balance.current_moment_scale = current_moment_scale;
  for (std::size_t i = 0U; i < 3U; ++i) {
    const double physical_moment =
        nonlinear ? frame.balance.current_internal_moment[i]
                  : frame.balance.reference_internal_moment[i];
    const double moment_scale =
        nonlinear ? current_moment_scale[i] : reference_moment_scale[i];
    if (!RoundoffZero(frame.balance.internal_force_sum[i], force_scale[i]) ||
        !RoundoffZero(physical_moment, moment_scale))
      return Failure(
          domain, "solid-static-internal-balance-failed",
          "Internal force or appropriate moment-arm identity failed.");
    const double force_identity = frame.global_equilibrium[i] -
                                  frame.balance.internal_force_sum[i] +
                                  free_force_sum[i];
    const double moment_identity = frame.global_equilibrium[3U + i] -
                                   frame.balance.current_internal_moment[i] +
                                   frame.balance.free_residual_moment[i];
    if (!RoundoffZero(force_identity, balance_scale[i]) ||
        !RoundoffZero(moment_identity, balance_scale[3U + i] +
                                           current_moment_scale[i] +
                                           identity_moment_scale[i]))
      return Failure(
          domain, "solid-static-residual-balance-failed",
          "Applied/support balance does not explain the free residual.");
    if (!std::isfinite(frame.global_equilibrium[i]) ||
        !std::isfinite(frame.global_equilibrium[3U + i]))
      return Failure(domain, "nonfinite-solid-static-recovery",
                     "Static current-position balance must remain finite.");
  }
  return Status::Ok();
}
}  // namespace

SolidStaticProcedureExecutor::SolidStaticProcedureExecutor(
    const ParallelFor& parallel_for, LinearSolver& linear_solver,
    TangentSolver& tangent_solver)
    : parallel_for_(parallel_for),
      linear_solver_(linear_solver),
      tangent_solver_(tangent_solver) {}

Status SolidStaticProcedureExecutor::ExecuteInto(
    const Domain& domain, SolidStaticState& state) const {
  if (state.initialized_)
    return Failure(domain, "solid-static-observer-not-fresh",
                   "Execution requires a fresh static observer.");
  if (!domain.IsSolid() || domain.Frequency() != nullptr ||
      domain.Dynamic() != nullptr || domain.Steps().Size() != 1U)
    return Failure(domain, "invalid-solid-static-procedure",
                   "Solid static execution requires one Solid STATIC step.",
                   FailureCategory::kModel);
  auto model = AnalysisModel::Create(domain);
  if (!model.HasValue()) return model.GetStatus();
  std::vector<std::unique_ptr<Element>> elements;
  ElementView view;
  for (auto e : model.Value().ActiveElements()) {
    auto element = ElementFactory{}.Create(domain.Elements()[e], domain);
    if (!element.HasValue()) return element.GetStatus();
    view.emplace_back(*element.Value());
    elements.push_back(std::move(element.Value()));
  }
  DofManager dofs;
  const Status mapping = dofs.Build(model.Value(), view);
  if (!mapping.IsOk()) return mapping;
  auto baseline = SolidStaticState::Create(domain, dofs);
  if (!baseline.HasValue()) return baseline.GetStatus();
  state = std::move(baseline.Value());
  std::vector<RuntimeSolid> solids;
  for (const auto& element : elements) {
    const auto* hex = dynamic_cast<const Hex8Element*>(element.get());
    if (hex == nullptr)
      return Failure(domain, "invalid-solid-static-element",
                     "Solid STATIC requires full-integration Hex8 kernels.",
                     FailureCategory::kModel);
    auto scatter = dofs.ElementScatter(hex->DofLayout());
    if (!scatter.HasValue()) return scatter.GetStatus();
    solids.push_back({hex, std::move(scatter.Value())});
  }
  auto stiffness = SparseAssembler::Assemble(view, dofs, parallel_for_);
  if (!stiffness.HasValue()) return stiffness.GetStatus();
  const EssentialConstraintPolicy constraints;
  const auto& step = model.Value().Step();
  const bool nonlinear = step.IsNonlinearStatic();
  // No support can remove any rigid mode; fail before a perturbed
  // factorization.
  if (dofs.ConstrainedDofCount() == 0U)
    return Failure(domain, "solid-static-rigid-mechanism",
                   "An unsupported Solid static free system has rigid modes.");
  if (!nonlinear) {
    auto partition = constraints.Partition(stiffness.Value(), dofs);
    if (!partition.HasValue()) return partition.GetStatus();
    Status status = linear_solver_.Factorize(partition.Value().k_ff);
    if (!status.IsOk()) return status;
    // Preserve the required factorization-before-load-assembly lifecycle.
    auto load = LoadAssembler::AssembleFullNodalLoad(model.Value(), dofs);
    if (!load.HasValue()) return load.GetStatus();
    auto rhs = LoadAssembler::EffectiveFreeRhs(
        load.Value(), partition.Value().k_fc, dofs.PrescribedValues(), dofs);
    if (!rhs.HasValue()) return rhs.GetStatus();
    Vector free{dofs.FreeDofCount()};
    status = linear_solver_.Solve(rhs.Value(), free);
    if (!status.IsOk()) return status;
    SolidStaticFrame frame = state.committed_;
    frame.frame_index = 0U;
    frame.increment_index = step.IsDirectStatic() ? step.IncrementCount() : 1U;
    frame.load_factor = 1.0;
    frame.is_factor_one = true;
    frame.displacement =
        constraints.ReconstructFull(free, dofs.PrescribedValues(), dofs);
    auto recovery =
        Assemble(solids, dofs, frame.displacement, false, parallel_for_);
    if (!recovery.HasValue()) return recovery.GetStatus();
    frame.internal_force = std::move(recovery.Value().force);
    frame.external_force = std::move(load.Value());
    frame.solid_results = std::move(recovery.Value().points);
    frame.physical_strain_energy = recovery.Value().energy;
    status = CompleteFrame(domain, dofs, false, stiffness.Value(), frame);
    return status.IsOk() ? state.CommitAcceptedFrame(frame) : status;
  }
  auto factors = Factors(domain);
  if (!factors.HasValue()) return factors.GetStatus();
  auto base_load = LoadAssembler::AssembleFullNodalLoad(model.Value(), dofs);
  if (!base_load.HasValue()) return base_load.GetStatus();
  for (const auto& increment : factors.Value()) {
    const Vector start = state.committed_.displacement;
    Vector trial = start;
    for (std::size_t c = 0U; c < dofs.ConstrainedDofCount(); ++c)
      trial[dofs.ConstrainedDofs()[c]] =
          increment.load_factor * dofs.PrescribedValues()[c];
    for (std::uint32_t iteration = 1U; iteration <= kMaximumIterations;
         ++iteration) {
      SolidStaticIterationRecord record;
      record.increment_index = increment.increment_index;
      record.load_factor = increment.load_factor;
      record.iteration = iteration;
      record.stage = "entered";
      const auto fail = [&](const Status& failure) {
        record.diagnostic_code = failure.Diagnostics().empty()
                                     ? "solid-static-failure"
                                     : failure.Diagnostics()[0].code;
        record.failure_reason = failure.Diagnostics().empty()
                                    ? "Static trial failed."
                                    : failure.Diagnostics()[0].message;
        state.iterations_.push_back(record);
        return failure;
      };
      record.stage = "trial-recovery";
      auto assembly = Assemble(solids, dofs, trial, true, parallel_for_);
      if (!assembly.HasValue()) return fail(assembly.GetStatus());
      SolidStaticFrame frame = state.committed_;
      frame.balance = {};
      frame.global_equilibrium = {};
      frame.frame_index =
          static_cast<std::uint32_t>(state.accepted_frames_.size());
      frame.increment_index = increment.increment_index;
      frame.load_factor = increment.load_factor;
      frame.iteration_count = iteration;
      frame.is_factor_one = increment.is_factor_one;
      frame.displacement = trial;
      frame.internal_force = std::move(assembly.Value().force);
      frame.external_force = base_load.Value();
      frame.external_force.Scale(increment.load_factor);
      frame.solid_results = std::move(assembly.Value().points);
      frame.physical_strain_energy = assembly.Value().energy;
      auto status = CompleteFrame(domain, dofs, true, stiffness.Value(), frame);
      if (!status.IsOk()) return fail(status);
      record.stage = "tangent-partition";
      auto partition = constraints.Partition(assembly.Value().tangent, dofs);
      if (!partition.HasValue()) return fail(partition.GetStatus());
      Vector correction{dofs.FreeDofCount()};
      // Empty free is vacuous; prescribed geometry/recovery above still runs.
      if (dofs.FreeDofCount() != 0U) {
        record.stage = "tangent-factorization";
        status = tangent_solver_.FactorizeFresh(partition.Value().k_ff);
        if (!status.IsOk()) return fail(status);
        auto rhs = constraints.GatherFree(frame.residual, dofs);
        rhs.Scale(-1.0);
        record.stage = "correction-solve";
        status = tangent_solver_.SolveCorrection(rhs, correction);
        if (!status.IsOk()) return fail(status);
      }
      Vector current_increment{dofs.FreeDofCount()};
      for (std::size_t f = 0U; f < dofs.FreeDofCount(); ++f) {
        const auto full = dofs.FreeDofs()[f];
        current_increment[f] = trial[full] - start[full];
      }
      record.stage = "convergence";
      auto convergence = EvaluateSolidStaticConvergence(
          constraints.GatherFree(frame.residual, dofs),
          constraints.GatherFree(frame.internal_force, dofs),
          constraints.GatherFree(frame.external_force, dofs), correction,
          current_increment);
      if (!convergence.HasValue()) return fail(convergence.GetStatus());
      record.force_residual = convergence.Value().force_residual;
      record.translation_correction =
          convergence.Value().translation_correction;
      record.convergence_evaluated = true;
      if (record.force_residual.passed &&
          record.translation_correction.passed) {
        record.stage = "commit";
        status = state.CommitAcceptedFrame(frame);
        if (!status.IsOk()) return fail(status);
        record.accepted = true;
        state.iterations_.push_back(record);
        break;
      }
      if (iteration == kMaximumIterations)
        return fail(Failure(domain, "solid-static-nonconvergence",
                            "Solid STATIC exceeded 100 Newton iterations."));
      for (std::size_t f = 0U; f < dofs.FreeDofCount(); ++f)
        trial[dofs.FreeDofs()[f]] += correction[f];
      record.stage = "continue";
      state.iterations_.push_back(record);
    }
  }
  return Status::Ok();
}
Result<SolidStaticExecutionResult> SolidStaticProcedureExecutor::Execute(
    const Domain& domain) const {
  SolidStaticState state;
  const Status status = ExecuteInto(domain, state);
  if (!status.IsOk())
    return Result<SolidStaticExecutionResult>::Failure(status);
  auto diagnostics = domain.Warnings();
  SortDiagnostics(diagnostics);
  return Result<SolidStaticExecutionResult>::Success(
      {std::move(state), std::move(diagnostics)});
}
}  // namespace fesa