// Benchmarks for special functions beyond what bench_cwise_math.cpp covers.
// Includes Bessel functions, two-argument functions (igamma, betainc),
// and additional functions (lgamma, digamma, zeta, polygamma).
// SPDX-FileCopyrightText: The Eigen Authors
// SPDX-License-Identifier: MPL-2.0

#include <benchmark/benchmark.h>
#include <Eigen/Core>
#include <contrib/Eigen/SpecialFunctions>

using namespace Eigen;

// Macro for unary special functions on arrays.
#define BENCH_SPECIAL_UNARY(NAME, EXPR, LO, HI)                                                          \
  template <typename Scalar>                                                                             \
  static void BM_##NAME(benchmark::State& state) {                                                       \
    const Index n = state.range(0);                                                                      \
    using Arr = Array<Scalar, Dynamic, 1>;                                                               \
    Arr a = (Arr::Random(n) + Scalar(1)) * Scalar((double(HI) - double(LO)) / 2.0) + Scalar(LO);         \
    Arr b(n);                                                                                            \
    for (auto _ : state) {                                                                               \
      b = EXPR;                                                                                          \
      benchmark::DoNotOptimize(b.data());                                                                \
    }                                                                                                    \
    state.counters["Elements/s"] = benchmark::Counter(n, benchmark::Counter::kIsIterationInvariantRate); \
    state.SetBytesProcessed(state.iterations() * n * sizeof(Scalar) * 2);                                \
  }

// Macro for binary special functions on arrays.
#define BENCH_SPECIAL_BINARY(NAME, EXPR, LO_A, HI_A, LO_B, HI_B)                                         \
  template <typename Scalar>                                                                             \
  static void BM_##NAME(benchmark::State& state) {                                                       \
    const Index n = state.range(0);                                                                      \
    using Arr = Array<Scalar, Dynamic, 1>;                                                               \
    Arr a = (Arr::Random(n) + Scalar(1)) * Scalar((double(HI_A) - double(LO_A)) / 2.0) + Scalar(LO_A);   \
    Arr b = (Arr::Random(n) + Scalar(1)) * Scalar((double(HI_B) - double(LO_B)) / 2.0) + Scalar(LO_B);   \
    Arr c(n);                                                                                            \
    for (auto _ : state) {                                                                               \
      c = EXPR;                                                                                          \
      benchmark::DoNotOptimize(c.data());                                                                \
    }                                                                                                    \
    state.counters["Elements/s"] = benchmark::Counter(n, benchmark::Counter::kIsIterationInvariantRate); \
    state.SetBytesProcessed(state.iterations() * n * sizeof(Scalar) * 3);                                \
  }

// --- Unary special functions ---
BENCH_SPECIAL_UNARY(Lgamma, Eigen::lgamma(a), 0.1, 20)
BENCH_SPECIAL_UNARY(Digamma, Eigen::digamma(a), 0.1, 20)

// --- Bessel functions (first kind) ---
BENCH_SPECIAL_UNARY(BesselI0, Eigen::bessel_i0(a), 0, 10)
BENCH_SPECIAL_UNARY(BesselI1, Eigen::bessel_i1(a), 0, 10)
BENCH_SPECIAL_UNARY(BesselI0e, Eigen::bessel_i0e(a), 0, 100)
BENCH_SPECIAL_UNARY(BesselI1e, Eigen::bessel_i1e(a), 0, 100)
BENCH_SPECIAL_UNARY(BesselJ0, Eigen::bessel_j0(a), 0, 20)
BENCH_SPECIAL_UNARY(BesselJ1, Eigen::bessel_j1(a), 0, 20)

// --- Bessel functions (second kind) ---
BENCH_SPECIAL_UNARY(BesselY0, Eigen::bessel_y0(a), 0.1, 20)
BENCH_SPECIAL_UNARY(BesselY1, Eigen::bessel_y1(a), 0.1, 20)
BENCH_SPECIAL_UNARY(BesselK0, Eigen::bessel_k0(a), 0.1, 20)
BENCH_SPECIAL_UNARY(BesselK1, Eigen::bessel_k1(a), 0.1, 20)
BENCH_SPECIAL_UNARY(BesselK0e, Eigen::bessel_k0e(a), 0.1, 100)
BENCH_SPECIAL_UNARY(BesselK1e, Eigen::bessel_k1e(a), 0.1, 100)

// --- Two-argument functions ---
BENCH_SPECIAL_BINARY(Igamma, Eigen::igamma(a, b), 0.1, 10, 0.1, 10)
BENCH_SPECIAL_BINARY(Igammac, Eigen::igammac(a, b), 0.1, 10, 0.1, 10)
BENCH_SPECIAL_BINARY(Zeta, Eigen::zeta(a, b), 1.1, 10, 0.1, 10)
BENCH_SPECIAL_BINARY(Polygamma, Eigen::polygamma(a, b), 1, 4, 0.1, 10)

// --- Ternary: betainc ---
template <typename Scalar>
static void BM_Betainc(benchmark::State& state) {
  const Index n = state.range(0);
  using Arr = Array<Scalar, Dynamic, 1>;
  Arr a = (Arr::Random(n) + Scalar(1)) * Scalar(2.5) + Scalar(0.5);  // [0.5, 5.5]
  Arr b = (Arr::Random(n) + Scalar(1)) * Scalar(2.5) + Scalar(0.5);
  Arr x = (Arr::Random(n) + Scalar(1)) * Scalar(0.5);  // [0, 1]
  Arr result(n);
  for (auto _ : state) {
    result = Eigen::betainc(a, b, x);
    benchmark::DoNotOptimize(result.data());
  }
  state.counters["Elements/s"] = benchmark::Counter(n, benchmark::Counter::kIsIterationInvariantRate);
  state.SetBytesProcessed(state.iterations() * n * sizeof(Scalar) * 4);
}

#define SPECIAL_SIZES ->Arg(256)->Arg(4096)->Arg(65536)->Arg(1048576)

// --- Register float ---
BENCHMARK(BM_Lgamma<float>) SPECIAL_SIZES->Name("Lgamma_float");
BENCHMARK(BM_Digamma<float>) SPECIAL_SIZES->Name("Digamma_float");
BENCHMARK(BM_BesselI0<float>) SPECIAL_SIZES->Name("BesselI0_float");
BENCHMARK(BM_BesselI1<float>) SPECIAL_SIZES->Name("BesselI1_float");
BENCHMARK(BM_BesselI0e<float>) SPECIAL_SIZES->Name("BesselI0e_float");
BENCHMARK(BM_BesselI1e<float>) SPECIAL_SIZES->Name("BesselI1e_float");
BENCHMARK(BM_BesselJ0<float>) SPECIAL_SIZES->Name("BesselJ0_float");
BENCHMARK(BM_BesselJ1<float>) SPECIAL_SIZES->Name("BesselJ1_float");
BENCHMARK(BM_BesselY0<float>) SPECIAL_SIZES->Name("BesselY0_float");
BENCHMARK(BM_BesselY1<float>) SPECIAL_SIZES->Name("BesselY1_float");
BENCHMARK(BM_BesselK0<float>) SPECIAL_SIZES->Name("BesselK0_float");
BENCHMARK(BM_BesselK1<float>) SPECIAL_SIZES->Name("BesselK1_float");
BENCHMARK(BM_BesselK0e<float>) SPECIAL_SIZES->Name("BesselK0e_float");
BENCHMARK(BM_BesselK1e<float>) SPECIAL_SIZES->Name("BesselK1e_float");
BENCHMARK(BM_Igamma<float>) SPECIAL_SIZES->Name("Igamma_float");
BENCHMARK(BM_Igammac<float>) SPECIAL_SIZES->Name("Igammac_float");
BENCHMARK(BM_Betainc<float>) SPECIAL_SIZES->Name("Betainc_float");
BENCHMARK(BM_Zeta<float>) SPECIAL_SIZES->Name("Zeta_float");
BENCHMARK(BM_Polygamma<float>) SPECIAL_SIZES->Name("Polygamma_float");

// --- Register double ---
BENCHMARK(BM_Lgamma<double>) SPECIAL_SIZES->Name("Lgamma_double");
BENCHMARK(BM_Digamma<double>) SPECIAL_SIZES->Name("Digamma_double");
BENCHMARK(BM_BesselI0<double>) SPECIAL_SIZES->Name("BesselI0_double");
BENCHMARK(BM_BesselI1<double>) SPECIAL_SIZES->Name("BesselI1_double");
BENCHMARK(BM_BesselI0e<double>) SPECIAL_SIZES->Name("BesselI0e_double");
BENCHMARK(BM_BesselI1e<double>) SPECIAL_SIZES->Name("BesselI1e_double");
BENCHMARK(BM_BesselJ0<double>) SPECIAL_SIZES->Name("BesselJ0_double");
BENCHMARK(BM_BesselJ1<double>) SPECIAL_SIZES->Name("BesselJ1_double");
BENCHMARK(BM_BesselY0<double>) SPECIAL_SIZES->Name("BesselY0_double");
BENCHMARK(BM_BesselY1<double>) SPECIAL_SIZES->Name("BesselY1_double");
BENCHMARK(BM_BesselK0<double>) SPECIAL_SIZES->Name("BesselK0_double");
BENCHMARK(BM_BesselK1<double>) SPECIAL_SIZES->Name("BesselK1_double");
BENCHMARK(BM_BesselK0e<double>) SPECIAL_SIZES->Name("BesselK0e_double");
BENCHMARK(BM_BesselK1e<double>) SPECIAL_SIZES->Name("BesselK1e_double");
BENCHMARK(BM_Igamma<double>) SPECIAL_SIZES->Name("Igamma_double");
BENCHMARK(BM_Igammac<double>) SPECIAL_SIZES->Name("Igammac_double");
BENCHMARK(BM_Betainc<double>) SPECIAL_SIZES->Name("Betainc_double");
BENCHMARK(BM_Zeta<double>) SPECIAL_SIZES->Name("Zeta_double");
BENCHMARK(BM_Polygamma<double>) SPECIAL_SIZES->Name("Polygamma_double");