// Build with clang++ -ltbb -pthread -g -O3 -Wall --std=c++17  -o sort Sort.cpp

#include <execution>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <iostream>
#include <random>
#include <thread>
#include <unordered_set>
#include <vector>
#include <parallel/algorithm>

using numberType = unsigned;
using numberVec = std::vector<numberType>;
constexpr unsigned partitionBits = 9;
constexpr unsigned partitionCount = 1 << partitionBits;
constexpr unsigned partitionShift = sizeof(numberType) * 8 - partitionBits;

numberVec partitionSort(numberVec::iterator begin, numberVec::iterator end) {
   // Step 1: Go over the data and build a histogram
   std::vector<unsigned> hist(partitionCount);
   for (auto it = begin; it != end; ++it) {
      numberType r = *it;
      unsigned bucket = r >> partitionShift;
      hist[bucket]++;
   }

   // Step 2: Use the histograms to calculate a prefix sum
   unsigned sum = 0;
   for (auto &h : hist) {
      auto tmp = h;
      h = sum;
      sum += tmp;
   }

   // Step 3: Partition the data into a new empty array
   numberVec keys(end - begin);
   for (auto it = begin; it != end; ++it) {
      numberType r = *it;
      unsigned bucket = r >> partitionShift;
      auto pos = hist[bucket];
      keys[pos] = r;
      hist[bucket]++;
   }

   // Step 4: Sort each partition in place with std::sort
   unsigned start = 0;
   for (unsigned i = 0; i < partitionCount; ++i) {
      std::sort(keys.begin() + start, keys.begin() + hist[i]);
      start = hist[i];
   }
   return keys;
}

// All information you need for the thread
struct threadInfo {
   numberVec::iterator begin, end;
   numberVec *out;
   std::vector<unsigned> hist;
};

// Build the histogram thread local (multi threaded)
void histLocal(threadInfo &p) {
   p.hist.resize(partitionCount);
   for (auto it = p.begin; it != p.end; ++it) {
      numberType r = *it;
      unsigned bucket = r >> partitionShift;
      p.hist[bucket]++;
   }
}

// exchange all histograms into a prefix sum (single threaded)
void exchangePrefixSum(std::vector<threadInfo> &ts) {
   unsigned sum = 0;
   unsigned hs = ts[0].hist.size();
   for (unsigned h = 0; h < hs; ++h) {
      for (unsigned t = 0; t < ts.size(); ++t) {
         auto tmp = ts[t].hist[h];
         ts[t].hist[h] = sum;
         sum += tmp;
      }
   }
}

// partition thread local unsynchronized (multi threaded)
void partitionLocal(threadInfo &i) {
   // Step 3 partition
   numberType *keys = i.out->data();
   for (auto it = i.begin; it != i.end; ++it) {
      numberType r = *it;
      unsigned bucket = r >> partitionShift;
      auto pos = i.hist[bucket];
      keys[pos] = r;
      i.hist[bucket]++;
   }
}

numberVec sortParallel(numberVec::iterator begin, numberVec::iterator end) {
   // Initialize threadInfo for all threads
   unsigned threadCount = 20;
   std::vector<threadInfo> info(threadCount);
   std::vector<std::thread> threads;
   unsigned partitionSize = (end - begin) / threadCount;
   // Run histLocal in parallel
   for (unsigned t = 0; t < threadCount; ++t) {
      auto &i = info[t];
      i.begin = begin + t * partitionSize;
      if (t + 1 == threadCount) {
         i.end = end;
      } else {
         i.end = begin + (t + 1) * partitionSize;
      }
      threads.emplace_back([&i]() { histLocal(i); });
   }
   for (auto &t : threads) {
      t.join();
   }
   threads.clear();
   // Generate Prefix sums single threaded
   exchangePrefixSum(info);
   // Partition multithreded
   numberVec output(end - begin);
   for (unsigned i = 0; i < threadCount; ++i) {
      info[i].out = &output;
      threads.emplace_back([i, &info]() { partitionLocal(info[i]); });
   }
   for (auto &t : threads) {
      t.join();
   }
   threads.clear();
   // Sort the partitions multi threaded with work stealing
   std::atomic<unsigned> pCount{0};
   auto sortPartition = [&]() {
      while (true) {
         unsigned p = pCount++;
         if (p >= partitionCount) {
            return;
         }
         auto &offsets = info[threadCount-1].hist;
         auto begin = p == 0 ? 0 : offsets[p-1];
         auto end = offsets[p];
         std::sort(output.begin() + begin, output.begin() + end);
      }
   };
   for (unsigned i = 0; i < threadCount; ++i) {
      threads.emplace_back(sortPartition);
   }
   for (auto &t : threads) {
      t.join();
   }
   return output;
}

numberVec tester() {
   // Random Number generator
   std::mt19937 gen(42);
   // Distribution
   std::uniform_int_distribution<numberType> uni(
       std::numeric_limits<numberType>::min(),
       std::numeric_limits<numberType>::max());
   // Fill vector
   numberVec keys(1e9);
   for (auto &k : keys) {
      k = uni(gen);
   }
   return keys;
}

bool checkSorted(numberVec keys) {
   auto c = keys[0];
   for (auto &k : keys) {
      if (c > k) return false;
      c = k;
   }
   return true;
}

int main(int argc, char *argv[]) {
   std::cout << "Generating Test Data" << std::endl;
   auto keys = tester();
   std::cout << "Running Test" << std::endl;
   auto start = std::chrono::high_resolution_clock::now();
   keys = sortParallel(keys.begin(), keys.end());
   // __gnu_parallel::sort(keys.begin(), keys.end());
   // std::sort(std::execution::par, keys.begin(), keys.end());
   auto end = std::chrono::high_resolution_clock::now();
   std::cout << "Time elapsed: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << " ms" << std::endl;
   // return 0 if sorted correct
   return !checkSorted(keys);
}

