#include "math/dense_blas_internal.h"

#include <gtest/gtest.h>

#include <array>
#include <cfenv>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits>
#include <stdexcept>

#if defined(_M_X64) || defined(_M_IX86)
#include <xmmintrin.h>
#endif

namespace fesa::dense_blas_internal {
namespace {

template <std::size_t N>
void ExpectRankOneContract(const std::array<double, N>& left,
                           const std::array<double, N>& right) {
  const auto saved_left = left;
  const auto saved_right = right;
  const int rounding = std::fegetround();
#if defined(_M_X64) || defined(_M_IX86)
  const unsigned control = _mm_getcsr() & ~0x3fU;
#endif
  std::array<double, N * N> product{};
  std::array<double, N * N> original{};
  product.fill(std::numeric_limits<double>::quiet_NaN());
  const ScopedLocalDenseThreads scope{};
  ASSERT_TRUE(MultiplyTransposeRight(left.data(), N, right.data(), N, 1U,
                                     original.data()));
  ASSERT_TRUE(OuterProduct(left.data(), N, right.data(), N, product.data()));
  std::size_t zero_sign_differences = 0U;
  for (std::size_t i = 0U; i < product.size(); ++i) {
    std::uint64_t original_bits = 0U;
    std::uint64_t product_bits = 0U;
    std::memcpy(&original_bits, &original[i], sizeof(original_bits));
    std::memcpy(&product_bits, &product[i], sizeof(product_bits));
    if (original[i] == 0.0 && product[i] == 0.0) {
      zero_sign_differences += original_bits != product_bits ? 1U : 0U;
      EXPECT_EQ(original[i], product[i]);
    } else {
      EXPECT_EQ(original_bits, product_bits)
          << "entry=" << i << " left=" << left[i / N]
          << " right=" << right[i % N]
          << " original=" << original[i] << " candidate=" << product[i];
    }
  }
  std::cout << "DGER boundary shape=" << N << "x" << N
            << " zero_sign_differences=" << zero_sign_differences << '\n';
  EXPECT_EQ(std::memcmp(saved_left.data(), left.data(), sizeof(left)), 0);
  EXPECT_EQ(std::memcmp(saved_right.data(), right.data(), sizeof(right)), 0);
  EXPECT_EQ(std::fegetround(), rounding);
#if defined(_M_X64) || defined(_M_IX86)
  // Arithmetic may raise sticky exception flags; control bits must be retained.
  EXPECT_EQ(_mm_getcsr() & ~0x3fU, control);
#endif
}

TEST(DenseBlasInternal,
     OuterProductUsesRectangularStorageAndPreservesRejection) {
  const std::array<double, 2U> left{2.0, -3.0};
  const std::array<double, 3U> right{4.0, -5.0, 0.25};
  std::array<double, 6U> product{};
  ASSERT_TRUE(OuterProduct(left.data(), left.size(), right.data(), right.size(),
                           product.data()));
  const std::array<double, 6U> expected{8.0, -10.0, 0.5, -12.0, 15.0, -0.75};
  EXPECT_EQ(product, expected);
  const auto accepted = product;
  const std::array<double, 2U> invalid{
      1.0, std::numeric_limits<double>::quiet_NaN()};
  EXPECT_FALSE(OuterProduct(invalid.data(), invalid.size(), right.data(),
                            right.size(), product.data()));
  EXPECT_EQ(product, accepted);
  EXPECT_TRUE(OuterProduct(nullptr, 0U, nullptr, 3U, nullptr));
  EXPECT_TRUE(OuterProduct(nullptr, 2U, nullptr, 0U, nullptr));
  EXPECT_THROW(OuterProduct(nullptr, 2U, right.data(), 3U, product.data()),
               std::invalid_argument);
  const auto too_large =
      static_cast<std::size_t>((std::numeric_limits<MKL_INT>::max)()) + 1U;
  EXPECT_THROW(OuterProduct(nullptr, too_large, nullptr, 1U, nullptr),
               std::length_error);
  const auto too_many_elements =
      std::numeric_limits<std::size_t>::max() / sizeof(double) + 1U;
  EXPECT_THROW(OuterProduct(nullptr, too_many_elements, nullptr, 1U, nullptr),
               std::length_error);
  product = accepted;
  EXPECT_THROW(OuterProduct(left.data(), too_large, right.data(), 1U,
                            product.data()), std::length_error);
  EXPECT_EQ(product, accepted);
  EXPECT_THROW(OuterProduct(nullptr, 2U, right.data(), 3U, product.data()),
               std::invalid_argument);
  EXPECT_EQ(product, accepted);
}

TEST(DenseBlasInternal, LocalDenseScopeRestoresTlsAfterProductsAndUnwind) {
  const int previous = mkl_set_num_threads_local(0);
  struct Restore {
    int previous;
    ~Restore() { static_cast<void>(mkl_set_num_threads_local(previous)); }
  } restore{previous};
  const int default_threads = mkl_get_max_threads();
  const std::array<double, 2U> left{0.0, 2.0};
  const std::array<double, 2U> right{3.0, 4.0};
  std::array<double, 4U> product{};
  product.fill(std::numeric_limits<double>::quiet_NaN());
  {
    const ScopedLocalDenseThreads scope{};
    EXPECT_EQ(mkl_get_max_threads(), 1);
    EXPECT_TRUE(
        OuterProduct(left.data(), 2U, right.data(), 2U, product.data()));
    EXPECT_DOUBLE_EQ(product[0U], 0.0);
    EXPECT_DOUBLE_EQ(product[3U], 8.0);
  }
  EXPECT_EQ(mkl_get_max_threads(), default_threads);
  EXPECT_EQ(mkl_set_num_threads_local(0), 0);
  try {
    const ScopedLocalDenseThreads scope{};
    OuterProduct(nullptr, 2U, right.data(), 2U, product.data());
  } catch (const std::invalid_argument&) {
  }
  EXPECT_EQ(mkl_set_num_threads_local(0), 0);
  static_cast<void>(mkl_set_num_threads_local(3));
  {
    const ScopedLocalDenseThreads scope{};
    EXPECT_EQ(mkl_get_max_threads(), 3);
  }
  EXPECT_EQ(mkl_get_max_threads(), 3);
}

TEST(DenseBlasInternal, OuterProductLeavesFiniteInputOverflowForOwnerGate) {
  const std::array<double, 1U> left{std::numeric_limits<double>::max()};
  const std::array<double, 1U> right{2.0};
  std::array<double, 1U> product{};
  EXPECT_TRUE(OuterProduct(left.data(), 1U, right.data(), 1U, product.data()));
  EXPECT_TRUE(std::isinf(product[0U]));
}

TEST(DenseBlasInternal, OuterProductMatchesK1ForSignedZeroAndExtremeScale) {
  const std::array<double, 6U> left{
      0.0, -0.0, std::numeric_limits<double>::denorm_min(),
      std::ldexp(1.0, -500), std::numeric_limits<double>::max(),
      -std::numeric_limits<double>::max()};
  const std::array<double, 6U> right{
      0.0, -0.0, 2.0, -2.0, std::ldexp(1.0, 500),
      std::numeric_limits<double>::denorm_min()};
  ExpectRankOneContract(left, right);
}

TEST(DenseBlasInternal, OuterProduct24PreservesNonzeroAndFailureBoundaries) {
  std::array<double, 24U> left{};
  std::array<double, 24U> right{};
  for (std::size_t i = 0U; i < left.size(); ++i) {
    left[i] = std::ldexp(static_cast<double>(i + 1U), -8);
    right[i] = std::ldexp(-static_cast<double>(i + 2U), -10);
  }
  left[0U] = 0.0;
  left[1U] = -0.0;
  left[2U] = std::numeric_limits<double>::denorm_min();
  left[3U] = -std::numeric_limits<double>::denorm_min();
  left[4U] = std::ldexp(1.0, -500);
  left[5U] = std::numeric_limits<double>::max();
  right[0U] = -0.0;
  right[1U] = 0.0;
  right[2U] = std::numeric_limits<double>::denorm_min();
  right[3U] = 2.0;
  right[4U] = -2.0;
  right[5U] = std::ldexp(1.0, 500);
  ExpectRankOneContract(left, right);
}

TEST(DenseBlasInternal, OuterProductRejectsEitherNonfiniteInputBeforeWriting) {
  std::array<double, 2U> left{2.0, -3.0};
  std::array<double, 3U> right{4.0, -5.0, 0.25};
  std::array<double, 6U> destination{1.0, -2.0, 3.0, -4.0, 5.0, -6.0};
  const auto accepted = destination;
  for (const double invalid : {std::numeric_limits<double>::quiet_NaN(),
                               std::numeric_limits<double>::infinity(),
                               -std::numeric_limits<double>::infinity()}) {
    right[2U] = invalid;
    EXPECT_FALSE(OuterProduct(left.data(), left.size(), right.data(),
                              right.size(), destination.data()));
    EXPECT_EQ(destination, accepted);
    right[2U] = 0.25;
    left[1U] = invalid;
    EXPECT_FALSE(OuterProduct(left.data(), left.size(), right.data(),
                              right.size(), destination.data()));
    EXPECT_EQ(destination, accepted);
    left[1U] = -3.0;
  }
}

TEST(DenseBlasInternal, TransposeRightProductUsesInnerAndRectangularStrides) {
  const std::array<double, 6U> left{1, 2, 3, 4, 5, 6};
  const std::array<double, 12U> right{7,  8,  9,  10, 11, 12,
                                      13, 14, 15, 16, 17, 18};
  std::array<double, 8U> product;
  ASSERT_TRUE(MultiplyTransposeRight(left.data(), 2U, right.data(), 4U, 3U,
                                     product.data()));
  const std::array<double, 8U> expected{50, 68, 86, 104, 122, 167, 212, 257};
  EXPECT_EQ(product, expected);
  EXPECT_TRUE(
      MultiplyTransposeRight(nullptr, 2U, nullptr, 4U, 0U, product.data()));
  for (const auto value : product) EXPECT_DOUBLE_EQ(value, 0.0);
  auto nonfinite = right;
  nonfinite[8U] = std::numeric_limits<double>::infinity();
  product.fill(31.0);
  EXPECT_FALSE(MultiplyTransposeRight(left.data(), 2U, nonfinite.data(), 4U, 3U,
                                      product.data()));
  for (const auto value : product) EXPECT_DOUBLE_EQ(value, 31.0);
}

TEST(DenseBlasInternal, CDup003ConvertsZeroNormalAndOverflowLengths) {
  const auto zero = ToMklSize(0U);
  ASSERT_TRUE(zero.HasValue());
  EXPECT_EQ(zero.Value(), 0);

  const auto normal = ToMklSize(3U);
  ASSERT_TRUE(normal.HasValue());
  EXPECT_EQ(normal.Value(), 3);

  const std::size_t overflow =
      static_cast<std::size_t>((std::numeric_limits<MKL_INT>::max)()) + 1U;
  const auto rejected = ToMklSize(overflow);
  EXPECT_FALSE(rejected.HasValue());
  EXPECT_FALSE(rejected.GetStatus().IsOk());
}

TEST(DenseBlasInternal, CDup003CopiesZeroAndNonzeroContiguousValues) {
  EXPECT_NO_THROW(CopyValues(nullptr, 0U, nullptr));

  const std::array<double, 3U> source{1.25, -2.5, 4.0};
  std::array<double, 3U> destination{};
  CopyValues(source.data(), source.size(), destination.data());

  EXPECT_EQ(destination, source);
}

TEST(DenseBlasInternal, TransposeProductChecksShapeAndKnownRectangularAction) {
  Matrix left{3U, 2U};
  Matrix right{3U, 1U};
  for (std::size_t row = 0; row < 3U; ++row) {
    left(row, 0U) = static_cast<double>(row + 1U);
    left(row, 1U) = 2.0;
    right(row, 0U) = static_cast<double>(row);
  }
  const auto product = MultiplyTransposeLeft(left, right);
  ASSERT_EQ(product.Rows(), 2U);
  ASSERT_EQ(product.Columns(), 1U);
  EXPECT_DOUBLE_EQ(product(0U, 0U), 8.0);
  EXPECT_DOUBLE_EQ(product(1U, 0U), 6.0);
  EXPECT_THROW(static_cast<void>(MultiplyTransposeLeft(left, Matrix{2U, 1U})),
               std::invalid_argument);
  const auto empty = MultiplyTransposeLeft(Matrix{0U, 2U}, Matrix{0U, 3U});
  EXPECT_EQ(empty.Rows(), 2U);
  EXPECT_EQ(empty.Columns(), 3U);
  EXPECT_DOUBLE_EQ(empty(1U, 2U), 0.0);
  left(0U, 0U) = std::numeric_limits<double>::quiet_NaN();
  EXPECT_THROW(static_cast<void>(MultiplyTransposeLeft(left, right)),
               std::invalid_argument);
  left(0U, 0U) = 1.0;
  right(0U, 0U) = std::numeric_limits<double>::infinity();
  EXPECT_THROW(static_cast<void>(MultiplyTransposeLeft(left, right)),
               std::invalid_argument);
}

TEST(DenseBlasInternal, DiagonalProductKeepsEntrywiseCancellationAndChecks) {
  Vector diagonal{3U};
  diagonal[0U] = 4;
  diagonal[1U] = 2;
  diagonal[2U] = -1;
  Vector input{3U};
  const double large = std::ldexp(1.0, 56);
  input[0U] = 1;
  input[1U] = -large;
  input[2U] = 2 * large;
  const auto product = MultiplyDiagonal(diagonal, input);
  EXPECT_DOUBLE_EQ(product[0U] + (product[1U] - product[2U]), 4);
  EXPECT_EQ(MultiplyDiagonal(Vector{0U}, Vector{0U}).Size(), 0U);
  EXPECT_THROW(static_cast<void>(MultiplyDiagonal(diagonal, Vector{2U})),
               std::invalid_argument);
  input[0U] = std::numeric_limits<double>::infinity();
  EXPECT_THROW(static_cast<void>(MultiplyDiagonal(diagonal, input)),
               std::invalid_argument);
  input[0U] = 1;
  diagonal[0U] = std::numeric_limits<double>::quiet_NaN();
  EXPECT_THROW(static_cast<void>(MultiplyDiagonal(diagonal, input)),
               std::invalid_argument);
}

TEST(DenseBlasInternal, CholeskyAndInverseMatchKnownFactorsAndKeepFailedOutput) {
  Matrix matrix{2U, 2U};
  matrix(0U, 0U) = 4;
  matrix(0U, 1U) = 2;
  matrix(1U, 0U) = 2;
  matrix(1U, 1U) = 5;
  Matrix lower{0U, 0U}, inverse{0U, 0U};
  ASSERT_TRUE(CholeskyLowerAndInverse(matrix, lower, inverse));
  EXPECT_DOUBLE_EQ(lower(0U, 0U), 2);
  EXPECT_DOUBLE_EQ(lower(1U, 0U), 1);
  EXPECT_DOUBLE_EQ(lower(1U, 1U), 2);
  EXPECT_DOUBLE_EQ(lower(0U, 1U), 0);
  EXPECT_DOUBLE_EQ(inverse(0U, 0U), .5);
  EXPECT_DOUBLE_EQ(inverse(1U, 0U), -.25);
  EXPECT_DOUBLE_EQ(inverse(1U, 1U), .5);
  matrix(1U, 1U) = 0;
  EXPECT_FALSE(CholeskyLowerAndInverse(matrix, lower, inverse));
  EXPECT_DOUBLE_EQ(lower(1U, 1U), 2);
  EXPECT_DOUBLE_EQ(inverse(1U, 1U), .5);
  EXPECT_THROW(static_cast<void>(CholeskyLowerAndInverse(
                   Matrix{2U, 1U}, lower, inverse)), std::invalid_argument);
  matrix(0U, 0U) = std::numeric_limits<double>::quiet_NaN();
  EXPECT_THROW(static_cast<void>(CholeskyLowerAndInverse(matrix, lower, inverse)),
               std::invalid_argument);
}

TEST(DenseBlasInternal,
     CholeskyFailureReasonPreservesCandidatesAndResetsOnReuse) {
  Matrix matrix{2U, 2U};
  matrix(0U, 0U) = 4.0;
  matrix(1U, 1U) = 9.0;
  Matrix lower{0U, 0U};
  Matrix inverse{0U, 0U};
  auto reason = CholeskyFailure::kInverse;
  ASSERT_TRUE(CholeskyLowerAndInverse(matrix, lower, inverse, &reason));
  EXPECT_EQ(reason, CholeskyFailure::kNone);
  matrix(1U, 1U) = -1.0;
  EXPECT_FALSE(CholeskyLowerAndInverse(matrix, lower, inverse, &reason));
  EXPECT_EQ(reason, CholeskyFailure::kPivot);
  EXPECT_DOUBLE_EQ(lower(1U, 1U), 3.0);
  EXPECT_DOUBLE_EQ(inverse(1U, 1U), 1.0 / 3.0);
  matrix(1U, 1U) = 9.0;
  EXPECT_TRUE(CholeskyLowerAndInverse(matrix, lower, inverse, &reason));
  EXPECT_EQ(reason, CholeskyFailure::kNone);
  reason = CholeskyFailure::kInverse;
  EXPECT_THROW(CholeskyLowerAndInverse(Matrix{2U, 1U}, lower, inverse, &reason),
               std::invalid_argument);
  EXPECT_EQ(reason, CholeskyFailure::kNone);
}

TEST(DenseBlasInternal, SingleEquationScopeRestoresTlsAndRespectsOverrides) {
  const int previous = mkl_set_num_threads_local(0);
  struct Restore {
    int previous;
    ~Restore() { static_cast<void>(mkl_set_num_threads_local(previous)); }
  } restore{previous};
  const int default_threads = mkl_get_max_threads();
  {
    const ScopedSingleEquationThreads scope{1U};
    EXPECT_EQ(mkl_get_max_threads(), 1);
  }
  EXPECT_EQ(mkl_get_max_threads(), default_threads);
  EXPECT_EQ(mkl_set_num_threads_local(0), 0);
  {
    const ScopedSingleEquationThreads scope{2U};
    EXPECT_EQ(mkl_get_max_threads(), default_threads);
  }
  try {
    const ScopedSingleEquationThreads scope{1U};
    throw std::runtime_error{"scope failure"};
  } catch (const std::runtime_error&) {
  }
  EXPECT_EQ(mkl_set_num_threads_local(0), 0);
  static_cast<void>(mkl_set_num_threads_local(3));
  {
    const ScopedSingleEquationThreads scope{1U};
    EXPECT_EQ(mkl_get_max_threads(), 3);
  }
  EXPECT_EQ(mkl_get_max_threads(), 3);
}

}  // namespace
}  // namespace fesa::dense_blas_internal