#include <fesa/fem/dof_manager.hpp>

#include <algorithm>
#include <stdexcept>
#include <vector>

namespace fesa {
namespace {

constexpr std::size_t dofs_per_node = 6;

}  // namespace

DofManager DofManager::build(const Domain& domain) {
    DofManager manager;

    std::vector<NodeId> node_ids;
    node_ids.reserve(domain.nodes().size());
    for (const Node& node : domain.nodes()) {
        node_ids.push_back(node.id);
    }
    std::ranges::sort(node_ids);

    for (std::size_t index = 0; index < node_ids.size(); ++index) {
        manager.node_full_dof_bases_.emplace(
            node_ids[index].value(), index * dofs_per_node);
    }

    const std::size_t full_dof_count = node_ids.size() * dofs_per_node;
    manager.equations_.resize(full_dof_count);
    manager.prescribed_values_.resize(full_dof_count);
    std::vector<bool> constrained(full_dof_count, false);

    for (const PrescribedDof& prescribed : domain.step().prescribed_dofs) {
        const NodeDof dof =
            static_cast<NodeDof>(prescribed.dof - std::uint8_t{1});
        const std::size_t full =
            manager.full_dof({prescribed.node, dof});
        constrained[full] = true;
        manager.prescribed_values_[full] = prescribed.value;
    }

    for (std::size_t full = 0; full < full_dof_count; ++full) {
        if (!constrained[full]) {
            manager.equations_[full] = manager.free_equation_count_;
            ++manager.free_equation_count_;
        }
    }

    return manager;
}

std::size_t DofManager::full_dof_count() const noexcept {
    return equations_.size();
}

std::size_t DofManager::free_equation_count() const noexcept {
    return free_equation_count_;
}

std::optional<std::size_t> DofManager::equation(
    const DofAddress address) const {
    return equation(full_dof(address));
}

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

std::optional<double> DofManager::prescribed_value(
    const DofAddress address) const {
    return prescribed_value(full_dof(address));
}

std::optional<double> DofManager::prescribed_value(
    const std::size_t full_dof) const {
    if (equations_.at(full_dof).has_value()) {
        return std::nullopt;
    }
    return prescribed_values_.at(full_dof);
}

std::array<std::size_t, 12> DofManager::element_full_dofs(
    const BeamElement& element) const {
    std::array<std::size_t, 12> full_dofs{};
    for (std::size_t node = 0; node < element.nodes.size(); ++node) {
        const std::size_t base = node_full_dof_base(element.nodes[node]);
        for (std::size_t component = 0; component < dofs_per_node;
             ++component) {
            full_dofs[node * dofs_per_node + component] = base + component;
        }
    }
    return full_dofs;
}

std::vector<double> DofManager::reconstruct_full(
    const std::span<const double> reduced) const {
    if (reduced.size() != free_equation_count_) {
        throw std::invalid_argument{
            "Reduced vector size must equal the free equation count."};
    }

    std::vector<double> full = prescribed_values_;
    for (std::size_t index = 0; index < equations_.size(); ++index) {
        if (equations_[index].has_value()) {
            full[index] = reduced[*equations_[index]];
        }
    }
    return full;
}

std::size_t DofManager::full_dof(const DofAddress address) const {
    const auto component = static_cast<std::uint8_t>(address.dof);
    if (component >= dofs_per_node) {
        throw std::invalid_argument{"Node DOF must be one of ux through rz."};
    }
    return node_full_dof_base(address.node) + component;
}

std::size_t DofManager::node_full_dof_base(const NodeId node) const {
    const auto found = node_full_dof_bases_.find(node.value());
    if (found == node_full_dof_bases_.end()) {
        throw std::out_of_range{"Node ID is not present in the DofManager."};
    }
    return found->second;
}

}  // namespace fesa