#include "math/dense_blas_internal.h"

#include <cmath>
#include <limits>
#include <stdexcept>
#include <utility>
#include <vector>

namespace fesa::dense_blas_internal {

ScopedSingleEquationThreads::ScopedSingleEquationThreads(
    const std::size_t equation_count) noexcept {
  // The measured 1x1 PARDISO case has backend/thread overhead and an
  // intermittently failing default-thread lifecycle. Broader sizes are not
  // inferred from this auxiliary evidence.
  if (equation_count != 1U) return;
  const int previous = mkl_set_num_threads_local(1);
  if (previous == 0) {
    restore_default_ = true;
  } else {
    static_cast<void>(mkl_set_num_threads_local(previous));
  }
}

ScopedSingleEquationThreads::~ScopedSingleEquationThreads() {
  if (restore_default_) static_cast<void>(mkl_set_num_threads_local(0));
}

bool CholeskyLowerAndInverse(const Matrix& matrix, Matrix& lower,
                            Matrix& inverse_lower) {
  if (matrix.Rows() != matrix.Columns())
    throw std::invalid_argument{"Cholesky requires a square matrix."};
  const auto size = ToMklSize(matrix.Rows());
  if (!size.HasValue())
    throw std::length_error{"Cholesky size exceeds the MKL integer range."};
  Matrix candidate_lower{matrix.Rows(), matrix.Columns()};
  Matrix candidate_inverse{matrix.Rows(), matrix.Columns()};
  if (matrix.Rows() == 0U) {
    lower = std::move(candidate_lower);
    inverse_lower = std::move(candidate_inverse);
    return true;
  }
  std::vector<double> values(matrix.Rows() * matrix.Columns());
  for (std::size_t row = 0; row < matrix.Rows(); ++row)
    for (std::size_t column = 0; column < matrix.Columns(); ++column) {
      if (!std::isfinite(matrix(row, column)))
        throw std::invalid_argument{"Cholesky inputs must be finite."};
      values[row * matrix.Columns() + column] = matrix(row, column);
    }
  if (LAPACKE_dpotrf(LAPACK_ROW_MAJOR, 'L', size.Value(), values.data(),
                     size.Value()) != 0) return false;
  for (std::size_t row = 0; row < matrix.Rows(); ++row)
    for (std::size_t column = 0; column <= row; ++column)
      candidate_lower(row, column) = values[row * matrix.Columns() + column];
  if (LAPACKE_dtrtri(LAPACK_ROW_MAJOR, 'L', 'N', size.Value(), values.data(),
                     size.Value()) != 0) return false;
  for (std::size_t row = 0; row < matrix.Rows(); ++row)
    for (std::size_t column = 0; column <= row; ++column) {
      const auto value = values[row * matrix.Columns() + column];
      if (!std::isfinite(candidate_lower(row, column)) || !std::isfinite(value))
        return false;
      candidate_inverse(row, column) = value;
    }
  lower = std::move(candidate_lower);
  inverse_lower = std::move(candidate_inverse);
  return true;
}

Vector MultiplyDiagonal(const Vector& diagonal, const Vector& rhs) {
  if (diagonal.Size() != rhs.Size())
    throw std::invalid_argument{"Diagonal product dimensions differ."};
  const auto size = ToMklSize(diagonal.Size());
  if (!size.HasValue())
    throw std::length_error{"Diagonal size exceeds the MKL integer range."};
  for (std::size_t i = 0; i < rhs.Size(); ++i) {
    if (!std::isfinite(diagonal[i]) || !std::isfinite(rhs[i]))
      throw std::invalid_argument{"Diagonal product inputs must be finite."};
  }
  Vector candidate{rhs.Size()};
  if (rhs.Size() == 0U) return candidate;
  // k=0 stores exactly one coefficient per independent row. MKL performs
  // each product without regrouping a caller's subsequent certified sums.
  cblas_dsbmv(CblasColMajor, CblasUpper, size.Value(), 0, 1.0, diagonal.Data(),
              1, rhs.Data(), 1, 0.0, candidate.Data(), 1);
  return candidate;
}

Matrix MultiplyTransposeLeft(const Matrix& left, const Matrix& right) {
  if (left.Rows() != right.Rows()) {
    throw std::invalid_argument{"Transpose product dimensions differ."};
  }
  Matrix transposed{left.Columns(), left.Rows()};
  for (std::size_t row = 0; row < left.Rows(); ++row)
    for (std::size_t column = 0; column < left.Columns(); ++column)
      transposed(column, row) = left(row, column);
  return transposed.Multiply(right);
}

Result<MKL_INT> ToMklSize(const std::size_t size) {
  if (size > static_cast<std::size_t>((std::numeric_limits<MKL_INT>::max)())) {
    return Result<MKL_INT>::Failure(Status::Failure(
        {{Severity::kError,
          "dense-blas-size-overflow",
          {},
          "",
          "",
          "Dense operation size exceeds the MKL integer range."}}));
  }
  return Result<MKL_INT>::Success(static_cast<MKL_INT>(size));
}

void CopyValues(const double* source, const std::size_t size,
                double* destination) {
  if (size == 0U) {
    return;
  }

  const auto mkl_size = ToMklSize(size);
  if (!mkl_size.HasValue()) {
    throw std::length_error{"Dense copy size exceeds the MKL integer range."};
  }
  cblas_dcopy(mkl_size.Value(), source, 1, destination, 1);
}

}  // namespace fesa::dense_blas_internal