#pragma once

#include "fesa/Core/Dof.hpp"
#include "fesa/Core/Domain.hpp"
#include "fesa/Core/Validation.hpp"

#include <array>
#include <map>
#include <set>
#include <utility>
#include <vector>

namespace fesa {

struct DofAddress {
  GlobalId node_id = 0;
  Dof dof = Dof::UX;
};

class DofManager {
 public:
  explicit DofManager(const Domain& domain) {
    for (const auto& [node_id, node] : domain.nodes) {
      (void)node;
      node_ids_.push_back(node_id);
      for (Dof dof : allDofs()) {
        const LocalIndex full_index = static_cast<LocalIndex>(all_dofs_.size());
        const auto key = std::make_pair(node_id, dofIndex(dof));
        all_dofs_.push_back(key);
        full_index_by_key_[key] = full_index;
      }
    }
    for (const BoundaryCondition& boundary : domain.boundary_conditions) {
      if (!validAbaqusDofRange(boundary.first_dof, boundary.last_dof)) {
        continue;
      }
      for (GlobalId node_id : resolveNodeTarget(domain, boundary.target)) {
        for (int dof = boundary.first_dof; dof <= boundary.last_dof; ++dof) {
          constrained_.insert(std::make_pair(node_id, dof - 1));
        }
      }
    }
    for (const auto& key : all_dofs_) {
      const LocalIndex full_index = full_index_by_key_.at(key);
      if (constrained_.count(key) == 0) {
        equation_by_key_[key] = static_cast<EquationId>(free_full_indices_.size());
        free_full_indices_.push_back(full_index);
      } else {
        equation_by_key_[key] = -1;
        constrained_full_indices_.push_back(full_index);
      }
    }
  }

  LocalIndex fullDofCount() const {
    return static_cast<LocalIndex>(all_dofs_.size());
  }

  LocalIndex freeDofCount() const {
    return static_cast<LocalIndex>(free_full_indices_.size());
  }

  LocalIndex constrainedDofCount() const {
    return static_cast<LocalIndex>(constrained_full_indices_.size());
  }

  const std::vector<GlobalId>& nodeIds() const {
    return node_ids_;
  }

  const std::vector<LocalIndex>& freeFullIndices() const {
    return free_full_indices_;
  }

  const std::vector<LocalIndex>& constrainedFullIndices() const {
    return constrained_full_indices_;
  }

  DofAddress fullDof(LocalIndex full_index) const {
    const auto& key = all_dofs_.at(static_cast<std::size_t>(full_index));
    return {key.first, static_cast<Dof>(key.second)};
  }

  LocalIndex fullIndex(GlobalId node_id, Dof dof) const {
    return full_index_by_key_.at(std::make_pair(node_id, dofIndex(dof)));
  }

  EquationId equation(GlobalId node_id, Dof dof) const {
    return equation_by_key_.at(std::make_pair(node_id, dofIndex(dof)));
  }

  bool isConstrained(GlobalId node_id, Dof dof) const {
    return constrained_.count(std::make_pair(node_id, dofIndex(dof))) != 0;
  }

  std::vector<Real> reduceFullVector(const std::vector<Real>& full) const {
    std::vector<Real> reduced;
    reduced.reserve(free_full_indices_.size());
    for (LocalIndex full_index : free_full_indices_) {
      reduced.push_back(full.at(static_cast<std::size_t>(full_index)));
    }
    return reduced;
  }

  std::vector<Real> reconstructFullVector(const std::vector<Real>& reduced) const {
    std::vector<Real> full(static_cast<std::size_t>(fullDofCount()), 0.0);
    for (std::size_t i = 0; i < free_full_indices_.size(); ++i) {
      full[static_cast<std::size_t>(free_full_indices_[i])] = reduced.at(i);
    }
    return full;
  }

  std::array<LocalIndex, 24> elementFullDofIndices(const Element& element) const {
    std::array<LocalIndex, 24> indices{};
    for (LocalIndex node = 0; node < 4; ++node) {
      for (Dof dof : allDofs()) {
        const LocalIndex local = 6 * node + dofIndex(dof);
        indices[static_cast<std::size_t>(local)] = fullIndex(element.node_ids[static_cast<std::size_t>(node)], dof);
      }
    }
    return indices;
  }

  std::array<EquationId, 24> elementEquationIds(const Element& element) const {
    std::array<EquationId, 24> equations{};
    for (LocalIndex node = 0; node < 4; ++node) {
      for (Dof dof : allDofs()) {
        const LocalIndex local = 6 * node + dofIndex(dof);
        equations[static_cast<std::size_t>(local)] = equation(element.node_ids[static_cast<std::size_t>(node)], dof);
      }
    }
    return equations;
  }

 private:
  std::vector<GlobalId> node_ids_;
  std::vector<std::pair<GlobalId, int>> all_dofs_;
  std::set<std::pair<GlobalId, int>> constrained_;
  std::map<std::pair<GlobalId, int>, LocalIndex> full_index_by_key_;
  std::map<std::pair<GlobalId, int>, EquationId> equation_by_key_;
  std::vector<LocalIndex> free_full_indices_;
  std::vector<LocalIndex> constrained_full_indices_;
};

}  // namespace fesa