#include <fesa/assembly/serial_assembler.hpp>

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

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

namespace fesa {
namespace {

struct NumericContribution final {
    std::size_t row;
    std::size_t column;
    const EntityOrigin* element_origin;
    std::size_t local_order;
    double value;
};

std::int32_t csr_index(const std::size_t value) {
    if (value >
        static_cast<std::size_t>(
            std::numeric_limits<std::int32_t>::max())) {
        throw std::overflow_error{
            "Symmetric CSR exceeds the 32-bit index range."};
    }
    return static_cast<std::int32_t>(value);
}

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

SymmetricCsr build_sparsity_pattern(
    const Domain& domain,
    const DofManager& dofs) {
    using Coordinate = std::pair<std::size_t, std::size_t>;
    std::vector<Coordinate> coordinates;
    coordinates.reserve(
        dofs.full_dof_count() +
        domain.beam_elements().size() * 78);

    for (std::size_t full = 0; full < dofs.full_dof_count(); ++full) {
        coordinates.emplace_back(full, full);
    }

    for (const BeamElement& element : domain.beam_elements()) {
        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) {
                coordinates.emplace_back(
                    std::min(
                        full_dofs[local_row],
                        full_dofs[local_column]),
                    std::max(
                        full_dofs[local_row],
                        full_dofs[local_column]));
            }
        }
    }

    std::ranges::sort(coordinates);
    coordinates.erase(
        std::ranges::unique(coordinates).begin(),
        coordinates.end());

    SymmetricCsr pattern{
        dofs.full_dof_count(),
        std::vector<std::int32_t>(dofs.full_dof_count() + 1, 0),
        {},
        {},
    };
    pattern.column_indices.reserve(coordinates.size());
    for (const auto [row, column] : coordinates) {
        if (row >= pattern.order || column >= pattern.order) {
            throw std::invalid_argument{
                "DofManager element mapping exceeds the full system order."};
        }
        ++pattern.row_offsets[row + 1];
        pattern.column_indices.push_back(csr_index(column));
    }
    for (std::size_t row = 0; row < pattern.order; ++row) {
        const std::size_t offset =
            static_cast<std::size_t>(pattern.row_offsets[row]) +
            static_cast<std::size_t>(pattern.row_offsets[row + 1]);
        pattern.row_offsets[row + 1] = csr_index(offset);
    }
    pattern.values.resize(pattern.column_indices.size(), 0.0);
    return pattern;
}

std::runtime_error kernel_error(
    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 std::runtime_error{std::move(message)};
}

std::vector<NumericContribution> collect_numeric_contributions(
    const Domain& domain,
    const DofManager& dofs) {
    std::vector<NumericContribution> contributions;
    contributions.reserve(domain.beam_elements().size() * 78);

    for (const BeamElement& element : domain.beam_elements()) {
        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()) {
            throw kernel_error(element, result);
        }

        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) {
                contributions.push_back({
                    std::min(
                        full_dofs[local_row],
                        full_dofs[local_column]),
                    std::max(
                        full_dofs[local_row],
                        full_dofs[local_column]),
                    &element.origin,
                    local_row * full_dofs.size() + local_column,
                    result.contribution
                        ->global_stiffness[local_row][local_column],
                });
            }
        }
    }

    std::ranges::sort(
        contributions,
        [](const NumericContribution& left,
           const NumericContribution& right) {
            return std::tuple{
                       left.row,
                       left.column,
                       origin_key(*left.element_origin),
                       left.local_order} <
                   std::tuple{
                       right.row,
                       right.column,
                       origin_key(*right.element_origin),
                       right.local_order};
        });
    return contributions;
}

void merge_numeric_contributions(
    SymmetricCsr& matrix,
    const std::vector<NumericContribution>& contributions) {
    std::size_t contribution_index = 0;
    while (contribution_index < contributions.size()) {
        const NumericContribution& first =
            contributions[contribution_index];
        double value = 0.0;
        do {
            value += contributions[contribution_index].value;
            ++contribution_index;
        } while (
            contribution_index < contributions.size() &&
            contributions[contribution_index].row == first.row &&
            contributions[contribution_index].column == first.column);

        const auto row_begin =
            matrix.column_indices.begin() +
            matrix.row_offsets[first.row];
        const auto row_end =
            matrix.column_indices.begin() +
            matrix.row_offsets[first.row + 1];
        const auto entry = std::lower_bound(
            row_begin,
            row_end,
            csr_index(first.column));
        if (entry == row_end || *entry != csr_index(first.column)) {
            throw std::logic_error{
                "Numeric contribution is absent from the CSR pattern."};
        }
        matrix.values[static_cast<std::size_t>(
            std::distance(matrix.column_indices.begin(), entry))] = value;
    }
}

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;
}

}  // namespace

EquationSystem assemble_serial(
    const Domain& domain,
    const DofManager& dofs) {
    SymmetricCsr stiffness = build_sparsity_pattern(domain, dofs);
    const std::vector<NumericContribution> contributions =
        collect_numeric_contributions(domain, dofs);
    merge_numeric_contributions(stiffness, contributions);
    return {
        std::move(stiffness),
        assemble_force(domain, dofs),
    };
}

}  // namespace fesa
