#include <algorithm>
#include <cassert>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <vector>

constexpr size_t length = 2097152;
std::vector<unsigned long long> population_vector, sample_vector;

unsigned long long max_value = 0, min_value = 0xfffffffff;

inline void store_into_vector(unsigned long long value) {
  if (max_value < value) {
    max_value = value;
  }
  if (min_value > value) {
    min_value = value;
  }
  population_vector.push_back(value);
}

typedef unsigned long long key_type;

typedef unsigned long long *key_type_ptr;

constexpr long SAMPLE_LENGTH = 1024;

constexpr long TEST_SIZE = 1024;

__device__ key_type *SampleItem, *PopulationItem;
__device__ bool *cdf_result;
__device__ unsigned insert_value;
#ifdef TEST_BOUNDS
__device__ unsigned int index_max, index_min;
#endif

__device__ const key_type *cudaBinarySearch(key_type *start, key_type *end,
                                            const key_type val) {
  auto begin = start;
  key_type *last_known_point = nullptr;
  while (begin < end) {
    auto mid = (end - begin) / 2;
    auto mid_val = *(start + mid);
    if (val == mid_val) {
      return start + mid;
    } else if (val > mid_val) {
      begin = begin + mid + 1;
    } else {
      end = end - mid - 1;
    }
    last_known_point = begin;
  }
  return last_known_point;
}

__device__ double sample_cdf(double x) {
  auto it = cudaBinarySearch(SampleItem, SampleItem + SAMPLE_LENGTH + 2, x);
  if (it == SampleItem + SAMPLE_LENGTH) {
    return 1;
  }
  if (it == SampleItem) {
    return 0;
  }
  auto it_prev = it - 1;
  return (double(it_prev - SampleItem) +
          (x - (double)*it_prev) / (double)(*it - *it_prev)) /
         double(SAMPLE_LENGTH - 1);
}

__global__ void initStorage(unsigned scale, unsigned test_size) {
  cudaFree(cdf_result);
  cudaMalloc(&cdf_result, scale * test_size * sizeof(bool));
  memset(cdf_result, 0, scale * test_size * sizeof(bool));
  insert_value = 0;
  // printf("init storage\n");
#ifdef TEST_BOUNDS
  index_max = 0;
  index_min = 0x7fffffff;
#endif
}

__global__ void init(key_type *sample_item, key_type *population_item) {
  SampleItem = sample_item;
  PopulationItem = population_item;
  // cudaMalloc(&SampleItem, (SAMPLE_LENGTH + 2) * sizeof(unsigned long long));
  //  cudaMalloc(&Storage, TEST_SIZE * sizeof(key_type));
  cdf_result = nullptr;
}

__global__ void kernel(unsigned long step, const double slice_size,
                       const long split_size) {
  for (int i = 0; i < step; i++) {
    auto tid = step * gridDim.x * blockDim.x + blockIdx.x * blockDim.x +
               threadIdx.x + i;

    // printf("%d\n", tid);
    auto index = (int)(sample_cdf(PopulationItem[tid]) / slice_size);
    if (cdf_result[index]) {
      while (cdf_result[++index]) {
        assert(index < split_size);
      }
    }
    cdf_result[index] = true;
    // atomicAdd(&insert_value, 1);

#ifdef TEST_BOUNDS
    while (true) {
      unsigned tmp = index_min;
      if (tmp < tid) {
        break;
      }
      if (atomicCAS(&index_min, tmp, tid) == tmp) {
        break;
      }
    }
    while (true) {
      unsigned tmp = index_max;
      if (tmp > tid) {
        break;
      }
      if (atomicCAS(&index_max, tmp, tid) == tmp) {
        break;
      }
    }
#endif
  }
}

__global__ void print_function() {
  // printf("%u\n", insert_value);
#ifdef TEST_BOUNDS
  printf("%u %u\n", index_min, index_max);
#endif
}

int main(int argc, char const *argv[]) {
  size_t test_size;
  if (argc == 1) {
    test_size = 1048576;
  } else {
    try {
      test_size = std::stol(argv[1]);
    } catch (...) {
      return 1;
    }
  }

  FILE *file = fopen("normal_distribution.txt", "r");
  assert(file);
  for (long long i; fscanf(file, "%lld ", &i) != EOF; store_into_vector(i))
    ;
  fclose(file);

  assert(population_vector.size() == length);

  sample_vector = std::vector<unsigned long long>(
      population_vector.begin(), population_vector.begin() + SAMPLE_LENGTH);
  sample_vector.push_back(min_value);
  sample_vector.push_back(max_value);

  std::sort(sample_vector.begin(), sample_vector.end());

  key_type *cudaSample = nullptr, *cudaPopulation = nullptr;
  cudaMalloc(&cudaSample, sizeof(key_type) * (SAMPLE_LENGTH + 2));
  cudaMalloc(&cudaPopulation, sizeof(key_type) * TEST_SIZE);
  cudaMemcpy(cudaSample, &sample_vector[0],
             sizeof(key_type) * sample_vector.size(), cudaMemcpyHostToDevice);
  cudaMemcpy(cudaPopulation, &population_vector[SAMPLE_LENGTH],
             sizeof(key_type) * TEST_SIZE, cudaMemcpyHostToDevice);
  // memcpy(sample_heap, sample_vector.cbegin(),sizeof(key_type) *(SAMPLE_LENGTH
  // + 2));

  init<<<1, 1>>>(cudaSample, cudaPopulation);
  cudaDeviceSynchronize();

  dim3 grid_dim = 2, block_dim = 512;

  unsigned long step = test_size / (grid_dim.x * block_dim.x);
  printf("step: %lu\n", step);
  assert(!(test_size % (grid_dim.x * block_dim.x)));

  for (int scale = 2; scale <= 8; scale++) {
    const long split_size = test_size * scale;
    const auto slice_size = 1.0 / (double)split_size;
    printf("scale: %i ", scale);
    initStorage<<<1, 1>>>(scale, test_size);
    cudaDeviceSynchronize();
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start, nullptr);
    kernel<<<grid_dim, block_dim>>>(step, slice_size, split_size);
    cudaDeviceSynchronize();
    cudaEventRecord(stop, nullptr);
    cudaEventSynchronize(stop);
    float time;
    cudaEventElapsedTime(&time, start, stop);
    cudaEventDestroy(start);
    cudaEventDestroy(stop);

    printf("time: %lf\n", time);
    print_function<<<1, 1>>>();
    cudaDeviceSynchronize();
  }
  return 0;
}