#include "fesa/fem/dof_manager.h"

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

#include "fesa/model/source_target_resolver.h"

namespace fesa {
namespace {

constexpr std::size_t kDofsPerNode = 6U;

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

bool SameSourceIdentity(const SourceEntityId& left,
                        const SourceEntityId& right) {
  return left.instance_name == right.instance_name &&
         left.source_label == right.source_label &&
         left.source_label_text == right.source_label_text;
}

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();
}

bool HasAdjacentDuplicate(const std::vector<std::size_t>& values) {
  return std::adjacent_find(values.begin(), values.end()) != values.end();
}

std::vector<DofComponent> FullNodeComponents() {
  return {DofComponent::kUx,  DofComponent::kUy,  DofComponent::kUz,
          DofComponent::kUrx, DofComponent::kUry, DofComponent::kUrz};
}

std::vector<EntityIndex> ExpandBoundaryTarget(
    const SourceTargetResolver& resolver, const BoundaryCondition& boundary) {
  auto resolved =
      resolver.Resolve({SourceEntityKind::kNode, "", boundary.target});
  if (!resolved.HasValue()) {
    return {};
  }
  std::vector<EntityIndex> indices;
  indices.reserve(resolved.Value().size());
  for (const auto& target : resolved.Value()) {
    indices.push_back(target.entity_index);
  }
  return indices;
}

void AppendScatter(std::vector<std::vector<std::size_t>>& columns_by_row,
                   const std::vector<std::size_t>& scatter) {
  for (const std::size_t row : scatter) {
    auto& columns = columns_by_row[row];
    columns.insert(columns.end(), scatter.begin(), scatter.end());
  }
}

/// @brief Builds sorted unique CSR columns by deterministic scatter traversal.
SparsePattern BuildSparsePattern(
    const std::size_t full_dof_count,
    const std::vector<std::vector<std::size_t>>& element_scatters) {
  std::vector<std::vector<std::size_t>> columns_by_row(full_dof_count);
  for (const auto& scatter : element_scatters) {
    AppendScatter(columns_by_row, scatter);
  }

  SparsePattern pattern;
  pattern.row_offsets.reserve(full_dof_count + 1U);
  pattern.row_offsets.push_back(0U);
  for (auto& columns : columns_by_row) {
    // Stable CSR structure is independent of element traversal duplicates.
    std::sort(columns.begin(), columns.end());
    columns.erase(std::unique(columns.begin(), columns.end()), columns.end());
    pattern.column_indices.insert(pattern.column_indices.end(), columns.begin(),
                                  columns.end());
    pattern.row_offsets.push_back(pattern.column_indices.size());
  }
  return pattern;
}

template <std::size_t kSize>
std::array<std::size_t, kSize> FixedScatter(
    const std::vector<std::size_t>& scatter) {
  if (scatter.size() != kSize) {
    throw std::out_of_range{
        "Stored element scatter does not match the compatibility shape."};
  }
  std::array<std::size_t, kSize> fixed{};
  std::copy(scatter.begin(), scatter.end(), fixed.begin());
  return fixed;
}

}  // namespace

Result<DofManager> DofManager::Create(const AnalysisModel& model) {
  const Domain& domain = model.GetDomain();
  std::vector<ElementDofLayout> layouts;
  layouts.reserve(model.ActiveElements().size());
  for (const EntityIndex element_index : model.ActiveElements()) {
    if (element_index >= domain.Elements().Size()) {
      return Result<DofManager>::Failure(
          DofFailure("invalid-element-layout-order", {domain.SourcePath(), 0U},
                     std::to_string(element_index),
                     "An active element index is outside the Domain."));
    }
    const auto& definition = domain.Elements()[element_index];
    layouts.push_back({definition.SourceId(), definition.NodeIndices(),
                       FullNodeComponents()});
  }

  DofManager dofs;
  const Status status = dofs.BuildLayouts(model, layouts);
  if (!status.IsOk()) {
    return Result<DofManager>::Failure(status);
  }
  return Result<DofManager>::Success(std::move(dofs));
}

Status DofManager::Build(const AnalysisModel& analysis_model,
                         const ElementView& elements) {
  std::vector<ElementDofLayout> layouts;
  layouts.reserve(elements.size());
  for (const auto& element : elements) {
    layouts.push_back(element.get().DofLayout());
  }
  return BuildLayouts(analysis_model, layouts);
}

Status DofManager::BuildLayouts(const AnalysisModel& analysis_model,
                                const std::vector<ElementDofLayout>& layouts) {
  const Domain& domain = analysis_model.GetDomain();
  if (domain.Nodes().size() >
      (std::numeric_limits<std::size_t>::max)() / kDofsPerNode) {
    return DofFailure(
        "invalid-dof-dimensions", {domain.SourcePath(), 0U},
        std::to_string(domain.Nodes().size()),
        "The node count cannot be represented in full-DOF storage.");
  }
  if (layouts.size() != analysis_model.ActiveElements().size()) {
    return DofFailure(
        "invalid-element-layout-order", {domain.SourcePath(), 0U},
        std::to_string(layouts.size()),
        "Runtime elements must match the active element inventory.");
  }

  for (std::size_t source_order = 0U; source_order < layouts.size();
       ++source_order) {
    const EntityIndex definition_index =
        analysis_model.ActiveElements()[source_order];
    if (definition_index >= domain.Elements().Size() ||
        !SameSourceIdentity(layouts[source_order].source_id,
                            domain.Elements()[definition_index].SourceId())) {
      return DofFailure(
          "invalid-element-layout-order", {domain.SourcePath(), 0U},
          std::to_string(source_order),
          "Runtime element layouts must preserve active source order and "
          "identity.");
    }
  }

  const SourceTargetIndex target_index = SourceTargetIndex::FromDomain(domain);
  const SourceTargetResolver target_resolver{target_index};
  const std::size_t full_count = domain.Nodes().size() * kDofsPerNode;

  std::vector<std::optional<double>> prescribed_by_full_dof(full_count);
  for (const EntityIndex boundary_index :
       analysis_model.ActiveBoundaryConditions()) {
    const auto& boundary = analysis_model.Step().boundaries.at(boundary_index);
    const auto target = ExpandBoundaryTarget(target_resolver, boundary);
    for (const EntityIndex node : target) {
      for (int component = boundary.first_dof; component <= boundary.last_dof;
           ++component) {
        const std::size_t full_dof =
            static_cast<std::size_t>(node) * kDofsPerNode +
            static_cast<std::size_t>(component - 1);
        auto& prescribed = prescribed_by_full_dof[full_dof];
        if (prescribed && *prescribed != boundary.value) {
          return Status::Failure(
              FailureCategory::kInput,
              {{Severity::kError, "conflicting-boundary-condition",
                boundary.location, "BOUNDARY", boundary.target,
                "Expanded boundary rows prescribe different values to one "
                "node/DOF."}});
        }
        prescribed = boundary.value;
      }
    }
  }

  std::vector<std::size_t> free_dofs;
  std::vector<std::size_t> constrained_dofs;
  std::vector<double> constrained_values;
  std::vector<std::optional<std::size_t>> free_equations(full_count);
  free_dofs.reserve(full_count);
  constrained_dofs.reserve(full_count);
  constrained_values.reserve(full_count);
  // A full-DOF scan fixes free equations, constrained DOFs, and dc in the
  // same stable order regardless of boundary declaration overlap.
  for (std::size_t full_dof = 0U; full_dof < full_count; ++full_dof) {
    if (prescribed_by_full_dof[full_dof]) {
      constrained_dofs.push_back(full_dof);
      constrained_values.push_back(*prescribed_by_full_dof[full_dof]);
    } else {
      free_equations[full_dof] = free_dofs.size();
      free_dofs.push_back(full_dof);
    }
  }

  Vector prescribed_values{constrained_values.size()};
  for (std::size_t index = 0U; index < constrained_values.size(); ++index) {
    prescribed_values[index] = constrained_values[index];
  }

  DofManager candidate{full_count,
                       std::move(free_equations),
                       {},
                       std::move(free_dofs),
                       std::move(constrained_dofs),
                       std::move(prescribed_values),
                       {}};
  candidate.element_scatters_.reserve(layouts.size());
  for (const auto& layout : layouts) {
    auto scatter = candidate.ElementScatter(layout);
    if (!scatter.HasValue()) {
      return scatter.GetStatus();
    }
    candidate.element_scatters_.push_back(std::move(scatter.Value()));
  }
  candidate.sparse_pattern_ =
      BuildSparsePattern(full_count, candidate.element_scatters_);

  const Status invariant_status = candidate.ValidateInvariants();
  if (!invariant_status.IsOk()) {
    return invariant_status;
  }
  *this = std::move(candidate);
  return Status::Ok();
}

std::size_t DofManager::FullDofCount() const noexcept {
  return full_dof_count_;
}

std::size_t DofManager::FreeDofCount() const noexcept {
  return free_dofs_.size();
}

std::size_t DofManager::ConstrainedDofCount() const noexcept {
  return constrained_dofs_.size();
}

std::size_t DofManager::FullDof(const EntityIndex node,
                                const DofComponent component) const {
  const std::size_t component_index = static_cast<std::size_t>(component);
  if (node >= full_dof_count_ / kDofsPerNode ||
      component_index >= kDofsPerNode) {
    throw std::out_of_range{"Node or DOF component is out of range."};
  }
  return static_cast<std::size_t>(node) * kDofsPerNode + component_index;
}

std::optional<std::size_t> DofManager::FreeEquation(
    const std::size_t full_dof) const {
  return free_equations_.at(full_dof);
}

Result<std::vector<std::size_t>> DofManager::ElementScatter(
    const ElementDofLayout& layout) const {
  if (layout.node_indices.empty() || layout.components_per_node.empty() ||
      layout.node_indices.size() > (std::numeric_limits<std::size_t>::max)() /
                                       layout.components_per_node.size()) {
    return Result<std::vector<std::size_t>>::Failure(DofFailure(
        "invalid-element-dof-layout", {}, layout.source_id.source_label_text,
        "An element DOF layout requires a representable nonempty topology "
        "and component inventory."));
  }

  std::vector<std::size_t> scatter;
  scatter.reserve(layout.node_indices.size() *
                  layout.components_per_node.size());
  for (const EntityIndex node : layout.node_indices) {
    for (const DofComponent component : layout.components_per_node) {
      const std::size_t component_index = static_cast<std::size_t>(component);
      if (node >= full_dof_count_ / kDofsPerNode ||
          component_index >= kDofsPerNode) {
        return Result<std::vector<std::size_t>>::Failure(DofFailure(
            "invalid-element-dof-layout", {},
            layout.source_id.source_label_text,
            "Element node and component identities must resolve in the full "
            "DOF range."));
      }
      const std::size_t full_dof =
          static_cast<std::size_t>(node) * kDofsPerNode + component_index;
      if (std::find(scatter.begin(), scatter.end(), full_dof) !=
          scatter.end()) {
        return Result<std::vector<std::size_t>>::Failure(DofFailure(
            "duplicate-element-dof", {}, layout.source_id.source_label_text,
            "An element layout must not repeat a full DOF."));
      }
      scatter.push_back(full_dof);
    }
  }
  return Result<std::vector<std::size_t>>::Success(std::move(scatter));
}

std::array<std::size_t, 12> DofManager::ElementScatter(
    const EntityIndex element) const {
  return FixedScatter<12U>(element_scatters_.at(element));
}

std::array<std::size_t, 24> DofManager::ShellElementScatter(
    const EntityIndex element) const {
  return FixedScatter<24U>(element_scatters_.at(element));
}

const std::vector<std::size_t>& DofManager::FreeDofs() const noexcept {
  return free_dofs_;
}

const std::vector<std::size_t>& DofManager::ConstrainedDofs() const noexcept {
  return constrained_dofs_;
}

const Vector& DofManager::PrescribedValues() const noexcept {
  return prescribed_values_;
}

const SparsePattern& DofManager::GetSparsePattern() const noexcept {
  return sparse_pattern_;
}

Status DofManager::ValidateInvariants() const {
  if (free_equations_.size() != full_dof_count_ ||
      free_dofs_.size() > full_dof_count_ ||
      constrained_dofs_.size() > full_dof_count_ ||
      free_dofs_.size() + constrained_dofs_.size() != full_dof_count_ ||
      prescribed_values_.Size() != constrained_dofs_.size()) {
    return DofFailure(
        "invalid-dof-dimensions", {}, std::to_string(full_dof_count_),
        "Full, free, constrained, prescribed, and equation dimensions must "
        "agree.");
  }
  if (HasAdjacentDuplicate(free_dofs_)) {
    return DofFailure("duplicate-dof-mapping", {},
                      std::to_string(full_dof_count_),
                      "Free DOF ownership must be unique.");
  }
  if (!IsStrictlyIncreasing(free_dofs_)) {
    return DofFailure("invalid-free-dof-mapping", {},
                      std::to_string(full_dof_count_),
                      "Free DOFs must use stable increasing full-DOF order.");
  }
  if (HasAdjacentDuplicate(constrained_dofs_)) {
    return DofFailure("duplicate-dof-mapping", {},
                      std::to_string(full_dof_count_),
                      "Constrained DOF ownership must be unique.");
  }
  if (!IsStrictlyIncreasing(constrained_dofs_)) {
    return DofFailure(
        "invalid-constrained-dof-mapping", {}, std::to_string(full_dof_count_),
        "Constrained DOFs must use stable increasing full-DOF order.");
  }

  std::vector<unsigned char> ownership(full_dof_count_, 0U);
  for (std::size_t equation = 0U; equation < free_dofs_.size(); ++equation) {
    const std::size_t full_dof = free_dofs_[equation];
    if (full_dof >= full_dof_count_) {
      return DofFailure("invalid-free-dof-mapping", {},
                        std::to_string(full_dof),
                        "A free DOF is outside the full range.");
    }
    if (ownership[full_dof] != 0U) {
      return DofFailure("duplicate-dof-mapping", {}, std::to_string(full_dof),
                        "Each full DOF must have one owner.");
    }
    if (free_equations_[full_dof] != equation) {
      return DofFailure(
          "invalid-equation-mapping", {}, std::to_string(full_dof),
          "Free equation numbering must match stable free-DOF order.");
    }
    ownership[full_dof] = 1U;
  }
  for (const std::size_t full_dof : constrained_dofs_) {
    if (full_dof >= full_dof_count_) {
      return DofFailure("invalid-constrained-dof-mapping", {},
                        std::to_string(full_dof),
                        "A constrained DOF is outside the full range.");
    }
    if (ownership[full_dof] != 0U) {
      return DofFailure("duplicate-dof-mapping", {}, std::to_string(full_dof),
                        "Each full DOF must have one owner.");
    }
    if (free_equations_[full_dof].has_value()) {
      return DofFailure("invalid-equation-mapping", {},
                        std::to_string(full_dof),
                        "Constrained DOFs must be absent from free equations.");
    }
    ownership[full_dof] = 2U;
  }
  if (std::find(ownership.begin(), ownership.end(), 0U) != ownership.end()) {
    return DofFailure(
        "invalid-dof-partition", {}, std::to_string(full_dof_count_),
        "Free and constrained DOFs must partition the complete full range.");
  }

  for (const auto& scatter : element_scatters_) {
    if (scatter.empty()) {
      return DofFailure("invalid-element-dof-layout", {}, {},
                        "Stored element scatters must be nonempty.");
    }
    for (std::size_t position = 0U; position < scatter.size(); ++position) {
      const std::size_t full_dof = scatter[position];
      if (full_dof >= full_dof_count_) {
        return DofFailure("invalid-element-dof-layout", {},
                          std::to_string(full_dof),
                          "Stored element scatters must stay in range.");
      }
      if (std::find(scatter.begin(),
                    scatter.begin() + static_cast<std::ptrdiff_t>(position),
                    full_dof) !=
          scatter.begin() + static_cast<std::ptrdiff_t>(position)) {
        return DofFailure("duplicate-element-dof", {}, std::to_string(full_dof),
                          "Stored element scatters must remain unique.");
      }
    }
  }

  const SparsePattern expected_pattern =
      BuildSparsePattern(full_dof_count_, element_scatters_);
  if (sparse_pattern_.row_offsets != expected_pattern.row_offsets ||
      sparse_pattern_.column_indices != expected_pattern.column_indices) {
    return DofFailure(
        "invalid-dof-sparse-pattern", {}, std::to_string(full_dof_count_),
        "The CSR pattern must exactly match the stored element scatters.");
  }
  return Status::Ok();
}

DofManager::DofManager(const std::size_t full_dof_count,
                       std::vector<std::optional<std::size_t>> free_equations,
                       std::vector<std::vector<std::size_t>> element_scatters,
                       std::vector<std::size_t> free_dofs,
                       std::vector<std::size_t> constrained_dofs,
                       Vector prescribed_values, SparsePattern sparse_pattern)
    : full_dof_count_{full_dof_count},
      free_equations_{std::move(free_equations)},
      element_scatters_{std::move(element_scatters)},
      free_dofs_{std::move(free_dofs)},
      constrained_dofs_{std::move(constrained_dofs)},
      prescribed_values_{std::move(prescribed_values)},
      sparse_pattern_{std::move(sparse_pattern)} {}

}  // namespace fesa