#include "fesa/math/sparse_matrix.h"

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

#include <mkl.h>

#include "math/dense_blas_internal.h"

namespace fesa {
namespace {

/// @brief Checks the actual private backend integer width before conversion.
MKL_INT SparseSize(const std::size_t size) {
  const auto converted = dense_blas_internal::ToMklSize(size);
  if (!converted.HasValue()) {
    throw std::length_error{"Sparse size exceeds the MKL integer range."};
  }
  return converted.Value();
}

/// @brief Releases a call-local sparse handle before its borrowed arrays die.
class SparseHandle {
 public:
  SparseHandle() = default;
  SparseHandle(const SparseHandle&) = delete;
  SparseHandle& operator=(const SparseHandle&) = delete;
  ~SparseHandle() {
    if (handle != nullptr) static_cast<void>(mkl_sparse_destroy(handle));
  }
  /// @brief Reports normal release failure before publishing the candidate.
  sparse_status_t Destroy() {
    const auto owned = handle;
    handle = nullptr;
    return mkl_sparse_destroy(owned);
  }
  sparse_matrix_t handle{nullptr};
};

/// @brief Fails explicitly instead of substituting a scalar backend.
void CheckSparseStatus(const sparse_status_t status, const char* operation) {
  if (status != SPARSE_STATUS_SUCCESS) {
    throw std::runtime_error{std::string{"Sparse backend "} + operation +
                             " failed."};
  }
}

/// @brief Builds a structured sparse-matrix model failure.
Status SparseFailure(const std::string& code, const std::string& identity,
                     const std::string& message) {
  return Status::Failure(
      FailureCategory::kModel,
      {{Severity::kError, code, {{}, 0U}, "SPARSE_MATRIX", identity, message}});
}

/// @brief Validates canonical CSR shape, order, index, and finite-value rules.
Status ValidateCsr(const std::size_t rows, const std::size_t columns,
                   const std::vector<std::size_t>& row_offsets,
                   const std::vector<std::size_t>& column_indices,
                   const std::vector<double>* const values) {
  if (rows == (std::numeric_limits<std::size_t>::max)() ||
      row_offsets.size() != rows + 1U) {
    return SparseFailure(
        "invalid-sparse-shape", "row-offset-count",
        "CSR row offsets must contain exactly rows plus one entries.");
  }
  if (row_offsets.empty() || row_offsets.front() != 0U ||
      row_offsets.back() != column_indices.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() != column_indices.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 = row_offsets[row];
    const std::size_t end = row_offsets[row + 1U];
    if (begin > end || end > column_indices.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 (column_indices[position] >= columns) {
        return SparseFailure(
            "invalid-sparse-index", std::to_string(position),
            "CSR column index is outside the matrix dimensions.");
      }
      if (position > begin &&
          column_indices[position - 1U] >= column_indices[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& expected_pattern) {
  const Status pattern_status =
      ValidateCsr(rows, columns, expected_pattern.row_offsets,
                  expected_pattern.column_indices, nullptr);
  if (!pattern_status.IsOk()) {
    return Result<SparseMatrix>::Failure(pattern_status);
  }

  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.element_order) + ":" +
                            std::to_string(contribution.local_order),
                        "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.element_order,
                        left.local_order) < std::tie(right.row, right.column,
                                                     right.element_order,
                                                     right.local_order);
      });

  std::vector<double> values(expected_pattern.column_indices.size(), 0.0);
  for (const auto& contribution : contributions) {
    const std::size_t begin = expected_pattern.row_offsets[contribution.row];
    const std::size_t end = expected_pattern.row_offsets[contribution.row + 1U];
    const auto first = expected_pattern.column_indices.begin() + begin;
    const auto last = expected_pattern.column_indices.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(expected_pattern.column_indices.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, expected_pattern.row_offsets,
                      expected_pattern.column_indices, 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 row_offsets_;
}

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

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."};
  }
  if (!Validate().IsOk()) {
    throw std::invalid_argument{"Sparse multiplication requires valid CSR."};
  }
  for (std::size_t i = 0; i < rhs.Size(); ++i) {
    if (!std::isfinite(rhs[i]))
      throw std::invalid_argument{"Sparse multiplication RHS must be finite."};
  }
  const MKL_INT rows = SparseSize(rows_);
  const MKL_INT columns = SparseSize(columns_);
  static_cast<void>(SparseSize(values_.size()));
  Vector result{rows_};
  if (rows == 0 || columns == 0 || values_.empty()) return result;

  // Each const call owns its converted arrays and handle. Neither copy/move
  // nor concurrent calls retain pointers into another call's storage.
  std::vector<MKL_INT> row_offsets;
  std::vector<MKL_INT> column_indices;
  row_offsets.reserve(row_offsets_.size());
  column_indices.reserve(column_indices_.size());
  for (const auto offset : row_offsets_) {
    row_offsets.push_back(SparseSize(offset));
  }
  for (const auto column : column_indices_) {
    column_indices.push_back(SparseSize(column));
  }
  auto values = values_;
  SparseHandle sparse;
  CheckSparseStatus(
      mkl_sparse_d_create_csr(&sparse.handle, SPARSE_INDEX_BASE_ZERO, rows,
                              columns, row_offsets.data(),
                              row_offsets.data() + 1, column_indices.data(),
                              values.data()),
      "create CSR");
  const matrix_descr descriptor{SPARSE_MATRIX_TYPE_GENERAL,
                                SPARSE_FILL_MODE_FULL,
                                SPARSE_DIAG_NON_UNIT};
  CheckSparseStatus(
      mkl_sparse_d_mv(SPARSE_OPERATION_NON_TRANSPOSE, 1.0, sparse.handle,
                     descriptor, rhs.Data(), 0.0, result.Data()),
      "multiply");
  CheckSparseStatus(sparse.Destroy(), "destroy CSR");
  return result;
}

Status SparseMatrix::Validate() const {
  return ValidateCsr(rows_, columns_, row_offsets_, column_indices_, &values_);
}

SparseMatrix::SparseMatrix(const std::size_t rows, const std::size_t columns,
                           std::vector<std::size_t> row_offsets,
                           std::vector<std::size_t> column_indices,
                           std::vector<double> values)
    : rows_{rows},
      columns_{columns},
      row_offsets_{std::move(row_offsets)},
      column_indices_{std::move(column_indices)},
      values_{std::move(values)} {}

}  // namespace fesa