#include <iostream>
#include <cstdlib>
#include <atomic>
#include <sys/mman.h>
#include <tbb/tbb.h>
#include <unordered_map>

using namespace tbb;
using namespace std;

// Allocate zero-initialized huge pages
void* mmap_huge(size_t size) {
   void* p = mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
   madvise(p, size, MADV_HUGEPAGE);
   return p;
}

inline uint64_t hashKey(uint64_t k) {
   // Fibonacci
   return k*11400714819323198485ull;
}

inline uint64_t hashKey0(uint64_t k) {
   // MurmurHash64A
   const uint64_t m = 0xc6a4a7935bd1e995;
   const int r = 47;
   uint64_t h = 0x8445d61a4e774912 ^ (8*m);
   k *= m;
   k ^= k >> r;
   k *= m;
   h ^= k;
   h *= m;
   h ^= h >> r;
   h *= m;
   h ^= h >> r;
   return h|(1ull<<63);
}

static const double sizeFactor=3.3;

class Chaining {
   public:
   // Chained tuple entry
   struct Entry {
      uint64_t key;
      uint64_t value;
      Entry* next;
   };


   // TODO...
  
   // Constructor
   Chaining(uint64_t size) {
      // TODO...
   }

   // Destructor
   ~Chaining() {
      // TODO...
   }

   inline Entry* lookup(uint64_t key) {
      // TODO...
   }

   inline void insert(Entry* entry) {
      // TODO...
   }
};

class LinearProbing {
   public:
   // Entry
   struct Entry {
      uint64_t key;
      uint64_t value;
      std::atomic<bool> marker;
   };

   // TODO...

   // Constructor
   LinearProbing(uint64_t size) {
      // TODO...
   }

   // Destructor
   ~LinearProbing() {
      // TODO...
   }

   inline Entry* lookup(uint64_t key) {
      // TODO...
   }

  inline void insert(uint64_t key,uint64_t value) {
      // TODO...
   }
};


int main(int argc,char** argv) {
   uint64_t sizeR = atoi(argv[1]); // first command line argument: size of build
   uint64_t sizeS = atoi(argv[2]); // second command line argument: size of probe
   unsigned threadCount = atoi(argv[3]); // third command line argument: number of threads

   task_scheduler_init init(threadCount);

   // Init build-side relation R
   uint64_t* R=static_cast<uint64_t*>(mmap_huge(sizeR*sizeof(uint64_t)));
   parallel_for(blocked_range<size_t>(0, sizeR), [&](const blocked_range<size_t>& range) {
         unsigned int seed=range.begin();
         for (size_t i=range.begin(); i!=range.end(); ++i)
            R[i]=rand_r(&seed)%sizeR;
      });

   // Init probe-side relation S
   uint64_t* S=static_cast<uint64_t*>(mmap_huge(sizeS*sizeof(uint64_t)));
   parallel_for(blocked_range<size_t>(0, sizeS), [&](const blocked_range<size_t>& range) {
         unsigned int seed=range.begin();
         for (size_t i=range.begin(); i!=range.end(); ++i)
            S[i]=rand_r(&seed)%sizeR;
      });

   // Chaining
   {
      // Build hash table
      tick_count buildTS=tick_count::now();
      Chaining ht(sizeR);
      Chaining::Entry* entries=static_cast<Chaining::Entry*>(mmap_huge(sizeR*sizeof(Chaining::Entry))); //pre-allocate entries
      parallel_for(blocked_range<size_t>(0, sizeR), [&](const blocked_range<size_t>& range) {
            for (size_t i=range.begin(); i!=range.end(); ++i) {
               Chaining::Entry* entry=entries+i;
               entry->key=R[i];
               ht.insert(entry);
            }
         });
      tick_count probeTS=tick_count::now();
      cout << "CHAINING build:" << (sizeR/1e6)/(probeTS-buildTS).seconds() << "MT/s ";

      // Probe hash table and count number of hits
      std::atomic<uint64_t> hitCounter;
      hitCounter=0;
      parallel_for(blocked_range<size_t>(0, sizeS), [&](const blocked_range<size_t>& range) {
            uint64_t localHitCounter=0;
            for (size_t i=range.begin(); i!=range.end(); ++i)
               for (Chaining::Entry* entry=ht.lookup(S[i]); entry; entry=entry->next)
                  if (entry->key==S[i])
                     localHitCounter++;
            hitCounter+=localHitCounter;
         });
      tick_count stopTS=tick_count::now();
      cout << "probe: " << (sizeS/1e6)/(stopTS-probeTS).seconds() << "MT/s "
           << "total: " << ((sizeR+sizeS)/1e6)/(stopTS-buildTS).seconds() << "MT/s "
           << "count: " << hitCounter << endl;
      munmap(entries,sizeR*sizeof(Chaining::Entry));
   }

   // Linear Probing
   {
      // Build hash table
      tick_count buildTS=tick_count::now();
      LinearProbing ht(sizeR);
      parallel_for(blocked_range<size_t>(0, sizeR), [&](const blocked_range<size_t>& range) {
            for (size_t i=range.begin(); i!=range.end(); ++i)
               ht.insert(R[i]);
         });
      tick_count probeTS=tick_count::now();
      cout << "LINPROB  build:" << (sizeR/1e6)/(probeTS-buildTS).seconds() << "MT/s ";

      // Probe hash table and count number of hits
      std::atomic<uint64_t> hitCounter;
      hitCounter=0;
      parallel_for(blocked_range<size_t>(0, sizeS), [&](const blocked_range<size_t>& range) {
            uint64_t localHitCounter=0;
            for (size_t i=range.begin(); i!=range.end(); ++i) {
               uint64_t pos=hashKey(S[i])&ht.mask;
               while (ht.array[pos].marker) {
                  if (ht.array[pos].key==S[i])
                     localHitCounter++;
                  pos=(pos+1)&ht.mask;
               }
            }
            hitCounter+=localHitCounter;
         });
      tick_count stopTS=tick_count::now();
      cout << "probe: " << (sizeS/1e6)/(stopTS-probeTS).seconds() << "MT/s "
           << "total: " << ((sizeR+sizeS)/1e6)/(stopTS-buildTS).seconds() << "MT/s "
           << "count: " << hitCounter << endl;
   }

   // STL (for comparison)
   {
      // Build hash table (single threaded)
      tick_count buildTS=tick_count::now();
      unordered_multimap<uint64_t,uint64_t> ht(sizeR);
      for (uint64_t i=0; i<sizeR; i++)
         ht.emplace(R[i],0);
      tick_count probeTS=tick_count::now();
      cout << "STL      build:" << (sizeR/1e6)/(probeTS-buildTS).seconds() << "MT/s ";

      // Probe hash table and count number of hits
      std::atomic<uint64_t> hitCounter;
      hitCounter=0;
      parallel_for(blocked_range<size_t>(0, sizeS), [&](const blocked_range<size_t>& range) {
            uint64_t localHitCounter=0;
            for (size_t i=range.begin(); i!=range.end(); ++i) {
               auto range=ht.equal_range(S[i]);
               for (unordered_multimap<uint64_t,uint64_t>::iterator it=range.first; it!=range.second; ++it)
                  localHitCounter++;
            }
            hitCounter+=localHitCounter;
         });
      tick_count stopTS=tick_count::now();
      cout << "probe: " << (sizeS/1e6)/(stopTS-probeTS).seconds() << "MT/s "
           << "total: " << ((sizeR+sizeS)/1e6)/(stopTS-buildTS).seconds() << "MT/s "
           << "count: " << hitCounter << endl;
   }

   return 0;
}
