#include <algorithm>
#include <array>
#include <blas_quda.h>
#include <color_spinor_field.h>
#include <mpi.h>
// #include <quda_internal.h>
#include <memory>
#include <stdio.h>
#include <stdlib.h>

#include <cassert>
#include <dirac_quda.h>
#include <gauge_tools.h>

using namespace quda;

QudaPrecision smoother_halo_prec = QUDA_INVALID_PRECISION;

// This is the MPI grid, i.e. the layout of ranks, not the lattice volume
std::array<int, 4> gridsize = {1, 1, 1, 4};

void initComms(int argc, char **argv, std::array<int, 4> const &commDims)
{
  // init MPI communication
  MPI_Init(&argc, &argv);

  // this maps coordinates to rank number
  auto lex_rank_from_coords = [](int const *coords, void *)
  {
    int rank = coords[0];
    for (int i = 1; i < 4; i++)
      rank = gridsize[i] * rank + coords[i];
    return rank;
  };

  initCommsGridQuda(4, commDims.data(), lex_rank_from_coords, nullptr);

  for (int d = 0; d < 4; d++)
    if (gridsize[d] > 1)
      commDimPartitionedSet(d);
}

// creates a random gauge field
cudaGaugeField make_gauge_field(std::array<int, 4> const &geom)
{
  GaugeFieldParam param;

  // dimension and type of the lattice object
  param.nDim = 4;
  param.x[0] = geom[0];
  param.x[1] = geom[1];
  param.x[2] = geom[2];
  param.x[3] = geom[3];

  // number of colors. potentially confusingly, QUDA sometimes uses the word "color" to
  // things unrelated with physical color. things like "nColor=32" do pop up in deflation
  // solvers where it (to my understanding) refers to the number of (parallely processed)
  // deflation vectors.
  param.nColor = 3;

  // boundary conditions (dont really care for benchmark)
  param.t_boundary = QUDA_PERIODIC_T;

  // for this benchmark we only need "SINGLE" and/or "DOUBLE" precision. But smaller
  // precisions are available in QUDA too
  param.setPrecision(QUDA_DOUBLE_PRECISION);

  // no even/odd subset, we want a full lattice
  param.siteSubset = QUDA_FULL_SITE_SUBSET;

  // what kind of 3x3 matrices the field contains. A proper gauge field has SU(3)
  // matrices, but (for example) smeared/thick links could have non-unitary links.
  param.link_type = QUDA_SU3_LINKS;

  // "NULL" does not initialize the field upon creation, "ZERO" would set everything to 0
  param.create = QUDA_NULL_FIELD_CREATE;

  // field should be allocated directly on the accelerator/GPU
  param.location = QUDA_CUDA_FIELD_LOCATION;

  // "reconstruct" here means reconstructing a SU(3) matrix from fewer than 18 real
  // numbers (=3x3 complex numbers). Great feature in production (saving
  // memory/cache/network bandwidth), not used for this benchmark.
  param.reconstruct = QUDA_RECONSTRUCT_NO;

  // "ghostExchange" would often be called "halo exchange" outside of Quda. This has
  // nothing to do with ghost fields from continuum/perturbative qcd.
  param.ghostExchange = QUDA_GHOST_EXCHANGE_NO;

  // This controls the physical order of elements. "float2" is the the default
  param.order = QUDA_FLOAT2_GAUGE_ORDER;

  // this means the field is a LORENTZ vector (which a gauge field must be). Has nothing
  // to do with spin.
  param.geometry = QUDA_VECTOR_GEOMETRY;

  // create the field and fill with random SU(3) matrices
  // std::cout << param << std::endl; // double-check parameters
  auto U = cudaGaugeField(param);
  gaugeGauss(U, /*seed=*/1234, 1.0);
  return U;
}

// create a random source vector
ColorSpinorField make_source(std::array<int, 4> const &geom)
{
  ColorSpinorParam param;
  param.nColor = 3;
  param.nSpin = 4;
  param.nVec = 1; // only a single vector
  param.pad = 0;
  param.siteSubset = QUDA_FULL_SITE_SUBSET;
  param.nDim = 4;
  param.x[0] = geom[0];
  param.x[1] = geom[1];
  param.x[2] = geom[2];
  param.x[3] = geom[3];
  param.x[4] = 1; // no fifth dimension
  param.pc_type = QUDA_4D_PC;
  param.siteOrder = QUDA_EVEN_ODD_SITE_ORDER;
  param.gammaBasis = QUDA_DEGRAND_ROSSI_GAMMA_BASIS;
  param.create = QUDA_NULL_FIELD_CREATE; // do not (zero-) initilize the field
  param.setPrecision(QUDA_DOUBLE_PRECISION);
  param.location = QUDA_CUDA_FIELD_LOCATION;

  // create the field and fill it with random values
  auto src = ColorSpinorField(param);
  quda::RNG rng(src, 1234);
  spinorNoise(src, rng, QUDA_NOISE_GAUSS);
  /*printfQuda(
      "created src with norm = %f (sanity check: should be close to %f) and %f bytes\n",
      blas::norm2(src), 2.0 * 12 * geom[0] * geom[1] * geom[2] * geom[3],
      src.Bytes() * 1.0);*/
  // src.PrintDims();

  return src;
}

void benchmark(int L, int niter)
{
  std::array<int, 4> geom = {L, L, L, L};

  printfQuda("=======================  benchmarking L=%d =======================\n", L);

  auto U = make_gauge_field(geom);
  printfQuda("created random gauge field, %.3f GiB (sanity check: should be %.3f)\n",
             U.Bytes() / 1024. / 1024. / 1024.,
             1.0 * L * L * L * L * 4 * 18 * 8 / 1024. / 1024. / 1024.);
  auto src = make_source(geom);
  printfQuda("created random source, %.3f GiB (sanity check: should be %.3f)\n",
             src.Bytes() / 1024. / 1024. / 1024.,
             1.0 * L * L * L * L * 12 * 2 * 8 / 1024. / 1024. / 1024.);

  // create (Wilson) dirac operator
  DiracParam param;
  param.kappa = 0.10;
  param.dagger = QUDA_DAG_NO;
  param.matpcType = QUDA_MATPC_EVEN_EVEN;
  auto dirac = DiracWilson(param);

  // insert gauge field into the dirac operator
  // (the additional nullptr's are for smeared links and fancy preconditioners and such.
  // Not used for simple Wilson fermions)
  dirac.updateFields(&U, nullptr, nullptr, nullptr);

  auto tmp = ColorSpinorField(ColorSpinorParam(src));

  printfQuda("benchmarking Dirac operator. geom=(%d,%d,%d,%d), niter=%d\n", geom[0],
             geom[1], geom[2], geom[3], niter);

  // couple iterations without timing to warm up
  printfQuda("warmup...\n");
  for (int iter = 0; iter < 20; ++iter)
    dirac.M(tmp, src);

  printfQuda("running...\n");
  dirac.Flops(); // reset flops counter
  device_timer_t device_timer;
  device_timer.start();
  for (int iter = 0; iter < niter; ++iter)
    dirac.M(tmp, src);
  device_timer.stop();

  double secs = device_timer.last();
  double gflops = (dirac.Flops() * 1e-9) / secs;
  printfQuda("Gflops = %6.2f\n", gflops);
}

void benchmark_axpy(int L)
{
  printfQuda("================ axpy L=%d ==============\n", L);

  ColorSpinorParam param;
  param.nColor = 3;
  param.nSpin = 4;
  param.nVec = 1;
  param.pad = 0;
  param.siteSubset = QUDA_FULL_SITE_SUBSET;
  param.nDim = 4;
  param.x[0] = L;
  param.x[1] = L;
  param.x[2] = L;
  param.x[3] = L;
  param.x[4] = 1; // no fifth dimension
  param.pc_type = QUDA_4D_PC;
  param.siteOrder = QUDA_EVEN_ODD_SITE_ORDER;
  param.gammaBasis = QUDA_DEGRAND_ROSSI_GAMMA_BASIS;
  param.create = QUDA_NULL_FIELD_CREATE; // do not (zero-) initilize the field
  param.setPrecision(QUDA_DOUBLE_PRECISION);
  param.location = QUDA_CUDA_FIELD_LOCATION;

  // create the field and fill it with random values
  auto fieldA = ColorSpinorField(param);
  auto fieldB = ColorSpinorField(param);
  quda::RNG rng(fieldA, 1234);
  auto size_bytes = size_t(8) * 2 * param.x[0] * param.x[1] * param.x[2] * param.x[3] *
                    param.nColor * param.nSpin;
  assert(fieldA.Bytes() == size_bytes); // sanity check
  assert(fieldB.Bytes() == size_bytes); // sanity check
  spinorNoise(fieldA, rng, QUDA_NOISE_GAUSS);
  spinorNoise(fieldB, rng, QUDA_NOISE_GAUSS);

  // number of (real) elements in the field = number of fma instructions to do
  double flops_per_iter =
      2 * param.x[0] * param.x[1] * param.x[2] * param.x[3] * param.nColor * param.nSpin;

  int niter = 20;

  printfQuda("warmup...\n");
  for (int iter = 0; iter < 10; ++iter)
    blas::axpy(1.234, fieldA, fieldB);

  printfQuda("running...\n");
  device_timer_t device_timer;
  device_timer.start();
  for (int iter = 0; iter < niter; ++iter)
    blas::axpy(1.234, fieldA, fieldB); // fieldB += 1.234*fieldA
  device_timer.stop();

  double secs = device_timer.last();
  double gflops = (flops_per_iter * niter) * 1e-9 / secs;
  printfQuda("Gflops = %6.2f\n", gflops);
  printfQuda("bytes = %6.2f GiB\n", 3. * fieldA.Bytes() / 1024. / 1024. / 1024.);
  printfQuda("bandwidth = %6.2f GiB/s\n",
             fieldA.Bytes() * 3 / 1024. / 1024. / 1024. * niter / secs);
}

int main(int argc, char **argv)
{
  initComms(argc, argv, gridsize);

  initQuda(-1); // -1 for multi-gpu. otherwise this selects the device to be used

  //  verbosity options are:
  //  SILENT, SUMMARIZE, VERBOSE, DEBUG_VERBOSE
  setVerbosity(QUDA_VERBOSE);

  for (int L : {8, 12, 16, 24, 32})
    benchmark_axpy(L);
  for (int L : {16, 24, 32, 48, 64})
    benchmark(L, 100);

  endQuda();
  quda::comm_finalize();
  MPI_Finalize();
}