// Benchmarks for Kronecker product (dense and sparse).
// SPDX-FileCopyrightText: The Eigen Authors
// SPDX-License-Identifier: MPL-2.0

#include <benchmark/benchmark.h>
#include <Eigen/Core>
#include <Eigen/Sparse>
#include <contrib/Eigen/KroneckerProduct>

using namespace Eigen;

typedef double Scalar;
typedef Matrix<Scalar, Dynamic, Dynamic> Mat;
typedef SparseMatrix<Scalar> SpMat;

// --- Dense Kronecker product ---
static void BM_KroneckerDense(benchmark::State& state) {
  int na = state.range(0);
  int nb = state.range(1);

  Mat A = Mat::Random(na, na);
  Mat B = Mat::Random(nb, nb);

  for (auto _ : state) {
    Mat C = kroneckerProduct(A, B).eval();
    benchmark::DoNotOptimize(C.data());
    benchmark::ClobberMemory();
  }
  int outSize = na * nb;
  state.counters["output_size"] = outSize;
}

// --- Sparse Kronecker product ---
static void BM_KroneckerSparse(benchmark::State& state) {
  int na = state.range(0);
  int nb = state.range(1);

  // Create sparse identity-like matrices with some fill.
  SpMat A(na, na);
  SpMat B(nb, nb);

  std::vector<Triplet<Scalar>> tripsA, tripsB;
  for (int i = 0; i < na; ++i) {
    tripsA.emplace_back(i, i, 2.0);
    if (i + 1 < na) {
      tripsA.emplace_back(i, i + 1, -1.0);
      tripsA.emplace_back(i + 1, i, -1.0);
    }
  }
  for (int i = 0; i < nb; ++i) {
    tripsB.emplace_back(i, i, 2.0);
    if (i + 1 < nb) {
      tripsB.emplace_back(i, i + 1, -1.0);
      tripsB.emplace_back(i + 1, i, -1.0);
    }
  }
  A.setFromTriplets(tripsA.begin(), tripsA.end());
  B.setFromTriplets(tripsB.begin(), tripsB.end());

  for (auto _ : state) {
    SpMat C = kroneckerProduct(A, B).eval();
    benchmark::DoNotOptimize(C.valuePtr());
    benchmark::ClobberMemory();
  }
  state.counters["output_size"] = na * nb;
}

#define KRONECKER_SIZES ->ArgsProduct({{4, 8, 16}, {4, 8, 16}})
#define KRONECKER_SPARSE_SIZES ->ArgsProduct({{16, 32, 64, 128}, {16, 32, 64, 128}})

BENCHMARK(BM_KroneckerDense) KRONECKER_SIZES;
BENCHMARK(BM_KroneckerSparse) KRONECKER_SPARSE_SIZES;