#define ENABLE_SORT_TEST
#define SORT_FIRST_VERSION
#include "sortlib.cuh"

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

std::vector<unsigned long long> population_vector, sample_vector;

constexpr size_t SAMPLE_LENGTH = 1023;
constexpr size_t TEST_LENGTH = 32'768;
// constexpr size_t TEST_LENGTH = 4096;
constexpr size_t RESERVED_BLOCK = 10'000;

unsigned long long max_value = 0, min_value = 0xfffffffff;
typedef std::pair<int, int> dim_pair_type;

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);
}
// #define TEST_BOUNDS

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

__global__ void initSample(key_type *sample, key_type *population) {
  cudaSampleItem = sample;
  cudaPopulationItem = population;
  // cudaMalloc(&cdf_result, sizeof(bool) * TEST_LENGTH);
  cdf_result = nullptr;
#ifdef TEST_BOUNDS
  index_max = 0;
  index_min = 0x7fffffff;
#endif
}

__global__ void kernel(unsigned long step, const double slice_size,
                       const unsigned long split_size) {
  auto custom_sort = CustomSort(SAMPLE_LENGTH, sizeof(key_type) * 8);
  // printf("kernel1 step: %ld\n", step);
  for (int i = 0; i < step; i++) {
    auto tid = step * (blockIdx.x * blockDim.x + threadIdx.x) + i;

    // printf("%d\n", tid);
    auto index =
        (int)(custom_sort.cdf(cudaSampleItem, cudaPopulationItem[tid]) /
              slice_size) -
        1;
    // printf("%llu %d\n", cudaPopulationItem[tid], index);
    /*if (cdf_result[index]) {
      while (cdf_result[++index]) {
        if (index >= split_size) {
          //printf("escape: %llu %lu\n", cudaPopulationItem[tid], split_size);
        }
        assert(index < (split_size + RESERVED_BLOCK));
      }
    }*/
    cdf_result[index] = true;
  }
}

__global__ void kernel2_real_binary(unsigned long step, const double slice_size,
                                    const unsigned long split_size) {
  auto custom_sort = CustomSort(SAMPLE_LENGTH, sizeof(key_type) * 8);
  // printf("%d\n", custom_sort.MOVE_OFFSET);
  //  printf("kernel2\n");

  for (int i = 0; i < step; i++) {
    auto tid = step * (blockIdx.x * blockDim.x + threadIdx.x) + i;
    // printf("%lu ", tid);

    // assert(tid < TEST_LENGTH);
    auto index = (int)(custom_sort.sample_cdf(cudaSampleItem,
                                              cudaSampleItem + SAMPLE_LENGTH,
                                              cudaPopulationItem[tid]) /
                       slice_size);
    // printf("%llu %d\n", cudaPopulationItem[tid], index);
    /*if (cdf_result[index]) {
      while (cdf_result[++index]) {
        assert(index < split_size);
      }
    }*/
    cdf_result[index] = true;
#ifdef TEST_BOUNDS
    atomicAdd(&insert_value, 1);

    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 print2() {
  /*for (int i = 0; i< SAMPLE_LENGTH; i++) {
    printf("%llu ", cudaSampleItem[i]);
  }
  printf("\n");*/
  printf("%lu\n", CustomSort::fast_log(SAMPLE_LENGTH));
}

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

__global__ void initStorage(unsigned scale, unsigned test_size) {
  cudaFree(cdf_result);
  // printf("test size: %u\n", test_size);
  cudaMalloc(&cdf_result, (scale * test_size + RESERVED_BLOCK) * sizeof(bool));
  memset(cdf_result, 0, (scale * test_size + RESERVED_BLOCK) * sizeof(bool));
}

__global__ void freeStorage() {
  cudaFree(cudaSampleItem);
  cudaFree(cudaPopulationItem);
  cudaFree(cdf_result);
}

__global__ void initCustomSample() {
  // printf("init\n");
  key_type *tmp = nullptr;
  auto custom_sort = CustomSort(SAMPLE_LENGTH, 0);
  cudaMalloc(&tmp, sizeof(key_type) * SAMPLE_LENGTH);
  for (size_t i = 0; i < SAMPLE_LENGTH; i++) {
    tmp[i] = cudaSampleItem[custom_sort.calculate_index(i) - 1];
  }
  // printf("copy\n");
  memcpy(cudaSampleItem, tmp, sizeof(key_type) * SAMPLE_LENGTH);
  cudaFree(tmp);
  tmp = nullptr;
  // memset(cdf_result, 0, sizeof(key_type) * TEST_LENGTH);
  cudaFree(cdf_result);
  cdf_result = nullptr;
  // printf("finalize\n");
}

const dim_pair_type DIM_PAIR[] = {dim_pair_type(16, 64), dim_pair_type(32, 32)};

void run_kernel(size_t test_size, bool custom = false) {
  for (auto &pair : DIM_PAIR) {
    dim3 grid_dim = pair.first, block_dim = pair.second;
    unsigned long step = test_size / (grid_dim.x * block_dim.x);
    printf("%scurrent dim: %d %d %lu\n", custom ? "custom " : "", pair.first,
           pair.second, step);
    for (int scale = 2; scale <= 16; scale += 2) {
      const unsigned 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);
      if (custom)
        kernel<<<grid_dim, block_dim>>>(step, slice_size, split_size);
      else
        kernel2_real_binary<<<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);
      if (time == 0) {
        // printf("last error: %u\n", cudaGetLastError());
      }
      printf("time: %lf\n", time);

      print_function<<<1, 1>>>();
      cudaDeviceSynchronize();
    }
  }
}

int main() {

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

  printf("test size: %d\n", test_size);
  sample_vector = std::vector<key_type>(
      population_vector.begin(), population_vector.begin() + SAMPLE_LENGTH - 2);
  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;
  cudaMalloc(&cudaSample, sizeof(key_type) * SAMPLE_LENGTH);
  assert(SAMPLE_LENGTH == sample_vector.size());
  cudaMemcpy(cudaSample, sample_vector.data(), sizeof(key_type) * SAMPLE_LENGTH,
             cudaMemcpyHostToDevice);
  cudaMalloc(&cudaPopulation, sizeof(key_type) * test_size);
  cudaMemcpy(cudaPopulation, population_vector.data() + SAMPLE_LENGTH,
             sizeof(key_type) * test_size, cudaMemcpyHostToDevice);

  // auto custom_sort = CustomSort(SAMPLE_LENGTH, sizeof(long) * 8);
  // custom_sort.testCalculation();
  testCustomCalculation<<<1, 1>>>(SAMPLE_LENGTH);
  cudaDeviceSynchronize();

  initSample<<<1, 1>>>(cudaSample, cudaPopulation);
  cudaDeviceSynchronize();
  run_kernel(test_size);

  initCustomSample<<<1, 1>>>();
  cudaDeviceSynchronize();

  run_kernel(test_size, true);
  freeStorage<<<1, 1>>>();
  cudaDeviceSynchronize();
}

/*
 * Sample output:
test size: 32768
current dim: 16 64 32
scale: 2 time: 0.100352
scale: 4 time: 0.102176
scale: 6 time: 0.104448
scale: 8 time: 0.100352
scale: 10 time: 0.100512
scale: 12 time: 0.104384
scale: 14 time: 0.101152
scale: 16 time: 0.099232
current dim: 32 32 32
scale: 2 time: 0.099168
scale: 4 time: 0.101376
scale: 6 time: 0.106496
scale: 8 time: 0.099584
scale: 10 time: 0.111616
scale: 12 time: 0.106464
scale: 14 time: 0.100352
scale: 16 time: 0.108544
custom current dim: 16 64 32
scale: 2 time: 0.113664
scale: 4 time: 0.112512
scale: 6 time: 0.115712
scale: 8 time: 0.114720
scale: 10 time: 0.107232
scale: 12 time: 0.109568
scale: 14 time: 0.107328
scale: 16 time: 0.107520
custom current dim: 32 32 32
scale: 2 time: 0.106496
scale: 4 time: 0.106368
scale: 6 time: 0.108384
scale: 8 time: 0.108384
scale: 10 time: 0.109792
scale: 12 time: 0.112704
scale: 14 time: 0.108672
scale: 16 time: 0.122848

 */