#include <fesa/fem/dof_manager.hpp>

#include <algorithm>
#include <cstddef>
#include <stdexcept>

namespace fesa::fem {

void DofManager::define_node_dofs(core::NodeId node_id, std::vector<model::DofComponent> components)
{
    for (const auto component : components) {
        DofKey key{node_id, component};
        if (find_record(key) == records_.end()) {
            records_.push_back({key});
        }
    }
}

void DofManager::apply_boundary_condition(const model::BoundaryCondition& bc)
{
    auto record = find_record({bc.node_id(), bc.component()});
    if (record == records_.end()) {
        throw std::invalid_argument("boundary condition references undefined dof");
    }
    record->constrained = true;
}

void DofManager::number_equations()
{
    std::sort(records_.begin(), records_.end(), [](const Record& lhs, const Record& rhs) {
        if (lhs.key.node_id.value != rhs.key.node_id.value) {
            return lhs.key.node_id.value < rhs.key.node_id.value;
        }
        return static_cast<int>(lhs.key.component) < static_cast<int>(rhs.key.component);
    });

    int free_id = 0;
    for (int equation_id = 0; equation_id < static_cast<int>(records_.size()); ++equation_id) {
        records_[equation_id].equation_id = equation_id;
        if (records_[equation_id].constrained) {
            records_[equation_id].free_equation_id = std::nullopt;
        } else {
            records_[equation_id].free_equation_id = free_id++;
        }
    }

    sparse_pattern_.clear();
    for (int row = 0; row < free_id; ++row) {
        for (int column = 0; column < free_id; ++column) {
            sparse_pattern_.push_back({row, column});
        }
    }
}

int DofManager::total_dof_count() const
{
    return static_cast<int>(records_.size());
}

int DofManager::free_dof_count() const
{
    return static_cast<int>(
        std::count_if(records_.begin(), records_.end(), [](const Record& record) {
            return !record.constrained;
        })
    );
}

int DofManager::constrained_dof_count() const
{
    return total_dof_count() - free_dof_count();
}

bool DofManager::is_constrained(DofKey key) const
{
    return require_record(key).constrained;
}

int DofManager::equation_id(DofKey key) const
{
    return require_record(key).equation_id;
}

std::optional<int> DofManager::free_equation_id(DofKey key) const
{
    return require_record(key).free_equation_id;
}

std::vector<double> DofManager::expand_free_vector(const std::vector<double>& free_values) const
{
    if (free_values.size() != static_cast<std::size_t>(free_dof_count())) {
        throw std::invalid_argument("free vector size does not match dof manager");
    }

    std::vector<double> full(records_.size(), 0.0);
    for (const auto& record : records_) {
        if (record.free_equation_id.has_value()) {
            full[static_cast<std::size_t>(record.equation_id)] =
                free_values[static_cast<std::size_t>(*record.free_equation_id)];
        }
    }
    return full;
}

const std::vector<std::pair<int, int>>& DofManager::sparse_pattern() const
{
    return sparse_pattern_;
}

std::vector<DofManager::Record>::iterator DofManager::find_record(DofKey key)
{
    return std::find_if(records_.begin(), records_.end(), [key](const Record& record) {
        return record.key == key;
    });
}

std::vector<DofManager::Record>::const_iterator DofManager::find_record(DofKey key) const
{
    return std::find_if(records_.begin(), records_.end(), [key](const Record& record) {
        return record.key == key;
    });
}

const DofManager::Record& DofManager::require_record(DofKey key) const
{
    const auto record = find_record(key);
    if (record == records_.end()) {
        throw std::invalid_argument("dof is not defined");
    }
    return *record;
}

} // namespace fesa::fem