#include <unordered_set>
#include <assert.h>
#include <stdio.h>
#include <errno.h>
#include <sys/user.h>
#include <sys/mman.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <vector>
#include <limits>
#include <thread>
#include <chrono>
#include <inttypes.h>
#include <set>

static std::set<volatile uint8_t *> page_set;

static size_t GB   (1024*1024*1024);
static size_t GBDEC(1000*1000*1000);

//#define SEQUENTIAL
//#define RANDOM

//#define  MAXPAGES ( 1.0*GBDEC/PAGE_SIZE)
#define  MAXPAGES (60.0*GBDEC/PAGE_SIZE)
#ifndef MAXPAGES
# define MAXPAGES std::numeric_limits<size_t>::max()
# define FREEPAGES (1.0*GBDEC/PAGE_SIZE)
#endif

static void stat(const char *msg) {
  printf("%s: %zu=%zuGBDEC\n",msg,page_set.size(),page_set.size()*PAGE_SIZE/GBDEC);
}

static void eatmem() {
  while (page_set.size()<MAXPAGES) {
    void *const page(mmap(nullptr/*addr*/,PAGE_SIZE,PROT_READ|PROT_WRITE,MAP_PRIVATE|MAP_ANONYMOUS|MAP_LOCKED|MAP_POPULATE,-1/*fd*/,0/*offset*/));
    if (page==MAP_FAILED) {
      if (errno==ENOMEM) {
	stat("ENOMEM");
	return;
      }
      if (errno==EAGAIN)
	printf("ulimit -l?\n");
      stat(strerror(errno));
      assert(0);
    }
    const auto pagecast(static_cast<volatile uint8_t *>(page));
    *pagecast=0x55;
    const auto successpair(page_set.insert(pagecast));
    assert(successpair.second);
  }
  stat("MAXPAGES");
}

static void freemem() {
#ifdef FREEPAGES
  for (size_t pageno=0;pageno<FREEPAGES;++pageno) {
    const auto it(page_set.begin());
    if (it==page_set.end()) {
      printf("freemem: Not enough memory\n");
      assert(0);
    }
    const auto page(const_cast<uint8_t *>(*it));
    int err;
    err=munmap(page,PAGE_SIZE);
    assert(!err);
  }
  stat("freemem");
#endif
}

class Dimm {
public:
  const uint64_t from,to_excl;
  Dimm(uint64_t from_,uint64_t to_excl_):from(from_),to_excl(to_excl_) {}
  bool contains(volatile uint8_t *const page) const {
    const uint64_t addr(reinterpret_cast<uint64_t>(page));
    return from<=addr&&addr<to_excl;
  }
};

static std::vector<Dimm> dimms;

static ssize_t dimmno_get_nocache(volatile uint8_t *const page) {
  for (auto it=dimms.cbegin();it!=dimms.cend();++it) {
    const auto &dimm(*it);
    if (dimm.contains(page)) {
#if defined SEQUENTIAL
      static size_t seq=0;
      ++seq;
      seq%=dimms.size();
//      return ((it-dimms.cbegin())&~3)+(seq%4);
      return seq;
#elif defined RANDOM
      return random()%dimms.size();
#else
      return it-dimms.cbegin();
//      return dimms.size()-1-(it-dimms.cbegin());
//      return ((it-dimms.cbegin())&~3)+(random()%4);
#endif
    }
  }
  //printf("Cannot find a DIMM for: %p\n",page);
  return -1;
}

static ssize_t dimmno_get(volatile uint8_t *const page) {
  static ssize_t cache_dimmno(0);
  if (cache_dimmno==-1||!dimms[cache_dimmno].contains(page))
    cache_dimmno=dimmno_get_nocache(page);
  return cache_dimmno;
}

static size_t dimmno_to_cpuno(size_t dimmno) {
  return 2*dimmno;
  //return dimmno*std::thread::hardware_concurrency()/dimms.size();
}

static std::vector<std::unordered_set<volatile uint8_t *>> dimmno_to_pages;

static void readmap() {
  std::vector<uint64_t> pagemapvec;
  volatile uint8_t *const page_set_front(*page_set. cbegin());
  volatile uint8_t *const page_set_back (*page_set.crbegin());
  { const int fd(open("/proc/self/pagemap",O_RDONLY));
    assert(fd!=-1);
    pagemapvec.resize((page_set_back-page_set_front)/PAGE_SIZE+1);
    const size_t count(pagemapvec.size()*sizeof(uint64_t));
    const ssize_t got(pread(fd,pagemapvec.data(),count,uintptr_t(page_set_front)/PAGE_SIZE*sizeof(uint64_t)));
    assert(size_t(got)==count);
    int err;
    err=close(fd);
    assert(!err);
  }
  dimmno_to_pages.resize(dimms.size());
  size_t bad_pfn_mask(0);
  size_t bad_dimmno(0);
  for (auto &it:page_set) {
    volatile uint8_t *const page(it);
    const uint64_t pagemapval(pagemapvec[(page-page_set_front)/PAGE_SIZE]);
    //printf("%p=0x%016lx\n",page,pagemapval);
    // vm/pagemap.txt
    const uint64_t pfn_mask((1ULL<<55)-1); // Bits 0-54  page frame number (PFN) if present
    if ((pagemapval&~pfn_mask)!=0x8600000000000000) {
      ++bad_pfn_mask;
      continue;
    }
    const uint64_t pfn(pagemapval&pfn_mask);
    const auto pfn_page(reinterpret_cast<volatile uint8_t *>(pfn*PAGE_SIZE));
    const ssize_t dimmno(dimmno_get(pfn_page));
    if (dimmno==-1) {
      ++bad_dimmno;
      continue;
    }
    auto successpair(dimmno_to_pages[dimmno].insert(page));
    assert(successpair.second);
  }
  printf("bad_pfn_mask: %zu\n",bad_pfn_mask);
  printf("bad_dimmno:   %zu\n",bad_dimmno  );
  for (size_t dimmno=0;dimmno<dimmno_to_pages.size();++dimmno) {
    const Dimm &dimm(dimms[dimmno]);
    const size_t cpuno(dimmno_to_cpuno(dimmno));
    printf("DIMM %2zu, CPU %2zu: 0x%10" PRIx64 "-0x%10" PRIx64 "=%2g-%2gGB: %7zu pages\n",dimmno,cpuno,
      dimm.from,dimm.to_excl,dimm.from/double(GB),dimm.to_excl/double(GB),
      dimmno_to_pages[dimmno].size());
  }
}

static void runthread(size_t dimmno) {
  const std::unordered_set<volatile uint8_t *> &pages(dimmno_to_pages[dimmno]);
  volatile size_t sum=0;
  for (volatile uint8_t *const page:pages)
    for (size_t offset=0;offset<PAGE_SIZE;offset+=64/8) // 64-bit FSB?
      sum+=page[offset];
}

static void runthreads() {
  const auto start(std::chrono::system_clock::now());
  for (size_t loops=0;loops<10;++loops) {
    std::vector<std::thread> threads;
    for (size_t dimmno=0;dimmno<dimmno_to_pages.size();++dimmno) {
      threads.push_back(std::thread(runthread,dimmno));
      std::thread &thread(threads.back());
      const size_t cpuno(dimmno_to_cpuno(dimmno));
      cpu_set_t *const cpuset(CPU_ALLOC(cpuno+1));
      assert(cpuset);
      size_t const cpuset_size(CPU_ALLOC_SIZE(cpuno+1));
      CPU_ZERO_S(cpuset_size,cpuset);
      CPU_SET_S(cpuno,cpuset_size,cpuset);
      auto native(thread.native_handle());
      int err;
      err=pthread_setaffinity_np(native,cpuset_size,cpuset);
      assert(!err);
      CPU_FREE(cpuset);
    }
    for (std::thread &thread:threads)
      thread.join();
  }
  const auto end  (std::chrono::system_clock::now());
  printf("%zu ms\n",std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count());
}

static uint64_t strtouint64(const std::string &str) {
  static_assert(sizeof(long)==sizeof(uint64_t),"!64-bit");
  char *end;
  const long l(strtol(str.c_str(),&end,0/*base*/));
  assert(l>=0);
  assert(*end==0);
  return l;
}

static void dmidecode() {
  FILE *const f(popen("dmidecode -q -t 20","r"));
  assert(f);
  ssize_t dimmno(-1);
  const auto uint64max(std::numeric_limits<uint64_t>::max());
  uint64_t dimm_start(uint64max),dimm_end(uint64max);
  char *linep(nullptr);
  size_t linen(0);
  for (;;) {
    errno=0;
    const ssize_t got(getline(&linep,&linen,f));
    if (got==-1) {
      assert_perror(errno);
      break;
    }
    assert(got>=0);
    assert(size_t(got)==strlen(linep));
    std::string line(linep);
    if (got==0)
      break;
    assert(line.back()=='\n');
    line.erase(line.length()-1);
    if (line=="Memory Device Mapped Address") {
      assert(dimmno==-1||dimm_end!=uint64max);
      ++dimmno;
      dimm_start=dimm_end=uint64max;
      continue;
    }
    { const std::string start_str("\tStarting Address: ");
      if (line.substr(0,start_str.length())==start_str) {
	assert(dimmno>=0);
	assert(dimm_start==uint64max);
	dimm_start=strtouint64(line.substr(start_str.length()));
	assert(dimm_start!=uint64max);
	continue;
      }
    }
    { const std::string end_str("\tEnding Address: ");
      if (line.substr(0,end_str.length())==end_str) {
	assert(dimmno>=0);
	assert(dimm_end==uint64max);
	dimm_end=strtouint64(line.substr(end_str.length()));
	assert(dimm_end!=uint64max);
	assert(dimm_start!=uint64max);
	if (dimm_start>=2*GB)
	  dimm_start+=2*GB;
	if (dimm_end>=2*GB)
	  dimm_end+=2*GB;
	dimms.push_back(Dimm(dimm_start,dimm_end));
	continue;
      }
    }
  }
  free(linep);
  int err;
  err=fclose(f);
  assert(!err);
}

int main() {
  setbuf(stdout,NULL);
  dmidecode();
  eatmem();
  freemem();
  readmap();
  runthreads();
}