#pragma once

#include <fesa/fem/dof_key.hpp>

#include <algorithm>
#include <optional>
#include <stdexcept>
#include <utility>
#include <vector>

namespace fesa::fem {

class DofManager {
public:
    void 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 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 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 total_dof_count() const
    {
        return static_cast<int>(records_.size());
    }

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

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

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

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

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

    std::vector<double> 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>>& sparse_pattern() const
    {
        return sparse_pattern_;
    }

private:
    struct Record {
        DofKey key;
        bool constrained = false;
        int equation_id = -1;
        std::optional<int> free_equation_id;
    };

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

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

    const Record& 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;
    }

    std::vector<Record> records_;
    std::vector<std::pair<int, int>> sparse_pattern_;
};

} // namespace fesa::fem