#include "fesa/assembly/load_assembler.h"

#include <algorithm>
#include <charconv>
#include <cmath>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

#include "fesa/constraints/essential_constraints.h"

namespace fesa {
namespace {

constexpr std::size_t kDofsPerNode = 6U;
constexpr double kShellMomentProjectionTolerance = 1.0e-12;

Status LoadFailure(const std::string& code, const SourceLocation& location,
                   const std::string& keyword, const std::string& identity,
                   const std::string& message) {
  return Status::Failure(
      FailureCategory::kModel,
      {{Severity::kError, code, location, keyword, identity, message}});
}

char AsciiLower(const char value) {
  if (value >= 'A' && value <= 'Z') {
    return static_cast<char>(value + ('a' - 'A'));
  }
  return value;
}

bool EqualName(const std::string& left, const std::string& right) {
  return left.size() == right.size() &&
         std::equal(left.begin(), left.end(), right.begin(),
                    [](const char left_value, const char right_value) {
                      return AsciiLower(left_value) == AsciiLower(right_value);
                    });
}

bool TryPositiveInteger(const std::string& text, std::int64_t& value) {
  const char* const first = text.data();
  const char* const last = first + text.size();
  const auto parsed = std::from_chars(first, last, value);
  return parsed.ec == std::errc{} && parsed.ptr == last && value > 0;
}

/// @brief Checks that equation-space indices preserve stable full-DOF order.
bool IsStrictlyIncreasing(const std::vector<std::size_t>& values) {
  return std::adjacent_find(
             values.begin(), values.end(),
             [](const std::size_t left, const std::size_t right) {
               return left >= right;
             }) == values.end();
}

/// @brief Validates the full/free/constrained partition used by load assembly.
Status ValidateDofOrder(const DofManager& dofs,
                        const std::size_t expected_full_count,
                        const SourceLocation& location) {
  const std::size_t full_count = dofs.FullDofCount();
  const auto& free_dofs = dofs.FreeDofs();
  const auto& constrained_dofs = dofs.ConstrainedDofs();
  if (full_count != expected_full_count ||
      free_dofs.size() != dofs.FreeDofCount() ||
      constrained_dofs.size() != dofs.ConstrainedDofCount() ||
      dofs.PrescribedValues().Size() != constrained_dofs.size() ||
      constrained_dofs.size() > full_count ||
      free_dofs.size() != full_count - constrained_dofs.size()) {
    return LoadFailure("invalid-load-dimensions", location, "LOAD_ASSEMBLER",
                       std::to_string(full_count),
                       "Full, free, constrained, prescribed, and model "
                       "dimensions must agree.");
  }
  if (!IsStrictlyIncreasing(free_dofs) ||
      !IsStrictlyIncreasing(constrained_dofs)) {
    return LoadFailure(
        "invalid-load-order", location, "LOAD_ASSEMBLER",
        std::to_string(full_count),
        "Free and constrained DOFs must use stable increasing full-DOF order.");
  }

  std::vector<unsigned char> ownership(full_count, 0U);
  try {
    for (std::size_t equation = 0U; equation < free_dofs.size(); ++equation) {
      const std::size_t full_dof = free_dofs[equation];
      if (full_dof >= full_count || ownership[full_dof] != 0U ||
          dofs.FreeEquation(full_dof) != equation) {
        return LoadFailure(
            "invalid-load-order", location, "LOAD_ASSEMBLER",
            std::to_string(full_dof),
            "Free equation numbering must match stable full-DOF order.");
      }
      ownership[full_dof] = 1U;
    }
    for (const std::size_t full_dof : constrained_dofs) {
      if (full_dof >= full_count || ownership[full_dof] != 0U ||
          dofs.FreeEquation(full_dof).has_value()) {
        return LoadFailure(
            "invalid-load-order", location, "LOAD_ASSEMBLER",
            std::to_string(full_dof),
            "Constrained DOFs must be unique and absent from free equations.");
      }
      ownership[full_dof] = 2U;
    }
  } catch (const std::out_of_range&) {
    return LoadFailure(
        "invalid-load-dimensions", location, "LOAD_ASSEMBLER",
        std::to_string(full_count),
        "DofManager equation storage must cover every full DOF.");
  }
  if (std::find(ownership.begin(), ownership.end(), 0U) != ownership.end()) {
    return LoadFailure(
        "invalid-load-order", location, "LOAD_ASSEMBLER",
        std::to_string(full_count),
        "Free and constrained DOFs must partition the full range.");
  }
  return Status::Ok();
}

Result<std::vector<EntityIndex>> ResolveTarget(const Domain& domain,
                                               const NodalLoad& load) {
  std::vector<const NodeSet*> matching_sets;
  for (const auto& set : domain.NodeSets()) {
    if (EqualName(set.name, load.target)) {
      matching_sets.push_back(&set);
    }
  }

  std::vector<EntityIndex> matching_nodes;
  std::int64_t label = 0;
  if (TryPositiveInteger(load.target, label)) {
    for (std::size_t index = 0U; index < domain.Nodes().size(); ++index) {
      if (domain.Nodes()[index].source_id.source_label == label) {
        matching_nodes.push_back(static_cast<EntityIndex>(index));
      }
    }
  }

  if (matching_sets.size() > 1U || matching_nodes.size() > 1U ||
      (!matching_sets.empty() && !matching_nodes.empty())) {
    return Result<std::vector<EntityIndex>>::Failure(
        LoadFailure("invalid-load-target", load.location, "CLOAD", load.target,
                    "The load target must resolve unambiguously to one node or "
                    "one expanded node set."));
  }
  if (!matching_sets.empty()) {
    const auto& nodes = matching_sets.front()->node_indices;
    std::vector<unsigned char> seen(domain.Nodes().size(), 0U);
    for (const EntityIndex node : nodes) {
      if (node >= domain.Nodes().size() || seen[node] != 0U) {
        return Result<std::vector<EntityIndex>>::Failure(LoadFailure(
            "invalid-load-target", load.location, "CLOAD", load.target,
            "The expanded node set must contain unique in-range stable node "
            "identities."));
      }
      seen[node] = 1U;
    }
    return Result<std::vector<EntityIndex>>::Success(nodes);
  }
  if (!matching_nodes.empty()) {
    return Result<std::vector<EntityIndex>>::Success(std::move(matching_nodes));
  }
  return Result<std::vector<EntityIndex>>::Failure(LoadFailure(
      "invalid-load-target", load.location, "CLOAD", load.target,
      "The load target must resolve to one semantic node or node set."));
}

Status ValidateFiniteVector(const Vector& values,
                            const SourceLocation& location,
                            const std::string& identity) {
  for (std::size_t index = 0U; index < values.Size(); ++index) {
    if (!std::isfinite(values[index])) {
      return LoadFailure("nonfinite-load-value", location, "LOAD_ASSEMBLER",
                         identity + ":" + std::to_string(index),
                         "Load and prescribed displacement vectors must "
                         "contain finite values.");
    }
  }
  return Status::Ok();
}

Status ValidateShellMoments(const Domain& domain, const Vector& full_load) {
  if (domain.ShellElements().empty()) {
    return Status::Ok();
  }

  std::vector<const ShellNodeInitialFrame*> frame_by_node(domain.Nodes().size(),
                                                          nullptr);
  for (const auto& frame : domain.ShellNodeInitialFrames()) {
    if (frame.node_index >= frame_by_node.size() ||
        frame_by_node[frame.node_index] != nullptr) {
      return LoadFailure(
          "invalid-shell-director", {domain.SourcePath(), 0U}, "NODE",
          std::to_string(frame.node_index),
          "Shell nodal directors must have unique in-range node identities.");
    }
    frame_by_node[frame.node_index] = &frame;
  }

  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const double moment_x = full_load[node * kDofsPerNode + 3U];
    const double moment_y = full_load[node * kDofsPerNode + 4U];
    const double moment_z = full_load[node * kDofsPerNode + 5U];
    if (moment_x == 0.0 && moment_y == 0.0 && moment_z == 0.0) {
      continue;
    }

    const auto* const frame = frame_by_node[node];
    if (frame == nullptr) {
      return LoadFailure(
          "invalid-shell-director", domain.Nodes()[node].location, "NODE",
          domain.Nodes()[node].source_id.source_label_text,
          "A loaded shell node must have an approved initial director.");
    }

    const double moment_scale = std::max(
        std::abs(moment_x), std::max(std::abs(moment_y), std::abs(moment_z)));
    const double scaled_x = moment_x / moment_scale;
    const double scaled_y = moment_y / moment_scale;
    const double scaled_z = moment_z / moment_scale;
    const double scaled_norm = std::hypot(scaled_x, scaled_y, scaled_z);
    const double scaled_dot = frame->director[0U] * scaled_x +
                              frame->director[1U] * scaled_y +
                              frame->director[2U] * scaled_z;
    const double projection_ratio = std::abs(scaled_dot) / scaled_norm;
    if (!(projection_ratio <= kShellMomentProjectionTolerance)) {
      return LoadFailure("unsupported-drilling-load",
                         domain.Nodes()[node].location, "CLOAD",
                         domain.Nodes()[node].source_id.source_label_text,
                         "The aggregate nodal moment has an unsupported "
                         "director-parallel component.");
    }
  }
  return Status::Ok();
}

}  // namespace

Result<Vector> LoadAssembler::AssembleFullNodalLoad(const AnalysisModel& model,
                                                    const DofManager& dofs) {
  const Domain& domain = model.GetDomain();
  if (domain.Nodes().size() >
      (std::numeric_limits<std::size_t>::max)() / kDofsPerNode) {
    return Result<Vector>::Failure(LoadFailure(
        "invalid-load-dimensions", {domain.SourcePath(), 0U}, "LOAD_ASSEMBLER",
        domain.SourceContentIdentity(),
        "The semantic node count cannot be represented in full-DOF order."));
  }
  const std::size_t expected_full_count = domain.Nodes().size() * kDofsPerNode;
  const Status dof_status =
      ValidateDofOrder(dofs, expected_full_count, {domain.SourcePath(), 0U});
  if (!dof_status.IsOk()) {
    return Result<Vector>::Failure(dof_status);
  }
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    for (std::size_t component = 0U; component < kDofsPerNode; ++component) {
      try {
        if (dofs.FullDof(static_cast<EntityIndex>(node),
                         static_cast<DofComponent>(component)) !=
            node * kDofsPerNode + component) {
          return Result<Vector>::Failure(LoadFailure(
              "invalid-load-order", domain.Nodes()[node].location,
              "LOAD_ASSEMBLER",
              domain.Nodes()[node].source_id.source_label_text,
              "DofManager node/component identity must match full-DOF order."));
        }
      } catch (const std::out_of_range&) {
        return Result<Vector>::Failure(LoadFailure(
            "invalid-load-dimensions", domain.Nodes()[node].location,
            "LOAD_ASSEMBLER", domain.Nodes()[node].source_id.source_label_text,
            "DofManager must provide all six DOFs for every semantic node."));
      }
    }
  }

  const auto& active_loads = model.ActiveLoads();
  const auto& loads = model.Step().loads;
  if (active_loads.size() != loads.size()) {
    return Result<Vector>::Failure(LoadFailure(
        "invalid-load-order", model.Step().location, "CLOAD", model.Step().name,
        "The active load view must include every sole-step load once."));
  }

  Vector full_load{expected_full_count};
  // Active load indices are required to be the original source order; this
  // loop is therefore also the fixed floating-point accumulation order.
  for (std::size_t source_order = 0U; source_order < active_loads.size();
       ++source_order) {
    const EntityIndex load_index = active_loads[source_order];
    if (static_cast<std::size_t>(load_index) != source_order ||
        load_index >= loads.size()) {
      return Result<Vector>::Failure(LoadFailure(
          "invalid-load-order", model.Step().location, "CLOAD",
          std::to_string(source_order),
          "Active loads must retain complete stable source order."));
    }
    const auto& load = loads[load_index];
    if (load.dof < 1 || load.dof > static_cast<int>(kDofsPerNode)) {
      return Result<Vector>::Failure(LoadFailure(
          "invalid-load-dof", load.location, "CLOAD", load.target,
          "A nodal load component must be in the range 1 through 6."));
    }
    if (!std::isfinite(load.magnitude)) {
      return Result<Vector>::Failure(
          LoadFailure("nonfinite-load-value", load.location, "CLOAD",
                      load.target, "A nodal load magnitude must be finite."));
    }

    auto target = ResolveTarget(domain, load);
    if (!target.HasValue()) {
      return Result<Vector>::Failure(target.GetStatus());
    }
    const auto component = static_cast<DofComponent>(load.dof - 1);
    for (const EntityIndex node : target.Value()) {
      const std::size_t full_dof = dofs.FullDof(node, component);
      const double accumulated = full_load[full_dof] + load.magnitude;
      if (!std::isfinite(accumulated)) {
        return Result<Vector>::Failure(LoadFailure(
            "nonfinite-load-accumulation", load.location, "CLOAD", load.target,
            "Source-order load accumulation produced a nonfinite value."));
      }
      full_load[full_dof] = accumulated;
    }
  }
  const Status shell_moment_status = ValidateShellMoments(domain, full_load);
  if (!shell_moment_status.IsOk()) {
    return Result<Vector>::Failure(shell_moment_status);
  }
  return Result<Vector>::Success(std::move(full_load));
}

Result<Vector> LoadAssembler::EffectiveFreeRhs(const Vector& full_load,
                                               const SparseMatrix& kfc,
                                               const Vector& prescribed_values,
                                               const DofManager& dofs) {
  const SourceLocation location{{}, 0U};
  const Status dof_status = ValidateDofOrder(dofs, full_load.Size(), location);
  if (!dof_status.IsOk()) {
    return Result<Vector>::Failure(dof_status);
  }
  if (kfc.Rows() != dofs.FreeDofCount() ||
      kfc.Columns() != dofs.ConstrainedDofCount() ||
      prescribed_values.Size() != dofs.ConstrainedDofCount()) {
    return Result<Vector>::Failure(LoadFailure(
        "invalid-load-dimensions", location, "LOAD_ASSEMBLER",
        std::to_string(kfc.Rows()) + "x" + std::to_string(kfc.Columns()),
        "Kfc rows/columns and prescribed values must match free/constrained "
        "order."));
  }
  const Status matrix_status = kfc.Validate();
  if (!matrix_status.IsOk()) {
    return Result<Vector>::Failure(matrix_status);
  }
  const Status load_status =
      ValidateFiniteVector(full_load, location, "full-load");
  if (!load_status.IsOk()) {
    return Result<Vector>::Failure(load_status);
  }
  const Status prescribed_status =
      ValidateFiniteVector(prescribed_values, location, "prescribed-values");
  if (!prescribed_status.IsOk()) {
    return Result<Vector>::Failure(prescribed_status);
  }

  Vector correction{kfc.Rows()};
  for (std::size_t row = 0U; row < kfc.Rows(); ++row) {
    double sum = 0.0;
    for (std::size_t position = kfc.RowOffsets()[row];
         position < kfc.RowOffsets()[row + 1U]; ++position) {
      const double product = kfc.Values()[position] *
                             prescribed_values[kfc.ColumnIndices()[position]];
      if (!std::isfinite(product)) {
        return Result<Vector>::Failure(LoadFailure(
            "nonfinite-load-accumulation", location, "LOAD_ASSEMBLER",
            std::to_string(row),
            "Kfc times prescribed displacement produced a nonfinite product."));
      }
      sum += product;
      if (!std::isfinite(sum)) {
        return Result<Vector>::Failure(LoadFailure(
            "nonfinite-load-accumulation", location, "LOAD_ASSEMBLER",
            std::to_string(row),
            "Kfc times prescribed displacement produced a nonfinite row sum."));
      }
    }
    correction[row] = sum;
  }

  Vector rhs = EssentialConstraints::GatherFree(full_load, dofs);
  // The constrained vector is already in DofManager order, so this is the
  // approved elimination equation rhs = Ff - Kfc*dc without reordering dc.
  for (std::size_t row = 0U; row < rhs.Size(); ++row) {
    const double value = rhs[row] - correction[row];
    if (!std::isfinite(value)) {
      return Result<Vector>::Failure(
          LoadFailure("nonfinite-load-accumulation", location, "LOAD_ASSEMBLER",
                      std::to_string(row),
                      "Effective RHS subtraction produced a nonfinite value."));
    }
    rhs[row] = value;
  }
  return Result<Vector>::Success(std::move(rhs));
}

}  // namespace fesa