#include <fesa/assembly/assembler.hpp>

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <limits>
#include <numeric>
#include <optional>
#include <stdexcept>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include <fesa/assembly/contribution.hpp>
#include <fesa/elements/beam/beam3d2.hpp>

#include <oneapi/tbb/blocked_range.h>
#include <oneapi/tbb/parallel_for.h>
#include <oneapi/tbb/task_arena.h>

namespace fesa {
namespace {

struct ElementEvaluation final {
    std::vector<MatrixContribution> contributions;
    std::optional<std::string> error;
};

auto origin_key(const EntityOrigin& origin) {
    return std::tie(
        origin.instance_name,
        origin.local_label,
        origin.part_name);
}

std::string kernel_error_message(
    const BeamElement& element,
    const BeamKernelResult& result) {
    std::string message =
        "Beam element " + std::to_string(element.origin.local_label) +
        " kernel failed";
    for (const Diagnostic& diagnostic : result.diagnostics) {
        message += ": " + diagnostic.code + " - " + diagnostic.message;
    }
    return message;
}

std::vector<std::size_t> canonical_element_order(const Domain& domain) {
    std::vector<std::size_t> order(domain.beam_elements().size());
    std::iota(order.begin(), order.end(), std::size_t{0});
    std::ranges::sort(
        order,
        [&domain](const std::size_t left, const std::size_t right) {
            return origin_key(domain.beam_elements()[left].origin) <
                   origin_key(domain.beam_elements()[right].origin);
        });
    return order;
}

std::vector<ElementId> canonical_contribution_element_ids(
    const std::vector<std::size_t>& order) {
    // Domain ElementIds are not ordered by input identity. These tie-break
    // IDs encode the existing serial assembler's element-origin order.
    std::vector<ElementId> ids(order.size(), ElementId{0});
    for (std::size_t rank = 0; rank < order.size(); ++rank) {
        ids[order[rank]] = ElementId{static_cast<std::int64_t>(rank)};
    }
    return ids;
}

ElementEvaluation evaluate_element(
    const Domain& domain,
    const DofManager& dofs,
    const BeamElement& element,
    const ElementId canonical_id) {
    const BeamKernelResult result = compute_beam3d2({
        {
            domain.node(element.nodes[0]).position,
            domain.node(element.nodes[1]).position,
        },
        element.nodes,
        domain.material(element.material),
        domain.section(element.section),
    });
    if (!result.contribution.has_value()) {
        return {{}, kernel_error_message(element, result)};
    }

    ElementEvaluation evaluation;
    evaluation.contributions.reserve(78);
    const std::array<std::size_t, 12> full_dofs =
        dofs.element_full_dofs(element);
    for (std::size_t local_row = 0; local_row < full_dofs.size();
         ++local_row) {
        for (std::size_t local_column = local_row;
             local_column < full_dofs.size();
             ++local_column) {
            evaluation.contributions.push_back({
                std::min(
                    full_dofs[local_row],
                    full_dofs[local_column]),
                std::max(
                    full_dofs[local_row],
                    full_dofs[local_column]),
                canonical_id,
                static_cast<std::uint16_t>(
                    local_row * full_dofs.size() + local_column),
                result.contribution
                    ->global_stiffness[local_row][local_column],
            });
        }
    }
    return evaluation;
}

std::vector<double> assemble_force(
    const Domain& domain,
    const DofManager& dofs) {
    std::vector<double> force(dofs.full_dof_count(), 0.0);
    for (const NodalLoad& load : domain.step().nodal_loads) {
        for (std::size_t component = 0; component < load.values.size();
             ++component) {
            const auto dof = static_cast<NodeDof>(component);
            force[dofs.full_dof({load.node, dof})] +=
                load.values[component];
        }
    }
    return force;
}

void validate_options(const AssemblyOptions options) {
    if (options.max_threads == 0) {
        throw std::invalid_argument{
            "Assembly max_threads must be greater than zero."};
    }
    if (options.max_threads >
        static_cast<std::size_t>(std::numeric_limits<int>::max())) {
        throw std::invalid_argument{
            "Assembly max_threads exceeds the TBB task arena range."};
    }
    if (options.grain_size == 0) {
        throw std::invalid_argument{
            "Assembly grain_size must be greater than zero."};
    }
}

}  // namespace

EquationSystem assemble_parallel(
    const Domain& domain,
    const DofManager& dofs,
    const AssemblyOptions options) {
    validate_options(options);

    const std::vector<std::size_t> element_order =
        canonical_element_order(domain);
    const std::vector<ElementId> element_ids =
        canonical_contribution_element_ids(element_order);
    std::vector<ElementEvaluation> evaluations(
        domain.beam_elements().size());

    oneapi::tbb::task_arena arena{
        static_cast<int>(options.max_threads)};
    arena.execute([&] {
        oneapi::tbb::parallel_for(
            oneapi::tbb::blocked_range<std::size_t>{
                0,
                evaluations.size(),
                options.grain_size,
            },
            [&](const oneapi::tbb::blocked_range<std::size_t>& range) {
                for (std::size_t index = range.begin();
                     index != range.end();
                     ++index) {
                    ElementEvaluation local = evaluate_element(
                        domain,
                        dofs,
                        domain.beam_elements()[index],
                        element_ids[index]);
                    evaluations[index] = std::move(local);
                }
            });
    });

    std::vector<MatrixContribution> contributions;
    contributions.reserve(domain.beam_elements().size() * 78);
    for (const std::size_t element_index : element_order) {
        ElementEvaluation& evaluation = evaluations[element_index];
        if (evaluation.error.has_value()) {
            throw std::runtime_error{std::move(*evaluation.error)};
        }
        contributions.insert(
            contributions.end(),
            std::make_move_iterator(evaluation.contributions.begin()),
            std::make_move_iterator(evaluation.contributions.end()));
    }

    const std::vector<MatrixContribution> canonical =
        canonicalize_contributions(contributions);
    return {
        merge_contributions(dofs.full_dof_count(), canonical),
        assemble_force(domain, dofs),
    };
}

}  // namespace fesa
