#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;

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.nColor = 3;
  param.x[0] = geom[0];
  param.x[1] = geom[1];
  param.x[2] = geom[2];
  param.x[3] = geom[3];
  param.t_boundary = QUDA_PERIODIC_T;
  param.siteSubset = QUDA_FULL_SITE_SUBSET; // no even/odd, just a full lattice
  param.link_type = QUDA_SU3_LINKS;
  param.setPrecision(QUDA_DOUBLE_PRECISION);

  param.create = QUDA_NULL_FIELD_CREATE;     // do not (zero-) initilize the fields
  param.location = QUDA_CUDA_FIELD_LOCATION; // field should live on the accelerator

  // turn off advanced features we dont care about for this benchmark
  param.reconstruct = QUDA_RECONSTRUCT_NO;
  param.ghostExchange = QUDA_GHOST_EXCHANGE_NO;

  // these control the physical data layout. Might be interesting to try out different
  // settings
  param.order = QUDA_FLOAT2_GAUGE_ORDER;
  param.geometry = QUDA_SCALAR_GEOMETRY;

  // create the field and fill with random SU(3) matrices
  auto U = cudaGaugeField(param);
  quda::RNG rng(U, /*seed=*/1234);
  gaugeGauss(U, rng, 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};

  auto U = make_gauge_field(geom);
  auto src = make_source(geom);

  // 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
  for (int iter = 0; iter < 20; ++iter)
    dirac.M(tmp, src);

  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.1f\n", gflops);
}

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

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

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

  for (int L : {8, 16, 24, 32})
    benchmark(L, 1000);

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