#define VKFFT_BACKEND 0  // vulkan

#include "fftw3.h"
#include "vkFFT.h"

#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include <chrono>
#include <vector>

using std::chrono::high_resolution_clock;
using std::chrono::duration_cast;
using std::chrono::duration;
using std::chrono::milliseconds;

#define CHECK(expr) do { \
  auto err = (expr); \
  if (err != 0) { \
    fprintf(stderr, "%s failed: %d\n", #expr, (int) err); \
    exit(1); \
  } \
} while (0)

#define CHECK_TRUE(expr) do { \
  auto ret = (expr); \
  if (!ret) { \
    fprintf(stderr, "%s failed: %d\n", #expr, (int) ret); \
    exit(1); \
  } \
} while (0)

void fftw_test() {
  printf("=== FFTW test ===\n");
  int N = 4096;
  auto *in  = (fftwf_complex*) fftwf_malloc(sizeof(fftwf_complex) * N);
  auto *out = (fftwf_complex*) fftwf_malloc(sizeof(fftwf_complex) * N);
  fftwf_plan p = fftwf_plan_dft_1d(N, in, out, FFTW_FORWARD, FFTW_EXHAUSTIVE);

  const int BATCH_SIZE = 32 * 1024;
  const int WARMUP_S = 2;
  const int MEASURE_S = 5;
  auto t0 = high_resolution_clock::now();
  while (high_resolution_clock::now() < t0 + std::chrono::seconds(WARMUP_S)) {
    for (int i = 0; i < BATCH_SIZE; i++) {
      fftwf_execute(p);
    }
  }
  auto t1 = high_resolution_clock::now();
  int n_iter = 0;
  while (high_resolution_clock::now() < t0 + std::chrono::seconds(WARMUP_S + MEASURE_S)) {
    for (int i = 0; i < BATCH_SIZE; i++) {
      fftwf_execute(p);
    }
    n_iter += BATCH_SIZE;
  }
  auto t2 = high_resolution_clock::now();

  auto ns_int = duration_cast<std::chrono::nanoseconds>(t2 - t1);
  printf("%d ns\n", static_cast<int>(ns_int.count() / n_iter));

  fftwf_destroy_plan(p);
  fftwf_free(in);
  fftwf_free(out);
}

void vkfft_test() {
  printf("=== vkFFT test ===\n");

  VkApplicationInfo ai = {};
  ai.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
  ai.apiVersion = VK_API_VERSION_1_1;  // match vkFFT test suite

  VkInstanceCreateInfo ici = {};
  ici.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
  ici.pApplicationInfo = &ai;
#ifdef __APPLE__
  const char* ici_ext[] = {
    VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME,
  };
  ici.enabledExtensionCount = 1;
  ici.ppEnabledExtensionNames = ici_ext;
  ici.flags |= VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;
#endif  // __APPLE__

  VkInstance instance;
  CHECK(vkCreateInstance(&ici, nullptr, &instance));

  uint32_t n_dev = 0;
  CHECK(vkEnumeratePhysicalDevices(instance, &n_dev, nullptr));
  CHECK_TRUE(n_dev > 0);
  std::vector<VkPhysicalDevice> devices(n_dev);
  CHECK(vkEnumeratePhysicalDevices(instance, &n_dev, devices.data()));
  VkPhysicalDevice phys_dev = devices[0];

  uint32_t n_fam = 0;
  vkGetPhysicalDeviceQueueFamilyProperties(phys_dev, &n_fam, nullptr);
  std::vector<VkQueueFamilyProperties> families(n_fam);
  vkGetPhysicalDeviceQueueFamilyProperties(phys_dev, &n_fam, families.data());

  // Find a compute queue.
  uint32_t family;
  for (family = 0; family < n_fam; family++) {
    if (families[family].queueFlags & VK_QUEUE_COMPUTE_BIT) {
      break;
    }
  }
  CHECK_TRUE(family < n_fam);
  const uint64_t ts_bits = families[family].timestampValidBits;
  CHECK_TRUE(ts_bits > 0);

  float priority = 1;
  VkDeviceQueueCreateInfo qci = {};
  qci.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
  qci.queueFamilyIndex = family;
  qci.queueCount = 1;
  qci.pQueuePriorities = &priority;

  VkDeviceCreateInfo dci ={};
  dci.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
  dci.queueCreateInfoCount = 1;
  dci.pQueueCreateInfos = &qci;

  VkDevice device;
  CHECK(vkCreateDevice(phys_dev, &dci, /*VkAllocationCallbacks*/ nullptr, &device));

  VkQueue queue;
  vkGetDeviceQueue(device, family, 0, &queue);

  // Get GPU timestamp query
  VkQueryPoolCreateInfo qi = {};
  qi.sType = VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO;
  qi.queryType = VK_QUERY_TYPE_TIMESTAMP;
  qi.queryCount = 2;

  VkQueryPool queries;
  CHECK(vkCreateQueryPool(device, &qi, nullptr, &queries));

  VkPhysicalDeviceProperties phys_dev_props;
  vkGetPhysicalDeviceProperties(phys_dev, &phys_dev_props);

  VkCommandPoolCreateInfo pci = {};
  pci.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
  pci.queueFamilyIndex = family;
  pci.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;

  VkCommandPool pool;
  CHECK(vkCreateCommandPool(device, &pci, nullptr, &pool));

  VkFenceCreateInfo fci = {};
  fci.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;

  VkFence fence;
  CHECK(vkCreateFence(device, &fci, nullptr, &fence));

  // Length-4096 complex input. All zeros.
  const uint64_t FFT_SIZE = 4096;
  const uint64_t BATCH_SIZE = 1024 * 4;
  std::vector<float> data(2 * FFT_SIZE * BATCH_SIZE, 0);
  uint64_t buf_sz = data.size() * sizeof(data[0]);

  VkBufferCreateInfo bci = {};
  bci.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
  bci.size = buf_sz;
  bci.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
  bci.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

  VkBuffer buffer;
  CHECK(vkCreateBuffer(device, &bci, nullptr, &buffer));

  VkMemoryRequirements mem_req;
  vkGetBufferMemoryRequirements(device, buffer, &mem_req);

  VkPhysicalDeviceMemoryProperties mem_props;
  vkGetPhysicalDeviceMemoryProperties(phys_dev, &mem_props);

  const VkMemoryPropertyFlags wanted =
    VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT |
    VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
    VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
  uint32_t mem_type;
  for (mem_type = 0; mem_type < mem_props.memoryTypeCount; mem_type++) {
    if (!(mem_req.memoryTypeBits & (1u << mem_type))) {
      continue;
    }
    if ((mem_props.memoryTypes[mem_type].propertyFlags & wanted) != wanted) {
      continue;
    }
    break;
  }
  CHECK_TRUE(mem_type < mem_props.memoryTypeCount);

  VkMemoryAllocateInfo mai = {};
  mai.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
  mai.allocationSize = mem_req.size;
  mai.memoryTypeIndex = mem_type;

  VkDeviceMemory memory;
  CHECK(vkAllocateMemory(device, &mai, nullptr, &memory));
  CHECK(vkBindBufferMemory(device, buffer, memory, 0));

  void *mapped;
  CHECK(vkMapMemory(device, memory, 0, buf_sz, 0, &mapped));
  memcpy(mapped, data.data(), buf_sz);

  // Finally into vkFFT shit now...
  VkFFTConfiguration config = {};
  config.FFTdim = 1;
  config.size[0] = FFT_SIZE;
  config.numberBatches = BATCH_SIZE;
  config.device = &device;
  config.physicalDevice = &phys_dev;
  config.queue = &queue;
  config.commandPool = &pool;
  config.fence = &fence;
  config.buffer = &buffer;
  config.bufferSize = &buf_sz;

  VkFFTApplication app = {};
  CHECK(initializeVkFFT(&app, config));

  VkCommandBufferAllocateInfo cai = {};
  cai.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
  cai.commandPool = pool;
  cai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
  cai.commandBufferCount = 1;

  VkCommandBuffer cmd;
  CHECK(vkAllocateCommandBuffers(device, &cai, &cmd));

  VkCommandBufferBeginInfo begin = {};
  begin.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
  CHECK(vkBeginCommandBuffer(cmd, &begin));

  vkCmdResetQueryPool(cmd, queries, 0, 2);
  vkCmdWriteTimestamp(cmd, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, queries, 0);

  VkFFTLaunchParams launch = {};
  launch.commandBuffer = &cmd;
  CHECK(VkFFTAppend(&app, -1, &launch));

  vkCmdWriteTimestamp(cmd, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, queries, 1);

  VkMemoryBarrier barrier = {};
  barrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
  barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
  barrier.dstAccessMask = VK_ACCESS_HOST_READ_BIT;
  vkCmdPipelineBarrier(
      cmd,
      VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
      VK_PIPELINE_STAGE_HOST_BIT,
      0, 1, &barrier, 0, nullptr, 0, nullptr);

  CHECK(vkEndCommandBuffer(cmd));

  VkSubmitInfo submit = {};
  submit.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
  submit.commandBufferCount = 1;
  submit.pCommandBuffers = &cmd;
  CHECK(vkQueueSubmit(queue, 1, &submit, VK_NULL_HANDLE));
  CHECK(vkQueueWaitIdle(queue));

  uint64_t timestamps[2] = {};
  CHECK(vkGetQueryPoolResults(
        device, queries, 0, 2,
        sizeof(timestamps), timestamps, sizeof(timestamps[0]),
        VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT));

  uint64_t mask = ts_bits == 64
    ? UINT64_MAX : (UINT64_C(1) << ts_bits) - 1;
  uint64_t ticks = (timestamps[1] - timestamps[0]) & mask;
  double ns = double(ticks) * phys_dev_props.limits.timestampPeriod;

  memcpy(data.data(), mapped, buf_sz);

  // TODO process result
  printf("Runtime (GPU): %.0f ns\n", ns / BATCH_SIZE);

  deleteVkFFT(&app);
  vkUnmapMemory(device, memory);
  vkDestroyBuffer(device, buffer, nullptr);
  vkFreeMemory(device, memory, nullptr);
  vkDestroyFence(device, fence, nullptr);
  vkDestroyCommandPool(device, pool, nullptr);
  vkDestroyQueryPool(device, queries, nullptr);
  vkDestroyDevice(device, nullptr);
  vkDestroyInstance(instance, nullptr);
}

int main() {
  //fftw_test();
  vkfft_test();
}