#include <fesa/assembly/contribution.hpp>

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

namespace fesa {
namespace {

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

bool contribution_less(
    const MatrixContribution& left,
    const MatrixContribution& right) {
    return std::tuple{
               left.row,
               left.column,
               left.element,
               left.local_order} <
           std::tuple{
               right.row,
               right.column,
               right.element,
               right.local_order};
}

}  // namespace

std::vector<MatrixContribution> canonicalize_contributions(
    const std::span<const MatrixContribution> contributions) {
    std::vector<MatrixContribution> canonical{
        contributions.begin(), contributions.end()};
    std::ranges::sort(canonical, contribution_less);
    return canonical;
}

SymmetricCsr merge_contributions(
    const std::size_t order,
    const std::span<const MatrixContribution> canonical) {
    if (!std::ranges::is_sorted(canonical, contribution_less)) {
        throw std::invalid_argument{
            "Matrix contributions are not in canonical order."};
    }

    using Coordinate = std::pair<std::size_t, std::size_t>;
    std::vector<Coordinate> coordinates;
    coordinates.reserve(order + canonical.size());
    for (std::size_t row = 0; row < order; ++row) {
        coordinates.emplace_back(row, row);
    }
    for (const MatrixContribution& contribution : canonical) {
        if (contribution.row > contribution.column ||
            contribution.column >= order) {
            throw std::invalid_argument{
                "Matrix contribution is outside the upper triangle."};
        }
        coordinates.emplace_back(
            contribution.row, contribution.column);
    }

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

    SymmetricCsr matrix{
        order,
        std::vector<std::int32_t>(order + 1, 0),
        {},
        {},
    };
    matrix.column_indices.reserve(coordinates.size());
    for (const auto [row, column] : coordinates) {
        ++matrix.row_offsets[row + 1];
        matrix.column_indices.push_back(csr_index(column));
    }
    for (std::size_t row = 0; row < order; ++row) {
        const std::size_t offset =
            static_cast<std::size_t>(matrix.row_offsets[row]) +
            static_cast<std::size_t>(matrix.row_offsets[row + 1]);
        matrix.row_offsets[row + 1] = csr_index(offset);
    }
    matrix.values.resize(matrix.column_indices.size(), 0.0);

    std::size_t contribution_index = 0;
    while (contribution_index < canonical.size()) {
        const MatrixContribution& first =
            canonical[contribution_index];
        double value = 0.0;
        do {
            value += canonical[contribution_index].value;
            ++contribution_index;
        } while (
            contribution_index < canonical.size() &&
            canonical[contribution_index].row == first.row &&
            canonical[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;
    }

    return matrix;
}

}  // namespace fesa
