#include "fesa/analysis/nonlinear_static_control.h"

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

#include "fesa/analysis/analysis_model.h"
#include "fesa/assembly/load_assembler.h"
#include "fesa/elements/finite_rotation.h"

namespace fesa {
namespace {

constexpr std::size_t kDofsPerNode = 6U;
constexpr std::size_t kTranslationDofCount = 3U;
constexpr double kMomentAdmissibilityThreshold = 1.0e-12;

/// @brief Creates a deterministic control-layer diagnostic status.
Status ControlFailure(const FailureCategory category, const std::string& code,
                      const SourceLocation& location,
                      const std::string& keyword, const std::string& identity,
                      const std::string& message) {
  return Status::Failure(category, {{Severity::kError, code, location, keyword,
                                     identity, message}});
}

/// @brief Verifies raw nonlinear static controls without deriving state.
Status ValidateNonlinearStaticControls(const StepDefinition& step) {
  if (!step.IsNonlinearStatic() || !std::isfinite(step.InitialIncrement()) ||
      !(step.InitialIncrement() > 0.0) || !std::isfinite(step.TimePeriod()) ||
      step.TimePeriod() != 1.0 || !std::isfinite(step.MinimumIncrement()) ||
      !(step.MinimumIncrement() > 0.0) ||
      !std::isfinite(step.MaximumIncrement()) ||
      !(step.MaximumIncrement() > 0.0)) {
    return ControlFailure(
        FailureCategory::kInput, "invalid-nonlinear-static-data",
        step.Location(), "STATIC", step.Name(),
        "A nonlinear static step requires four positive finite controls and "
        "an exactly factor-one time period.");
  }
  return Status::Ok();
}

/// @brief Returns a finite three-component vector from full trial storage.
Result<Vector3> RotationVectorForNode(const Vector& trial,
                                      const std::size_t node) {
  const std::size_t first = node * kDofsPerNode + kTranslationDofCount;
  if (first + kTranslationDofCount > trial.Size()) {
    return Result<Vector3>::Failure(ControlFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {{}, 0U},
        "NONLINEAR_STATIC_CONTROL", std::to_string(node),
        "The full trial coordinate vector does not contain every nodal "
        "rotation component."));
  }
  const Vector3 rotation{trial[first], trial[first + 1U], trial[first + 2U]};
  if (!rotation.IsFinite()) {
    return Result<Vector3>::Failure(ControlFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {{}, 0U},
        "NONLINEAR_STATIC_CONTROL", std::to_string(node),
        "The current trial rotation coordinate must be finite."));
  }
  return Result<Vector3>::Success(rotation);
}

/// @brief Builds a proper initial frame from its semantic nodal record.
Matrix InitialFrame(const ShellNodeInitialFrame& frame) {
  Matrix result{3U, 3U};
  for (std::size_t row = 0U; row < kTranslationDofCount; ++row) {
    result(row, 0U) = frame.tangent_a[row];
    result(row, 1U) = frame.tangent_b[row];
    result(row, 2U) = frame.director[row];
  }
  return result;
}

/// @brief Multiplies a 3-by-3 matrix by a Vector3.
Vector3 Multiply(const Matrix& matrix, const Vector3& vector) {
  return {matrix(0U, 0U) * vector.X() + matrix(0U, 1U) * vector.Y() +
              matrix(0U, 2U) * vector.Z(),
          matrix(1U, 0U) * vector.X() + matrix(1U, 1U) * vector.Y() +
              matrix(1U, 2U) * vector.Z(),
          matrix(2U, 0U) * vector.X() + matrix(2U, 1U) * vector.Y() +
              matrix(2U, 2U) * vector.Z()};
}

/// @brief Multiplies a 3-by-3 matrix transpose by a Vector3.
Vector3 MultiplyTranspose(const Matrix& matrix, const Vector3& vector) {
  return {matrix(0U, 0U) * vector.X() + matrix(1U, 0U) * vector.Y() +
              matrix(2U, 0U) * vector.Z(),
          matrix(0U, 1U) * vector.X() + matrix(1U, 1U) * vector.Y() +
              matrix(2U, 1U) * vector.Z(),
          matrix(0U, 2U) * vector.X() + matrix(1U, 2U) * vector.Y() +
              matrix(2U, 2U) * vector.Z()};
}

/// @brief Forms the three-dimensional skew matrix of a global vector.
Matrix Skew(const Vector3& vector) {
  Matrix result{3U, 3U};
  result(0U, 1U) = -vector.Z();
  result(0U, 2U) = vector.Y();
  result(1U, 0U) = vector.Z();
  result(1U, 2U) = -vector.X();
  result(2U, 0U) = -vector.Y();
  result(2U, 1U) = vector.X();
  return result;
}

/// @brief Scales a small dense matrix without changing its shape.
Matrix Scale(const Matrix& matrix, const double scale) {
  Matrix result{matrix.Rows(), matrix.Columns()};
  for (std::size_t row = 0U; row < matrix.Rows(); ++row) {
    for (std::size_t column = 0U; column < matrix.Columns(); ++column) {
      result(row, column) = scale * matrix(row, column);
    }
  }
  return result;
}

/// @brief Adds matching small dense matrices.
Matrix Add(const Matrix& left, const Matrix& right) {
  Matrix result{left.Rows(), left.Columns()};
  for (std::size_t row = 0U; row < left.Rows(); ++row) {
    for (std::size_t column = 0U; column < left.Columns(); ++column) {
      result(row, column) = left(row, column) + right(row, column);
    }
  }
  return result;
}

/// @brief Evaluates left-Jacobian coefficients and derivatives by theta
/// squared.
void LeftJacobianCoefficients(const double theta_squared, double& a, double& b,
                              double& da_dtheta_squared,
                              double& db_dtheta_squared) {
  if (theta_squared < 1.0e-8) {
    const double theta_fourth = theta_squared * theta_squared;
    const double theta_sixth = theta_fourth * theta_squared;
    a = 0.5 - theta_squared / 24.0 + theta_fourth / 720.0 -
        theta_sixth / 40320.0;
    b = 1.0 / 6.0 - theta_squared / 120.0 + theta_fourth / 5040.0 -
        theta_sixth / 362880.0;
    da_dtheta_squared =
        -1.0 / 24.0 + theta_squared / 360.0 - theta_fourth / 13440.0;
    db_dtheta_squared =
        -1.0 / 120.0 + theta_squared / 2520.0 - theta_fourth / 120960.0;
    return;
  }

  const double theta = std::sqrt(theta_squared);
  const double sine = std::sin(theta);
  const double cosine = std::cos(theta);
  a = (1.0 - cosine) / theta_squared;
  b = (theta - sine) / (theta_squared * theta);
  da_dtheta_squared = (theta * sine - 2.0 * (1.0 - cosine)) /
                      (2.0 * theta_squared * theta_squared);
  db_dtheta_squared =
      (1.0 - cosine) / (2.0 * theta_squared * theta_squared) -
      3.0 * (theta - sine) / (2.0 * theta_squared * theta_squared * theta);
}

/// @brief Differentiates the SO(3) left Jacobian along one coordinate vector.
Matrix LeftJacobianDerivative(const Vector3& phi, const Vector3& direction) {
  const Matrix skew = Skew(phi);
  const Matrix direction_skew = Skew(direction);
  const Matrix skew_squared = skew.Multiply(skew);
  double a = 0.0;
  double b = 0.0;
  double da_dtheta_squared = 0.0;
  double db_dtheta_squared = 0.0;
  LeftJacobianCoefficients(phi.Dot(phi), a, b, da_dtheta_squared,
                           db_dtheta_squared);
  const double derivative_scale = 2.0 * phi.Dot(direction);
  return Add(Add(Scale(skew, derivative_scale * da_dtheta_squared),
                 Scale(direction_skew, a)),
             Add(Scale(skew_squared, derivative_scale * db_dtheta_squared),
                 Scale(Add(direction_skew.Multiply(skew),
                           skew.Multiply(direction_skew)),
                       b)));
}

/// @brief Projects the immutable source moment into current director tangent
/// space.
Vector3 ProjectMoment(const Vector3& director, const Vector3& source_moment) {
  return source_moment - director.Dot(source_moment) * director;
}

/// @brief Differentiates the projected moment along a coordinate direction.
Vector3 ProjectMomentDerivative(const Vector3& director,
                                const Vector3& director_derivative,
                                const Vector3& source_moment) {
  return -director.Dot(source_moment) * director_derivative -
         director_derivative.Dot(source_moment) * director;
}

/// @brief Computes the infinity norm and reports finite input validity.
double InfinityNorm(const Vector& values, bool& is_finite) {
  double result = 0.0;
  for (std::size_t index = 0U; index < values.Size(); ++index) {
    if (!std::isfinite(values[index])) {
      is_finite = false;
      return std::numeric_limits<double>::quiet_NaN();
    }
    result = (std::max)(result, std::abs(values[index]));
  }
  return result;
}

/// @brief Evaluates one residual channel with its fixed physical zero-scale
/// unit.
NewtonChannelDecision EvaluateResidualChannel(
    const NewtonResidualChannelInput& input, const NewtonChannelUnit unit) {
  bool is_finite = input.free_residual.Size() == input.free_internal.Size() &&
                   input.free_residual.Size() == input.free_external.Size();
  const double numerator = InfinityNorm(input.free_residual, is_finite);
  const double internal = InfinityNorm(input.free_internal, is_finite);
  const double external = InfinityNorm(input.free_external, is_finite);
  if (!is_finite) {
    return {input.identity,
            numerator,
            external,
            NewtonConvergenceBranch::kZeroScale,
            kNewtonZeroScaleThreshold,
            unit,
            false};
  }
  if (external == 0.0) {
    return {input.identity,
            numerator,
            external,
            NewtonConvergenceBranch::kZeroScale,
            kNewtonZeroScaleThreshold,
            unit,
            numerator <= kNewtonZeroScaleThreshold};
  }
  const double denominator = (std::max)(internal, external);
  return {input.identity,
          numerator,
          denominator,
          NewtonConvergenceBranch::kRelative,
          kNewtonRelativeThreshold,
          NewtonChannelUnit::kDimensionless,
          numerator / denominator <= kNewtonRelativeThreshold};
}

/// @brief Evaluates one correction channel with its fixed physical zero-scale
/// unit.
NewtonChannelDecision EvaluateCorrectionChannel(
    const NewtonCorrectionChannelInput& input, const NewtonChannelUnit unit) {
  bool is_finite =
      input.fresh_correction.Size() == input.current_increment.Size();
  const double numerator = InfinityNorm(input.fresh_correction, is_finite);
  const double denominator = InfinityNorm(input.current_increment, is_finite);
  if (!is_finite) {
    return {input.identity,
            numerator,
            denominator,
            NewtonConvergenceBranch::kZeroScale,
            kNewtonZeroScaleThreshold,
            unit,
            false};
  }
  if (denominator == 0.0) {
    return {input.identity,
            numerator,
            denominator,
            NewtonConvergenceBranch::kZeroScale,
            kNewtonZeroScaleThreshold,
            unit,
            numerator <= kNewtonZeroScaleThreshold};
  }
  return {input.identity,
          numerator,
          denominator,
          NewtonConvergenceBranch::kRelative,
          kNewtonRelativeThreshold,
          NewtonChannelUnit::kDimensionless,
          numerator / denominator <= kNewtonRelativeThreshold};
}

}  // namespace

Result<std::vector<LoadFactorIncrement>> BuildFixedLoadFactorSequence(
    const StepDefinition& step) {
  const Status control_status = ValidateNonlinearStaticControls(step);
  if (!control_status.IsOk()) {
    return Result<std::vector<LoadFactorIncrement>>::Failure(control_status);
  }

  std::vector<LoadFactorIncrement> increments;
  std::uint32_t increment_index = 1U;
  while (true) {
    const double ordinal_factor =
        static_cast<double>(increment_index) * step.InitialIncrement();
    if (!std::isfinite(ordinal_factor)) {
      return Result<std::vector<LoadFactorIncrement>>::Failure(ControlFailure(
          FailureCategory::kInput, "invalid-nonlinear-static-data",
          step.Location(), "STATIC", step.Name(),
          "The fixed load-factor sequence must remain finite."));
    }
    if (!(ordinal_factor < 1.0)) {
      increments.push_back({increment_index, 1.0, true});
      break;
    }
    increments.push_back({increment_index, ordinal_factor, false});
    if (increment_index == (std::numeric_limits<std::uint32_t>::max)()) {
      return Result<std::vector<LoadFactorIncrement>>::Failure(ControlFailure(
          FailureCategory::kInput, "invalid-nonlinear-static-data",
          step.Location(), "STATIC", step.Name(),
          "The fixed load-factor sequence exceeds the supported increment "
          "identity range."));
    }
    ++increment_index;
  }
  return Result<std::vector<LoadFactorIncrement>>::Success(
      std::move(increments));
}

Result<Vector> ApplyAbsolutePrescribedDisplacementTarget(
    const StepDefinition& step, const double load_factor) {
  const Status control_status = ValidateNonlinearStaticControls(step);
  if (!control_status.IsOk()) {
    return Result<Vector>::Failure(control_status);
  }
  if (!std::isfinite(load_factor) || load_factor < 0.0 || load_factor > 1.0) {
    return Result<Vector>::Failure(ControlFailure(
        FailureCategory::kInput, "invalid-nonlinear-static-data",
        step.Location(), "BOUNDARY", step.Name(),
        "The prescribed-displacement path requires a finite load factor in "
        "the closed interval from zero to one."));
  }

  const auto& boundaries = step.PrescribedDisplacements();
  Vector targets{boundaries.Size()};
  for (std::size_t index = 0U; index < boundaries.Size(); ++index) {
    const double target = boundaries[index].PrescribedValue();
    const double value = load_factor * target;
    if (!std::isfinite(target) || !std::isfinite(value)) {
      return Result<Vector>::Failure(ControlFailure(
          FailureCategory::kInput, "invalid-nonlinear-static-data",
          boundaries[index].Location(), "BOUNDARY",
          boundaries[index].Target().target_name_or_label,
          "The factor-one prescribed displacement target must remain finite."));
    }
    targets[index] = value;
  }
  return Result<Vector>::Success(std::move(targets));
}

Status CheckMomentAdmissibility(const Vector3& director,
                                const Vector3& source_moment) {
  if (!source_moment.IsFinite()) {
    return ControlFailure(FailureCategory::kModel, "nonfinite-newton-data",
                          {{}, 0U}, "CLOAD", "",
                          "The fixed-global source moment must be finite.");
  }
  if (source_moment == Vector3{}) {
    return Status::Ok();
  }
  if (!director.IsFinite() || !(director.Norm() > 0.0)) {
    return ControlFailure(
        FailureCategory::kSolver, "nonfinite-newton-data", {{}, 0U}, "CLOAD",
        "", "The current shell director must be finite and nonzero.");
  }
  const double source_norm = source_moment.Norm();
  const double rho = std::abs(director.Dot(source_moment)) / source_norm;
  if (!std::isfinite(source_norm) || !(source_norm > 0.0) ||
      !std::isfinite(rho)) {
    return ControlFailure(
        FailureCategory::kModel, "nonfinite-newton-data", {{}, 0U}, "CLOAD", "",
        "The fixed-global source moment norm must be finite.");
  }
  if (!(rho <= kMomentAdmissibilityThreshold)) {
    return ControlFailure(
        FailureCategory::kModel, "unsupported-drilling-load", {{}, 0U}, "CLOAD",
        "",
        "The aggregate nodal moment has an unsupported director-parallel "
        "component.");
  }
  return Status::Ok();
}

Result<GlobalFixedLoadContribution> AssembleGlobalFixedLoadContribution(
    const Domain& domain, const DofManager& dofs,
    const Vector& trial_generalized_coordinates, const double load_factor) {
  if (!std::isfinite(load_factor) || load_factor < 0.0 || load_factor > 1.0 ||
      domain.Steps().Size() != 1U || !domain.Steps()[0U].IsNonlinearStatic()) {
    return Result<GlobalFixedLoadContribution>::Failure(ControlFailure(
        FailureCategory::kModel, "invalid-nonlinear-static-data",
        {domain.SourcePath(), 0U}, "NONLINEAR_STATIC_CONTROL",
        domain.SourceContentIdentity(),
        "Fixed-load assembly requires one nonlinear step and a finite load "
        "factor in the closed interval from zero to one."));
  }
  const Status dof_status = dofs.ValidateInvariants();
  if (!dof_status.IsOk()) {
    return Result<GlobalFixedLoadContribution>::Failure(dof_status);
  }
  if (dofs.FullDofCount() != trial_generalized_coordinates.Size() ||
      domain.Nodes().size() >
          (std::numeric_limits<std::size_t>::max)() / kDofsPerNode ||
      dofs.FullDofCount() != domain.Nodes().size() * kDofsPerNode) {
    return Result<GlobalFixedLoadContribution>::Failure(ControlFailure(
        FailureCategory::kSolver, "nonfinite-newton-data",
        {domain.SourcePath(), 0U}, "NONLINEAR_STATIC_CONTROL",
        domain.SourceContentIdentity(),
        "The full trial coordinates must match the stable Domain DOF "
        "dimension."));
  }
  for (std::size_t index = 0U; index < trial_generalized_coordinates.Size();
       ++index) {
    if (!std::isfinite(trial_generalized_coordinates[index])) {
      return Result<GlobalFixedLoadContribution>::Failure(ControlFailure(
          FailureCategory::kSolver, "nonfinite-newton-data",
          {domain.SourcePath(), 0U}, "NONLINEAR_STATIC_CONTROL",
          std::to_string(index),
          "The full trial coordinate vector must contain finite values."));
    }
  }

  auto model = AnalysisModel::Create(domain);
  if (!model.HasValue()) {
    return Result<GlobalFixedLoadContribution>::Failure(model.GetStatus());
  }
  auto source_load = LoadAssembler::AssembleFullNodalLoad(model.Value(), dofs);
  if (!source_load.HasValue()) {
    return Result<GlobalFixedLoadContribution>::Failure(
        source_load.GetStatus());
  }

  Vector external_force{dofs.FullDofCount()};
  for (std::size_t index = 0U; index < external_force.Size(); ++index) {
    external_force[index] = load_factor * source_load.Value()[index];
  }

  std::vector<const ShellNodeInitialFrame*> frames(domain.Nodes().size(),
                                                   nullptr);
  for (const auto& frame : domain.ShellNodeInitialFrames()) {
    if (frame.node_index >= frames.size() ||
        frames[frame.node_index] != nullptr) {
      return Result<GlobalFixedLoadContribution>::Failure(ControlFailure(
          FailureCategory::kModel, "invalid-shell-director",
          {domain.SourcePath(), 0U}, "NODE", std::to_string(frame.node_index),
          "Shell nodal frames must provide one unique record per node."));
    }
    frames[frame.node_index] = &frame;
  }

  std::vector<CooContribution> tangent_contributions;
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const std::size_t first_rotation =
        node * kDofsPerNode + kTranslationDofCount;
    const Vector3 source_moment{source_load.Value()[first_rotation],
                                source_load.Value()[first_rotation + 1U],
                                source_load.Value()[first_rotation + 2U]};
    if (source_moment == Vector3{}) {
      continue;
    }
    if (frames[node] == nullptr) {
      return Result<GlobalFixedLoadContribution>::Failure(ControlFailure(
          FailureCategory::kModel, "invalid-shell-director",
          domain.Nodes()[node].location, "NODE",
          domain.Nodes()[node].source_id.source_label_text,
          "A node carrying a moment requires an initial shell frame."));
    }
    auto phi = RotationVectorForNode(trial_generalized_coordinates, node);
    if (!phi.HasValue()) {
      return Result<GlobalFixedLoadContribution>::Failure(phi.GetStatus());
    }
    auto rotation =
        MakeFiniteRotationState(phi.Value(), InitialFrame(*frames[node]));
    if (!rotation.HasValue()) {
      return Result<GlobalFixedLoadContribution>::Failure(rotation.GetStatus());
    }
    const Status admissibility =
        CheckMomentAdmissibility(rotation.Value().director, source_moment);
    if (!admissibility.IsOk()) {
      return Result<GlobalFixedLoadContribution>::Failure(admissibility);
    }

    const Matrix left_jacobian = LeftJacobianSO3(phi.Value());
    const Vector3 projected =
        ProjectMoment(rotation.Value().director, source_moment);
    const Vector3 rotational_force =
        load_factor * MultiplyTranspose(left_jacobian, projected);
    for (std::size_t row = 0U; row < kTranslationDofCount; ++row) {
      external_force[first_rotation + row] = rotational_force[row];
    }

    for (std::size_t column = 0U; column < kTranslationDofCount; ++column) {
      Vector3 direction{};
      if (column == 0U) {
        direction = {1.0, 0.0, 0.0};
      } else if (column == 1U) {
        direction = {0.0, 1.0, 0.0};
      } else {
        direction = {0.0, 0.0, 1.0};
      }
      const Matrix left_jacobian_derivative =
          LeftJacobianDerivative(phi.Value(), direction);
      const Vector3 director_derivative =
          Multiply(left_jacobian, direction).Cross(rotation.Value().director);
      const Vector3 projected_derivative = ProjectMomentDerivative(
          rotation.Value().director, director_derivative, source_moment);
      const Vector3 tangent_column =
          load_factor *
          (MultiplyTranspose(left_jacobian_derivative, projected) +
           MultiplyTranspose(left_jacobian, projected_derivative));
      for (std::size_t row = 0U; row < kTranslationDofCount; ++row) {
        tangent_contributions.push_back(
            {first_rotation + row, first_rotation + column, tangent_column[row],
             node, row * kTranslationDofCount + column});
      }
    }
  }

  auto tangent = SparseMatrix::FromCoo(dofs.FullDofCount(), dofs.FullDofCount(),
                                       std::move(tangent_contributions),
                                       dofs.GetSparsePattern());
  if (!tangent.HasValue()) {
    return Result<GlobalFixedLoadContribution>::Failure(tangent.GetStatus());
  }
  return Result<GlobalFixedLoadContribution>::Success(
      {std::move(external_force), std::move(tangent.Value()), {}});
}

NewtonConvergenceResult EvaluateNewtonConvergence(
    const NewtonConvergenceInput& input) {
  NewtonConvergenceResult result{
      EvaluateResidualChannel(input.force_residual, NewtonChannelUnit::kForce),
      EvaluateResidualChannel(input.moment_residual,
                              NewtonChannelUnit::kForceLength),
      EvaluateCorrectionChannel(input.translation_correction,
                                NewtonChannelUnit::kLength),
      EvaluateCorrectionChannel(input.rotation_correction,
                                NewtonChannelUnit::kDimensionless),
      false};
  result.combined_passed =
      result.force_residual.passed && result.moment_residual.passed &&
      result.translation_correction.passed && result.rotation_correction.passed;
  return result;
}

}  // namespace fesa