#include "fesa/math/sparse_matrix.hpp"

#include "fesa/fem/dof_manager.hpp"

#include <algorithm>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <string>
#include <tuple>
#include <utility>

namespace fesa {
namespace {

Status sparseFailure(
    const std::string& code,
    const std::string& identity,
    const std::string& message) {
    return Status::failure(
        FailureCategory::model,
        {{Severity::error,
          code,
          {{}, 0U},
          "SPARSE_MATRIX",
          identity,
          message}});
}

Status validateCsr(
    const std::size_t rows,
    const std::size_t columns,
    const std::vector<std::size_t>& rowOffsets,
    const std::vector<std::size_t>& columnIndices,
    const std::vector<double>* const values) {
    if (rows == (std::numeric_limits<std::size_t>::max)() ||
        rowOffsets.size() != rows + 1U) {
        return sparseFailure(
            "invalid-sparse-shape",
            "row-offset-count",
            "CSR row offsets must contain exactly rows plus one entries.");
    }
    if (rowOffsets.empty() || rowOffsets.front() != 0U ||
        rowOffsets.back() != columnIndices.size()) {
        return sparseFailure(
            "invalid-sparse-pattern",
            "row-offset-range",
            "CSR row offsets must start at zero and end at the column count.");
    }
    if (values != nullptr && values->size() != columnIndices.size()) {
        return sparseFailure(
            "invalid-sparse-shape",
            "value-count",
            "CSR column and value arrays must have equal sizes.");
    }

    for (std::size_t row = 0U; row < rows; ++row) {
        const std::size_t begin = rowOffsets[row];
        const std::size_t end = rowOffsets[row + 1U];
        if (begin > end || end > columnIndices.size()) {
            return sparseFailure(
                "invalid-sparse-pattern",
                std::to_string(row),
                "CSR row offsets must be nondecreasing and remain in range.");
        }
        for (std::size_t position = begin; position < end; ++position) {
            if (columnIndices[position] >= columns) {
                return sparseFailure(
                    "invalid-sparse-index",
                    std::to_string(position),
                    "CSR column index is outside the matrix dimensions.");
            }
            if (position > begin &&
                columnIndices[position - 1U] >= columnIndices[position]) {
                return sparseFailure(
                    "invalid-sparse-pattern",
                    std::to_string(row),
                    "CSR columns must be sorted and unique within each row.");
            }
            if (values != nullptr && !std::isfinite((*values)[position])) {
                return sparseFailure(
                    "nonfinite-sparse-value",
                    std::to_string(position),
                    "CSR values must be finite.");
            }
        }
    }
    return Status::ok();
}

} // namespace

Result<SparseMatrix> SparseMatrix::fromCoo(
    const std::size_t rows,
    const std::size_t columns,
    std::vector<CooContribution> contributions,
    const SparsePattern& expectedPattern) {
    const Status patternStatus = validateCsr(
        rows,
        columns,
        expectedPattern.rowOffsets,
        expectedPattern.columnIndices,
        nullptr);
    if (!patternStatus.isOk()) {
        return Result<SparseMatrix>::failure(patternStatus);
    }

    for (const auto& contribution : contributions) {
        if (contribution.row >= rows || contribution.column >= columns) {
            return Result<SparseMatrix>::failure(sparseFailure(
                "invalid-sparse-index",
                std::to_string(contribution.row) + ":" +
                    std::to_string(contribution.column),
                "COO contribution index is outside the matrix dimensions."));
        }
        if (!std::isfinite(contribution.value)) {
            return Result<SparseMatrix>::failure(sparseFailure(
                "nonfinite-sparse-value",
                std::to_string(contribution.elementOrder) + ":" +
                    std::to_string(contribution.localOrder),
                "COO contribution values must be finite."));
        }
    }

    // The complete tuple fixes duplicate summation order independently of
    // worker completion order. stable_sort also preserves exact tuple ties.
    std::stable_sort(
        contributions.begin(),
        contributions.end(),
        [](const CooContribution& left, const CooContribution& right) {
            return std::tie(
                       left.row,
                       left.column,
                       left.elementOrder,
                       left.localOrder) <
                   std::tie(
                       right.row,
                       right.column,
                       right.elementOrder,
                       right.localOrder);
        });

    std::vector<double> values(expectedPattern.columnIndices.size(), 0.0);
    for (const auto& contribution : contributions) {
        const std::size_t begin = expectedPattern.rowOffsets[contribution.row];
        const std::size_t end = expectedPattern.rowOffsets[contribution.row + 1U];
        const auto first = expectedPattern.columnIndices.begin() + begin;
        const auto last = expectedPattern.columnIndices.begin() + end;
        const auto found = std::lower_bound(first, last, contribution.column);
        if (found == last || *found != contribution.column) {
            return Result<SparseMatrix>::failure(sparseFailure(
                "sparse-pattern-mismatch",
                std::to_string(contribution.row) + ":" +
                    std::to_string(contribution.column),
                "COO contribution is absent from the expected sparse pattern."));
        }

        const std::size_t position = static_cast<std::size_t>(
            std::distance(expectedPattern.columnIndices.begin(), found));
        values[position] += contribution.value;
        if (!std::isfinite(values[position])) {
            return Result<SparseMatrix>::failure(sparseFailure(
                "nonfinite-sparse-value",
                std::to_string(contribution.row) + ":" +
                    std::to_string(contribution.column),
                "Ordered COO duplicate summation produced a nonfinite value."));
        }
    }

    SparseMatrix matrix{
        rows,
        columns,
        expectedPattern.rowOffsets,
        expectedPattern.columnIndices,
        std::move(values)};
    const Status status = matrix.validate();
    if (!status.isOk()) {
        return Result<SparseMatrix>::failure(status);
    }
    return Result<SparseMatrix>::success(std::move(matrix));
}

std::size_t SparseMatrix::rows() const noexcept {
    return rows_;
}

std::size_t SparseMatrix::columns() const noexcept {
    return columns_;
}

const std::vector<std::size_t>& SparseMatrix::rowOffsets() const noexcept {
    return rowOffsets_;
}

const std::vector<std::size_t>& SparseMatrix::columnIndices() const noexcept {
    return columnIndices_;
}

const std::vector<double>& SparseMatrix::values() const noexcept {
    return values_;
}

Vector SparseMatrix::multiply(const Vector& rhs) const {
    if (columns_ != rhs.size()) {
        throw std::invalid_argument{
            "Sparse matrix-vector multiplication has incompatible dimensions."};
    }

    Vector result{rows_};
    for (std::size_t row = 0U; row < rows_; ++row) {
        double value = 0.0;
        for (std::size_t position = rowOffsets_[row];
             position < rowOffsets_[row + 1U];
             ++position) {
            value += values_[position] * rhs[columnIndices_[position]];
        }
        result[row] = value;
    }
    return result;
}

Status SparseMatrix::validate() const {
    return validateCsr(
        rows_, columns_, rowOffsets_, columnIndices_, &values_);
}

SparseMatrix::SparseMatrix(
    const std::size_t rows,
    const std::size_t columns,
    std::vector<std::size_t> rowOffsets,
    std::vector<std::size_t> columnIndices,
    std::vector<double> values)
    : rows_{rows},
      columns_{columns},
      rowOffsets_{std::move(rowOffsets)},
      columnIndices_{std::move(columnIndices)},
      values_{std::move(values)} {}

} // namespace fesa