#include "fesa/analysis/hht_dynamic_procedure.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <string>
#include <utility>

#include "analysis/hht_dynamic_internal.h"
#include "fesa/analysis/analysis_model.h"
#include "fesa/analysis/newton_iteration_log.h"
#include "fesa/solvers/linear/tangent_solver.h"

namespace fesa {
namespace {
using namespace hht_internal;

/// @brief Creates one procedure-owned diagnostic with the source step.
Status Failure(const Domain& domain, const std::string& code,
               const std::string& message) {
  return Status::Failure(
      FailureCategory::kSolver,
      {{Severity::kError, code, domain.Dynamic()->dynamic_location, "DYNAMIC",
        domain.Dynamic()->source_step_name_raw, message}});
}

/// @brief Checks one update against its uncancelled, dimensionally equal terms.
bool UpdateEqual(const double left, const std::initializer_list<double> terms) {
  double right = 0.0;
  double scale = 0.0;
  for (const double term : terms) {
    right += term;
    scale += std::abs(term);
  }
  const double defect = std::abs(left - right);
  return std::isfinite(defect) && std::isfinite(scale) &&
         (scale == 0.0 ? defect == 0.0 : defect / scale <= 1e-12);
}

/// @brief Verifies both absolute and increment forms before accepting history.
Status CheckUpdates(const Domain& domain, const DynamicAcceptedFrameState& old,
                    const DynamicAcceptedFrameState& frame,
                    const HhtPredictor& predictor, const HhtParameters& p,
                    const double h) {
  for (std::size_t i = 0U; i < frame.generalized_coordinates.Size(); ++i) {
    const double increment =
        frame.generalized_coordinates[i] - old.generalized_coordinates[i];
    const double old_term =
        h * h * (0.5 - p.beta) * old.algorithmic_acceleration[i];
    const double new_term = h * h * p.beta * frame.algorithmic_acceleration[i];
    if (!UpdateEqual(frame.generalized_coordinates[i],
                     {old.generalized_coordinates[i], h * old.velocity[i],
                      old_term, new_term}) ||
        !UpdateEqual(increment, {h * old.velocity[i], old_term, new_term}) ||
        !UpdateEqual(
            frame.generalized_coordinates[i] - predictor.coordinates[i],
            {new_term}) ||
        !UpdateEqual(frame.velocity[i],
                     {old.velocity[i],
                      h * (1 - p.gamma) * old.algorithmic_acceleration[i],
                      h * p.gamma * frame.algorithmic_acceleration[i]}) ||
        !UpdateEqual(frame.algorithmic_acceleration[i],
                     {(1 + p.alpha) * frame.acceleration[i],
                      -p.alpha * old.acceleration[i]})) {
      return Failure(
          domain, "dynamic-update-inconsistency",
          "HHT endpoint update or resolved increment exceeds 1e-12.");
    }
  }
  return Status::Ok();
}

/// @brief Fills the raw channel evidence without hiding a zero denominator.
void ResidualRecord(const DofManager& dofs,
                    const DynamicAcceptedFrameState& frame,
                    DynamicNewtonIterationLogRecord* record) {
  record->residual_numerator = {};
  record->residual_scale = {};
  for (const auto full : dofs.FreeDofs()) {
    const auto c = full % 6U < 3U ? 0U : 1U;
    record->residual_numerator[c] = (std::max)(record->residual_numerator[c],
                                               std::abs(frame.residual[full]));
    record->residual_scale[c] =
        (std::max)(record->residual_scale[c], frame.residual_scale[full]);
  }
  record->residual_pass = frame.free_residual_metrics[0U] <= 1e-10 &&
                          frame.free_residual_metrics[1U] <= 1e-10;
}

/// @brief Forms the four coordinate/rate correction checks in scaled units.
void CorrectionRecord(const Vector& correction, const DofManager& dofs,
                      const double length, const HhtParameters& p,
                      DynamicNewtonIterationLogRecord* record) {
  record->coordinate_correction = {};
  for (const auto full : dofs.FreeDofs()) {
    const auto c = full % 6U < 3U ? 0U : 1U;
    record->coordinate_correction[c] =
        (std::max)(record->coordinate_correction[c],
                   std::abs(correction[full]) / (c == 0U ? length : 1.0));
  }
  record->correction_pass = true;
  for (std::size_t c = 0; c < 2U; ++c) {
    const double dq = record->coordinate_correction[c];
    record->velocity_correction[c] = p.gamma / p.beta * dq;
    record->acceleration_correction[c] = dq / ((1 + p.alpha) * p.beta);
    record->algorithmic_correction[c] = dq / p.beta;
    for (const auto* channel :
         {&record->coordinate_correction, &record->velocity_correction,
          &record->acceleration_correction, &record->algorithmic_correction})
      record->correction_pass =
          record->correction_pass && (*channel)[c] <= 1e-10;
  }
}
}  // namespace

HhtDynamicProcedureExecutor::HhtDynamicProcedureExecutor(
    const ParallelFor& parallel_for, LinearSolver& positive_solver,
    TangentSolver& tangent_solver)
    : parallel_for_(parallel_for),
      positive_solver_(positive_solver),
      tangent_solver_(tangent_solver) {}

Result<DynamicExecutionResult> HhtDynamicProcedureExecutor::Execute(
    const Domain& domain, NewtonIterationLogWriter& log) const {
  if (domain.Dynamic() == nullptr) {
    return Result<DynamicExecutionResult>::Failure(Status::Failure(
        FailureCategory::kModel, {{Severity::kError,
                                   "invalid-dynamic-step",
                                   {},
                                   "DYNAMIC",
                                   {},
                                   "HHT requires exactly one dynamic step."}}));
  }
  auto model = AnalysisModel::Create(domain);
  if (!model.HasValue())
    return Result<DynamicExecutionResult>::Failure(model.GetStatus());
  auto dofs_result = DofManager::Create(model.Value());
  if (!dofs_result.HasValue())
    return Result<DynamicExecutionResult>::Failure(dofs_result.GetStatus());
  const auto& dofs = dofs_result.Value();
  auto state_result = DynamicAnalysisState::Create(domain, dofs);
  if (!state_result.HasValue())
    return Result<DynamicExecutionResult>::Failure(state_result.GetStatus());
  auto state = std::move(state_result.Value());
  auto shells = BuildRuntimeShells(domain, dofs);
  if (!shells.HasValue())
    return Result<DynamicExecutionResult>::Failure(shells.GetStatus());
  auto length_result = ReferenceLength(domain);
  if (!length_result.HasValue())
    return Result<DynamicExecutionResult>::Failure(length_result.GetStatus());
  const double length = length_result.Value();
  auto loads = AssembleLoads(domain, dofs);
  if (!loads.HasValue())
    return Result<DynamicExecutionResult>::Failure(loads.GetStatus());
  auto parameters = MakeHhtParameters(domain.Dynamic()->alpha);
  if (!parameters.HasValue())
    return Result<DynamicExecutionResult>::Failure(parameters.GetStatus());
  const auto p = parameters.Value();
  const double h = domain.Dynamic()->initial_increment;
  DynamicAcceptedFrameState frame{};
  frame.is_initial = true;
  for (auto* vector : {&frame.generalized_coordinates, &frame.velocity,
                       &frame.acceleration, &frame.algorithmic_acceleration})
    *vector = Vector(dofs.FullDofCount());
  frame.external_force = loads.Value()[0U];
  auto initial = EvaluateAssembly(domain, dofs, shells.Value(), frame,
                                  loads.Value()[1U], h, p, parallel_for_, true);
  if (!initial.HasValue())
    return Result<DynamicExecutionResult>::Failure(initial.GetStatus());
  const auto null_nodes = NullNodes(dofs, initial.Value().rotations);
  auto acceleration = InitializeAcceleration(
      domain, dofs, initial.Value(), length, positive_solver_, tangent_solver_);
  if (!acceleration.HasValue())
    return Result<DynamicExecutionResult>::Failure(acceleration.GetStatus());
  frame.acceleration = acceleration.Value();
  frame.algorithmic_acceleration = acceleration.Value();
  initial = EvaluateAssembly(domain, dofs, shells.Value(), frame,
                             loads.Value()[1U], h, p, parallel_for_, true);
  if (!initial.HasValue())
    return Result<DynamicExecutionResult>::Failure(initial.GetStatus());
  frame = initial.Value().frame;
  Status status = CompleteConstraintEvidence(
      domain, dofs, initial.Value(), length, null_nodes, positive_solver_,
      tangent_solver_, &frame);
  if (!status.IsOk()) return Result<DynamicExecutionResult>::Failure(status);
  status = CompleteEquilibrium(domain, dofs, initial.Value(), &frame);
  if (!status.IsOk()) return Result<DynamicExecutionResult>::Failure(status);
  const double initial_energy = frame.kinetic_energy +
                                frame.physical_strain_energy +
                                frame.drilling_energy;
  status = state.CommitAcceptedFrame(frame);
  if (!status.IsOk()) return Result<DynamicExecutionResult>::Failure(status);
  DynamicNewtonIterationLogRecord record{};
  record.stage = "accepted";
  record.outcome = "initial";
  record.correction_pass = true;
  ResidualRecord(dofs, frame, &record);
  status = log.Write(record);
  if (status.IsOk()) status = log.Flush();
  if (!status.IsOk()) return Result<DynamicExecutionResult>::Failure(status);
  for (std::uint32_t increment = 1U;
       increment <= domain.Dynamic()->increment_count; ++increment) {
    const auto& old = state.CommittedState();
    auto predictor = PredictHht(old.generalized_coordinates, old.velocity,
                                old.algorithmic_acceleration, h, p);
    if (!predictor.HasValue())
      return Result<DynamicExecutionResult>::Failure(predictor.GetStatus());
    frame = old;
    frame.is_initial = false;
    frame.frame_index = increment;
    frame.increment_index = increment;
    frame.time = static_cast<double>(increment) * h;
    // A large DIRECT step may put a kinematic predictor outside the shell
    // chart before Newton starts. The last accepted geometry is a valid
    // initial iterate; q_p still defines the unchanged HHT rate equations.
    frame.generalized_coordinates = old.generalized_coordinates;
    record = {};
    record.increment_index = increment;
    record.time = frame.time;
    record.correction_pass = true;
    bool accepted = false;
    for (std::uint32_t iteration = 0U; iteration <= 100U; ++iteration) {
      record.iteration_index = iteration;
      record.stage = "evaluation";
      record.outcome = "candidate";
      auto fail = [&](const Status& failure) -> Result<DynamicExecutionResult> {
        record.stage = "failed";
        record.outcome = "failed";
        record.failure_code = failure.Diagnostics().empty()
                                  ? "dynamic-solver-failure"
                                  : failure.Diagnostics().front().code;
        const Status written = log.Write(record);
        return Result<DynamicExecutionResult>::Failure(
            written.IsOk() ? failure : written);
      };
      auto rates = EvaluateHhtRates(frame.generalized_coordinates,
                                    predictor.Value(), old.acceleration, h, p);
      if (!rates.HasValue()) return fail(rates.GetStatus());
      frame.velocity = std::move(rates.Value().velocity);
      frame.acceleration = std::move(rates.Value().acceleration);
      frame.algorithmic_acceleration =
          std::move(rates.Value().algorithmic_acceleration);
      frame.iteration_count = iteration;
      auto assembly =
          EvaluateAssembly(domain, dofs, shells.Value(), frame,
                           loads.Value()[1U], h, p, parallel_for_, false);
      if (!assembly.HasValue()) return fail(assembly.GetStatus());
      frame = assembly.Value().frame;
      ResidualRecord(dofs, frame, &record);
      if (record.residual_pass && record.correction_pass) {
        assembly =
            EvaluateAssembly(domain, dofs, shells.Value(), frame,
                             loads.Value()[1U], h, p, parallel_for_, true);
        if (!assembly.HasValue()) return fail(assembly.GetStatus());
        frame = assembly.Value().frame;
        status = CompleteConstraintEvidence(
            domain, dofs, assembly.Value(), length, null_nodes,
            positive_solver_, tangent_solver_, &frame);
        if (!status.IsOk()) return fail(status);
        status = CompleteEquilibrium(domain, dofs, assembly.Value(), &frame);
        if (!status.IsOk()) return fail(status);
        status = CheckUpdates(domain, old, frame, predictor.Value(), p, h);
        if (!status.IsOk()) return fail(status);
        frame.external_work = 0.0;
        for (std::size_t i = 0; i < dofs.FullDofCount(); ++i)
          frame.external_work +=
              frame.external_force[i] * frame.generalized_coordinates[i];
        frame.energy_defect =
            frame.kinetic_energy + frame.physical_strain_energy +
            frame.drilling_energy - initial_energy - frame.external_work;
        status = state.CommitAcceptedFrame(frame);
        if (!status.IsOk()) return fail(status);
        record.stage = "accepted";
        record.outcome = "accepted";
        status = log.Write(record);
        if (status.IsOk()) status = log.Flush();
        if (!status.IsOk())
          return Result<DynamicExecutionResult>::Failure(status);
        accepted = true;
        break;
      }
      if (iteration == 100U)
        return fail(Failure(domain, "dynamic-nonconvergence",
                            "HHT exhausted 100 corrections without satisfying "
                            "every endpoint check."));
      status = log.Write(record);
      if (!status.IsOk())
        return Result<DynamicExecutionResult>::Failure(status);
      auto tangent =
          ScaledFree(assembly.Value().effective_tangent, dofs, length);
      if (!tangent.HasValue()) return fail(tangent.GetStatus());
      Vector rhs(dofs.FreeDofCount());
      Vector solution(dofs.FreeDofCount());
      Vector correction(dofs.FullDofCount());
      for (std::size_t i = 0; i < dofs.FreeDofCount(); ++i) {
        const auto full = dofs.FreeDofs()[i];
        rhs[i] = -frame.residual[full] * (full % 6U < 3U ? length : 1.0);
      }
      if (dofs.FreeDofCount() > 0U) {
        status = tangent_solver_.FactorizeFresh(tangent.Value());
        if (!status.IsOk()) return fail(status);
        status = tangent_solver_.SolveCorrection(rhs, solution);
        if (!status.IsOk()) return fail(status);
        status = CheckLinearSolve(tangent.Value(), rhs, solution);
        if (!status.IsOk()) return fail(status);
      }
      for (std::size_t i = 0; i < dofs.FreeDofCount(); ++i) {
        const auto full = dofs.FreeDofs()[i];
        correction[full] = solution[i] * (full % 6U < 3U ? length : 1.0);
        frame.generalized_coordinates[full] += correction[full];
      }
      CorrectionRecord(correction, dofs, length, p, &record);
    }
    if (!accepted)
      return Result<DynamicExecutionResult>::Failure(
          Failure(domain, "dynamic-nonconvergence",
                  "Dynamic increment was not accepted."));
  }
  return Result<DynamicExecutionResult>::Success(
      {std::move(state), domain.Warnings()});
}
}  // namespace fesa