#include "math/dense_blas_internal.h"

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

#include "math/finite_values.h"

namespace fesa::dense_blas_internal {

namespace {

using math_internal::IsFinite;

/// @brief Keeps dense failures backend-neutral for caller-specific diagnostics.
Status FactorFailure(const char* code, const char* message) {
  return Status::Failure({{Severity::kError, code, {}, "", "", message}});
}

/// @brief Preserves each value owner's size exception at the private boundary.
MKL_INT MklSizeOrThrow(const std::size_t size, const char* message) {
  const auto converted = ToMklSize(size);
  if (!converted.HasValue()) throw std::length_error{message};
  return converted.Value();
}

/// @brief Converts a contiguous BLAS vector length with its existing contract.
MKL_INT VectorSizeOrThrow(const std::size_t size) {
  return MklSizeOrThrow(size,
                       "Dense vector size exceeds the MKL integer range.");
}

/// @brief Converts a matrix dimension with its existing exception contract.
MKL_INT MatrixSizeOrThrow(const std::size_t size) {
  return MklSizeOrThrow(size,
                       "Dense matrix dimension exceeds the MKL integer range.");
}

/// @brief Rejects Matrix product inputs before allocation or empty returns.
void RequireFinite(const double* values, const std::size_t size) {
  for (std::size_t i = 0; i < size; ++i)
    if (!std::isfinite(values[i]))
      throw std::invalid_argument{"Matrix product inputs must be finite."};
}

}  // namespace

DenseFactor::DenseFactor(const Matrix& coefficients,
                         const bool positive_definite)
    : factors_{coefficients},
      pivots_(positive_definite ? 0U : coefficients.Rows()),
      positive_definite_{positive_definite} {
  for (std::size_t row = 0; row < coefficients.Rows(); ++row) {
    double row_sum = 0.0;
    for (std::size_t column = 0; column < coefficients.Columns(); ++column)
      row_sum += std::abs(coefficients(row, column));
    coefficient_norm_inf_ = (std::max)(coefficient_norm_inf_, row_sum);
  }
}

Result<DenseFactor> DenseFactor::Lu(const Matrix& coefficients) {
  return Create(coefficients, false);
}

Result<DenseFactor> DenseFactor::Cholesky(const Matrix& coefficients) {
  return Create(coefficients, true);
}

Result<DenseFactor> DenseFactor::Create(const Matrix& coefficients,
                                        const bool positive_definite) {
  if (coefficients.Rows() != coefficients.Columns())
    return Result<DenseFactor>::Failure(FactorFailure(
        "dense-factor-shape", "Dense factor requires a square matrix."));
  const auto size = ToMklSize(coefficients.Rows());
  if (!size.HasValue()) return Result<DenseFactor>::Failure(size.GetStatus());
  if (!IsFinite(coefficients))
    return Result<DenseFactor>::Failure(FactorFailure(
        "dense-factor-nonfinite", "Dense factor inputs must be finite."));
  DenseFactor candidate{coefficients, positive_definite};
  if (size.Value() != 0) {
    const auto info =
        positive_definite
            ? LAPACKE_dpotrf(LAPACK_ROW_MAJOR, 'L', size.Value(),
                             candidate.factors_.Data(), size.Value())
            : LAPACKE_dgetrf(LAPACK_ROW_MAJOR, size.Value(), size.Value(),
                             candidate.factors_.Data(), size.Value(),
                             candidate.pivots_.data());
    if (info != 0 || !IsFinite(candidate.factors_))
      return Result<DenseFactor>::Failure(FactorFailure(
          "dense-factor-pivot", "Dense numerical factorization failed."));
  }
  return Result<DenseFactor>::Success(std::move(candidate));
}

Result<Matrix> DenseFactor::Solve(const Matrix& rhs,
                                  const bool transpose) const {
  return SolveImpl(rhs, transpose, false, false);
}

Result<Matrix> DenseFactor::SolveRight(const Matrix& rhs,
                                       const bool transpose) const {
  return SolveImpl(rhs, transpose, true, false);
}

Result<Matrix> DenseFactor::SolveLower(const Matrix& rhs,
                                       const bool transpose) const {
  return SolveImpl(rhs, transpose, false, true);
}

Result<Matrix> DenseFactor::SolveLowerRight(const Matrix& rhs,
                                            const bool transpose) const {
  return SolveImpl(rhs, transpose, true, true);
}

Result<Matrix> DenseFactor::SolveImpl(const Matrix& rhs, const bool transpose,
                                      const bool right,
                                      const bool lower_only) const {
  const auto order = factors_.Rows();
  if ((right ? rhs.Columns() : rhs.Rows()) != order ||
      (lower_only && !positive_definite_))
    return Result<Matrix>::Failure(FactorFailure(
        "dense-solve-shape", "Dense solve shape or factor property differs."));
  const auto count = ToMklSize(right ? rhs.Rows() : rhs.Columns());
  if (!count.HasValue()) return Result<Matrix>::Failure(count.GetStatus());
  if (!IsFinite(rhs))
    return Result<Matrix>::Failure(FactorFailure(
        "dense-solve-nonfinite", "Dense solve RHS must be finite."));
  if (order == 0U || count.Value() == 0)
    return Result<Matrix>::Success(Matrix{rhs});
  Matrix candidate = right && !lower_only ? Matrix{order, rhs.Rows()}
                                         : Matrix{rhs};
  const auto size = ToMklSize(order).Value();
  if (lower_only) {
    cblas_dtrsm(CblasRowMajor, right ? CblasRight : CblasLeft, CblasLower,
                transpose ? CblasTrans : CblasNoTrans, CblasNonUnit,
                right ? count.Value() : size, right ? size : count.Value(), 1.0,
                factors_.Data(), size, candidate.Data(),
                right ? size : count.Value());
  } else {
    // X*op(A)=B becomes op(A)^T*X^T=B^T. GETRS uses a compact n-by-nrhs
    // row-major RHS with ldb=nrhs, including nonsquare multi-RHS storage.
    if (right) {
      for (std::size_t i = 0; i < rhs.Rows(); ++i)
        for (std::size_t j = 0; j < order; ++j) candidate(j, i) = rhs(i, j);
    }
    const auto info =
        positive_definite_
            ? LAPACKE_dpotrs(LAPACK_ROW_MAJOR, 'L', size, count.Value(),
                             factors_.Data(), size, candidate.Data(),
                             count.Value())
            : LAPACKE_dgetrs(LAPACK_ROW_MAJOR, transpose != right ? 'T' : 'N',
                             size, count.Value(), factors_.Data(), size,
                             pivots_.data(), candidate.Data(), count.Value());
    if (info != 0)
      return Result<Matrix>::Failure(
          FactorFailure("dense-solve-backend", "Dense backend solve failed."));
    if (right) {
      Matrix transposed{rhs.Rows(), order};
      for (std::size_t i = 0; i < rhs.Rows(); ++i)
        for (std::size_t j = 0; j < order; ++j)
          transposed(i, j) = candidate(j, i);
      candidate = std::move(transposed);
    }
  }
  if (!IsFinite(candidate))
    return Result<Matrix>::Failure(FactorFailure(
        "dense-solve-nonfinite", "Dense solve result is nonfinite."));
  return Result<Matrix>::Success(std::move(candidate));
}

Result<Matrix> DenseFactor::Lower() const {
  if (!positive_definite_)
    return Result<Matrix>::Failure(FactorFailure(
        "dense-factor-property", "An LU factor has no Cholesky lower factor."));
  Matrix lower{factors_.Rows(), factors_.Columns()};
  for (std::size_t i = 0; i < factors_.Rows(); ++i)
    for (std::size_t j = 0; j <= i; ++j) lower(i, j) = factors_(i, j);
  return Result<Matrix>::Success(std::move(lower));
}

double DenseFactor::Determinant() const {
  double determinant = 1.0;
  for (std::size_t i = 0; i < factors_.Rows(); ++i) {
    determinant *= factors_(i, i);
    if (positive_definite_)
      determinant *= factors_(i, i);
    else if (pivots_[i] != static_cast<MKL_INT>(i + 1U))
      determinant = -determinant;
  }
  return determinant;
}

Result<double> DenseFactor::ReciprocalConditionInfinity() const {
  if (!(coefficient_norm_inf_ > 0.0) || !std::isfinite(coefficient_norm_inf_))
    return Result<double>::Failure(FactorFailure(
        "dense-condition-norm", "Dense coefficient infinity norm is invalid."));
  const auto order = factors_.Rows();
  std::vector<double> row_sums(order, 0.0);
  Matrix rhs{order, 1U};
  for (std::size_t column = 0; column < order; ++column) {
    if (column != 0U) rhs(column - 1U, 0U) = 0.0;
    rhs(column, 0U) = 1.0;
    const auto solved = Solve(rhs);
    if (!solved.HasValue()) return Result<double>::Failure(solved.GetStatus());
    for (std::size_t row = 0; row < order; ++row)
      row_sums[row] += std::abs(solved.Value()(row, 0U));
  }
  double inverse_norm = 0.0;
  for (const auto row_sum : row_sums)
    inverse_norm = (std::max)(inverse_norm, row_sum);
  if (!(inverse_norm > 0.0) || !std::isfinite(inverse_norm))
    return Result<double>::Failure(FactorFailure(
        "dense-condition-norm", "Dense inverse infinity norm is invalid."));
  return Result<double>::Success((1.0 / coefficient_norm_inf_) / inverse_norm);
}

double DotValues(const double* left, const std::size_t size,
                  const double* right) {
  if (size == 0U) return 0.0;
  return cblas_ddot(VectorSizeOrThrow(size), left, 1, right, 1);
}

double NormValues(const double* values, const std::size_t size) {
  if (size == 0U) return 0.0;
  return cblas_dnrm2(VectorSizeOrThrow(size), values, 1);
}

void ScaleValues(const double alpha, double* values, const std::size_t size) {
  if (size == 0U) return;
  cblas_dscal(VectorSizeOrThrow(size), alpha, values, 1);
}

void AxpyValues(const double alpha, const double* x, const std::size_t size,
                double* y) {
  if (size == 0U) return;
  cblas_daxpy(VectorSizeOrThrow(size), alpha, x, 1, y, 1);
}

Vector Multiply(const Matrix& left, const Vector& right) {
  if (left.Columns() != right.Size())
    throw std::invalid_argument{
        "Matrix-vector multiplication has incompatible dimensions."};
  RequireFinite(left.Data(), left.Rows() * left.Columns());
  RequireFinite(right.Data(), right.Size());
  Vector result{left.Rows()};
  if (left.Rows() == 0U || left.Columns() == 0U) return result;
  const auto rows = MatrixSizeOrThrow(left.Rows());
  const auto columns = MatrixSizeOrThrow(left.Columns());
  cblas_dgemv(CblasRowMajor, CblasNoTrans, rows, columns, 1.0, left.Data(),
              columns, right.Data(), 1, 0.0, result.Data(), 1);
  return result;
}

Matrix Multiply(const Matrix& left, const Matrix& right) {
  if (left.Columns() != right.Rows())
    throw std::invalid_argument{
        "Matrix multiplication has incompatible dimensions."};
  RequireFinite(left.Data(), left.Rows() * left.Columns());
  RequireFinite(right.Data(), right.Rows() * right.Columns());
  Matrix result{left.Rows(), right.Columns()};
  if (left.Rows() == 0U || left.Columns() == 0U || right.Columns() == 0U)
    return result;
  const auto rows = MatrixSizeOrThrow(left.Rows());
  const auto columns = MatrixSizeOrThrow(right.Columns());
  const auto inner = MatrixSizeOrThrow(left.Columns());
  cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, rows, columns, inner,
              1.0, left.Data(), inner, right.Data(), columns, 0.0,
              result.Data(), columns);
  return result;
}

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

ScopedLocalDenseThreads::ScopedLocalDenseThreads() noexcept {
  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));
  }
}

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

LocalDensePolicy ReadScopedLocalDensePolicy() noexcept {
  const ScopedLocalDenseThreads scope{};
  return {mkl_domain_get_max_threads(MKL_DOMAIN_BLAS), mkl_get_dynamic()};
}

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."};
  }
  RequireFinite(left.Data(), left.Rows() * left.Columns());
  RequireFinite(right.Data(), right.Rows() * right.Columns());
  Matrix result{left.Columns(), right.Columns()};
  if (left.Columns() == 0U || left.Rows() == 0U || right.Columns() == 0U)
    return result;
  const auto rows = MatrixSizeOrThrow(left.Columns());
  const auto columns = MatrixSizeOrThrow(right.Columns());
  const auto inner = MatrixSizeOrThrow(left.Rows());
  // Row-major A is stored inner-by-rows even when GEMM consumes A-transpose.
  cblas_dgemm(CblasRowMajor, CblasTrans, CblasNoTrans, rows, columns, inner,
              1.0, left.Data(), rows, right.Data(), columns, 0.0,
              result.Data(), columns);
  return result;
}

bool OuterProduct(const double* left, const std::size_t rows,
                  const double* right, const std::size_t columns,
                  double* destination) {
  return MultiplyTransposeRight(left, rows, right, columns, 1U, destination);
}

bool MultiplyTransposeRight(const double* left, const std::size_t rows,
                            const double* right, const std::size_t columns,
                            const std::size_t inner, double* destination) {
  const auto m = ToMklSize(rows);
  const auto n = ToMklSize(columns);
  const auto k = ToMklSize(inner);
  const auto maximum = std::numeric_limits<std::size_t>::max() / sizeof(double);
  if (!m.HasValue() || !n.HasValue() || !k.HasValue() ||
      (columns != 0U && rows > maximum / columns) ||
      (inner != 0U && (rows > maximum / inner || columns > maximum / inner)))
    throw std::length_error{"Dense product dimensions exceed backend limits."};
  if (rows == 0U || columns == 0U) return true;
  if (destination == nullptr ||
      (inner != 0U && (left == nullptr || right == nullptr)))
    throw std::invalid_argument{"Dense product requires nonnull storage."};
  if (inner == 0U) {
    std::fill_n(destination, rows * columns, 0.0);
    return true;
  }
  for (std::size_t i = 0U; i < rows * inner; ++i)
    if (!std::isfinite(left[i])) return false;
  for (std::size_t i = 0U; i < columns * inner; ++i)
    if (!std::isfinite(right[i])) return false;
  cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m.Value(), n.Value(),
              k.Value(), 1.0, left, k.Value(), right, k.Value(), 0.0,
              destination, n.Value());
  return true;
}

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