text: #include "data/shortlist.h"
  #include "microsoft/shortlist/utils/ParameterTree.h"
  #include <queue>
  
  namespace marian {
  namespace data {
  
  // cast current void pointer to T pointer and move forward by num elements 
  template <typename T>
  const T* get(const void*& current, size_t num = 1) {
    const T* ptr = (const T*)current;
    current = (const T*)current + num;
    return ptr;
  }
  
  QuicksandShortlistGenerator::QuicksandShortlistGenerator(Ptr<Options> options,
                                                           Ptr<const Vocab> srcVocab,
                                                           Ptr<const Vocab> trgVocab,
                                                           size_t srcIdx,
                                                           size_t /*trgIdx*/,
                                                           bool /*shared*/)
      : options_(options),
        srcVocab_(srcVocab),
        trgVocab_(trgVocab),
        srcIdx_(srcIdx) {
    std::vector<std::string> vals = options_->get<std::vector<std::string>>("shortlist");
  
    ABORT_IF(vals.empty(), "No path to filter path given");
    std::string fname = vals[0];
  
    auto firstNum   = vals.size() > 1 ? std::stoi(vals[1]) : 0;
    auto bestNum    = vals.size() > 2 ? std::stoi(vals[2]) : 0;
    float threshold = vals.size() > 3 ? std::stof(vals[3]) : 0;
  
    if(firstNum != 0 || bestNum != 0 || threshold != 0) {
      LOG(warn, "You have provided additional parameters for the Quicksand shortlist, but they are ignored.");
    }
  
    mmap_ = mio::mmap_source(fname); // memory-map the binary file once
    const void* current = mmap_.data(); // pointer iterator over binary file
    
    // compare magic number in binary file to make sure we are reading the right thing
    const int32_t MAGIC_NUMBER = 1234567890;
    int32_t header_magic_number = *get<int32_t>(current);
    ABORT_IF(header_magic_number != MAGIC_NUMBER, "Trying to mmap Quicksand shortlist but encountered wrong magic number");
  
    auto config = ::quicksand::ParameterTree::FromBinaryReader(current);
    use16bit_ = config->GetBoolReq("use_16_bit");
    
    LOG(info, "[data] Mapping Quicksand shortlist from {}", fname);
  
    idSize_ = sizeof(int32_t);
    if (use16bit_) {
      idSize_ = sizeof(uint16_t);
    }
  
    // mmap the binary shortlist pieces
    numDefaultIds_        = *get<int32_t>(current);
    defaultIds_           =  get<int32_t>(current, numDefaultIds_);
    numSourceIds_         = *get<int32_t>(current);
    sourceLengths_        =  get<int32_t>(current, numSourceIds_);
    sourceOffsets_        =  get<int32_t>(current, numSourceIds_);
    numShortlistIds_      = *get<int32_t>(current);
    sourceToShortlistIds_ =  get<uint8_t>(current, idSize_ * numShortlistIds_);
    
    // display parameters
    LOG(info, 
        "[data] Quicksand shortlist has {} source ids, {} default ids and {} shortlist ids",
        numSourceIds_, 
        numDefaultIds_, 
        numShortlistIds_);
  }
  
  Ptr<Shortlist> QuicksandShortlistGenerator::generate(Ptr<data::CorpusBatch> batch) const {
    auto srcBatch = (*batch)[srcIdx_];
    auto maxShortlistSize = trgVocab_->size();
  
    std::unordered_set<int32_t> indexSet;
    for(int32_t i = 0; i < numDefaultIds_ && i < maxShortlistSize; ++i) {
      int32_t id = defaultIds_[i];
      indexSet.insert(id);
    }
  
    // State
    std::vector<std::pair<const uint8_t*, int32_t>> curShortlists(maxShortlistSize);
    auto curShortlistIt = curShortlists.begin();
  
    // Because we might fill up our shortlist before reaching max_shortlist_size, we fill the shortlist in order of rank.
    // E.g., first rank of word 0, first rank of word 1, ... second rank of word 0, ...
    int32_t maxLength = 0;
    for (Word word : srcBatch->data()) {
      int32_t sourceId = (int32_t)word.toWordIndex();
      srcVocab_->transcodeToShortlistInPlace((WordIndex*)&sourceId, 1);
  
      if (sourceId < numSourceIds_) { // if it's a valid source id
        const uint8_t* curShortlistIds = sourceToShortlistIds_ + idSize_ * sourceOffsets_[sourceId]; // start position for mapping
        int32_t length = sourceLengths_[sourceId]; // how many mappings are there
        curShortlistIt->first  = curShortlistIds;
        curShortlistIt->second = length;
        curShortlistIt++;
        
        if (length > maxLength)
          maxLength = length;
      }
    }
          
    // collect the actual shortlist mappings
    for (int32_t i = 0; i < maxLength && indexSet.size() < maxShortlistSize; i++) {
      for (int32_t j = 0; j < curShortlists.size() && indexSet.size() < maxShortlistSize; j++) {
        int32_t length = curShortlists[j].second;
        if (i < length) {
          const uint8_t* source_shortlist_ids_bytes = curShortlists[j].first;
          int32_t id = 0;
          if (use16bit_) {
            const uint16_t* source_shortlist_ids = reinterpret_cast<const uint16_t*>(source_shortlist_ids_bytes);
            id = (int32_t)source_shortlist_ids[i];
          }
          else {
            const int32_t* source_shortlist_ids = reinterpret_cast<const int32_t*>(source_shortlist_ids_bytes);
            id = source_shortlist_ids[i];
          }
          indexSet.insert(id);
        }
      }
    }
  
    // turn into vector and sort (selected indices)
    std::vector<WordIndex> indices;
    indices.reserve(indexSet.size());
    for(auto i : indexSet)
      indices.push_back((WordIndex)i);
  
    std::sort(indices.begin(), indices.end());
    return New<Shortlist>(indices);
  }
  
  Ptr<ShortlistGenerator> createShortlistGenerator(Ptr<Options> options,
                                                   Ptr<const Vocab> srcVocab,
                                                   Ptr<const Vocab> trgVocab,
                                                   size_t srcIdx,
                                                   size_t trgIdx,
                                                   bool shared) {
    std::vector<std::string> vals = options->get<std::vector<std::string>>("shortlist");
    ABORT_IF(vals.empty(), "No path to shortlist given");
    std::string fname = vals[0];
    if(isBinaryShortlist(fname)){
        return New<BinaryShortlistGenerator>(options, srcVocab, trgVocab, srcIdx, trgIdx, shared);
    } else if(filesystem::Path(fname).extension().string() == ".bin") {
      return New<QuicksandShortlistGenerator>(options, srcVocab, trgVocab, srcIdx, trgIdx, shared);
    } else {
      return New<LexicalShortlistGenerator>(options, srcVocab, trgVocab, srcIdx, trgIdx, shared);
    }
  }
  
  bool isBinaryShortlist(const std::string& fileName){
    uint64_t magic;
    io::InputFileStream in(fileName);
    in.read((char*)(&magic), sizeof(magic));
    return in && (magic == BINARY_SHORTLIST_MAGIC);
  }
  
  void BinaryShortlistGenerator::contentCheck() {
    bool failFlag = 0;
    // The offset table has to be within the size of shortlists.
    for(int i = 0; i < wordToOffsetSize_-1; i++)
      failFlag |= wordToOffset_[i] >= shortListsSize_;
  
    // The last element of wordToOffset_ must equal shortListsSize_
    failFlag |= wordToOffset_[wordToOffsetSize_-1] != shortListsSize_;
  
    // The vocabulary indices have to be within the vocabulary size.
    size_t vSize = trgVocab_->size();
    for(int j = 0; j < shortListsSize_; j++)
      failFlag |= shortLists_[j] >= vSize;
    ABORT_IF(failFlag, "Error: shortlist indices are out of bounds");
  }
  
  // load shortlist from buffer
  void BinaryShortlistGenerator::load(const void* ptr_void, size_t blobSize, bool check /*= true*/) {
    /* File layout:
     * header
     * wordToOffset array
     * shortLists array
     */
    ABORT_IF(blobSize < sizeof(Header), "Shortlist length {} too short to have a header", blobSize);
  
    const char *ptr = static_cast<const char*>(ptr_void);
    const Header &header = *reinterpret_cast<const Header*>(ptr);
    ptr += sizeof(Header);
    ABORT_IF(header.magic != BINARY_SHORTLIST_MAGIC, "Incorrect magic in binary shortlist");
  
    uint64_t expectedSize = sizeof(Header) + header.wordToOffsetSize * sizeof(uint64_t) + header.shortListsSize * sizeof(WordIndex);
    ABORT_IF(expectedSize != blobSize, "Shortlist header claims file size should be {} but file is {}", expectedSize, blobSize);
  
    if (check) {
      uint64_t checksumActual = util::hashMem<uint64_t, uint64_t>(&header.firstNum, (blobSize - sizeof(header.magic) - sizeof(header.checksum)) / sizeof(uint64_t));
      ABORT_IF(checksumActual != header.checksum, "checksum check failed: this binary shortlist is corrupted");
    }
  
    firstNum_ = header.firstNum;
    bestNum_ = header.bestNum;
    LOG(info, "[data] Lexical short list firstNum {} and bestNum {}", firstNum_, bestNum_);
  
    wordToOffsetSize_ = header.wordToOffsetSize;
    shortListsSize_ = header.shortListsSize;
  
    // Offsets right after header.
    wordToOffset_ = reinterpret_cast<const uint64_t*>(ptr);
    ptr += wordToOffsetSize_ * sizeof(uint64_t);
  
    shortLists_ = reinterpret_cast<const WordIndex*>(ptr);
  
    // Verify offsets and vocab ids are within bounds if requested by user.
    if(check)
      contentCheck();
  }
  
  // load shortlist from file
  void BinaryShortlistGenerator::load(const std::string& filename, bool check /*=true*/) {
    std::error_code error;
    mmapMem_.map(filename, error);
    ABORT_IF(error, "Error mapping file: {}", error.message());
    load(mmapMem_.data(), mmapMem_.mapped_length(), check);
  }
  
  BinaryShortlistGenerator::BinaryShortlistGenerator(Ptr<Options> options,
                                                     Ptr<const Vocab> srcVocab,
                                                     Ptr<const Vocab> trgVocab,
                                                     size_t srcIdx /*= 0*/,
                                                     size_t /*trgIdx = 1*/,
                                                     bool shared /*= false*/)
      : options_(options),
        srcVocab_(srcVocab),
        trgVocab_(trgVocab),
        srcIdx_(srcIdx),
        shared_(shared) {
  
    std::vector<std::string> vals = options_->get<std::vector<std::string>>("shortlist");
    ABORT_IF(vals.empty(), "No path to shortlist file given");
    std::string fname = vals[0];
  
    if(isBinaryShortlist(fname)){
      bool check = vals.size() > 1 ? std::stoi(vals[1]) : 1;
      LOG(info, "[data] Loading binary shortlist as {} {}", fname, check);
      load(fname, check);
    }
    else{
      firstNum_ = vals.size() > 1 ? std::stoi(vals[1]) : 100;
      bestNum_ = vals.size() > 2 ? std::stoi(vals[2]) : 100;
      float threshold = vals.size() > 3 ? std::stof(vals[3]) : 0;
      LOG(info, "[data] Importing text lexical shortlist as {} {} {} {}",
          fname, firstNum_, bestNum_, threshold);
      import(fname, threshold);
    }
  }
  
  BinaryShortlistGenerator::BinaryShortlistGenerator(const void *ptr_void,
                                                     const size_t blobSize,
                                                     Ptr<const Vocab> srcVocab,
                                                     Ptr<const Vocab> trgVocab,
                                                     size_t srcIdx /*= 0*/,
                                                     size_t /*trgIdx = 1*/,
                                                     bool shared /*= false*/,
                                                     bool check /*= true*/)
      : srcVocab_(srcVocab),
        trgVocab_(trgVocab),
        srcIdx_(srcIdx),
        shared_(shared) {
    load(ptr_void, blobSize, check);
  }
  
  Ptr<Shortlist> BinaryShortlistGenerator::generate(Ptr<data::CorpusBatch> batch) const {
    auto srcBatch = (*batch)[srcIdx_];
    size_t srcVocabSize = srcVocab_->size();
    size_t trgVocabSize = trgVocab_->size();
  
    // Since V=trgVocab_->size() is not large, anchor the time and space complexity to O(V).
    // Attempt to squeeze the truth tables into CPU cache
    std::vector<bool> srcTruthTable(srcVocabSize, 0);  // holds selected source words
    std::vector<bool> trgTruthTable(trgVocabSize, 0);  // holds selected target words
  
    // add firstNum most frequent words
    for(WordIndex i = 0; i < firstNum_ && i < trgVocabSize; ++i)
      trgTruthTable[i] = 1;
  
    // collect unique words from source
    // add aligned target words: mark trgTruthTable[word] to 1
    for(auto word : srcBatch->data()) {
      WordIndex srcIndex = word.toWordIndex();
      if(shared_)
        trgTruthTable[srcIndex] = 1;
      // If srcIndex has not been encountered, add the corresponding target words
      if (!srcTruthTable[srcIndex]) {
        for (uint64_t j = wordToOffset_[srcIndex]; j < wordToOffset_[srcIndex+1]; j++)
          trgTruthTable[shortLists_[j]] = 1;
        srcTruthTable[srcIndex] = 1;
      }
    }
  
    // Due to the 'multiple-of-eight' issue, the following O(N) patch is inserted
    size_t trgTruthTableOnes = 0;   // counter for no. of selected target words
    for (size_t i = 0; i < trgVocabSize; i++) {
      if(trgTruthTable[i])
        trgTruthTableOnes++;
    }
  
    // Ensure that the generated vocabulary items from a shortlist are a multiple-of-eight
    // This is necessary until intgemm supports non-multiple-of-eight matrices.
    for (size_t i = firstNum_; i < trgVocabSize && trgTruthTableOnes%8!=0; i++){
      if (!trgTruthTable[i]){
        trgTruthTable[i] = 1;
        trgTruthTableOnes++;
      }
    }
  
    // turn selected indices into vector and sort (Bucket sort: O(V))
    std::vector<WordIndex> indices;
    for (WordIndex i = 0; i < trgVocabSize; i++) {
      if(trgTruthTable[i])
        indices.push_back(i);
    }
  
    return New<Shortlist>(indices);
  }
  
  void BinaryShortlistGenerator::dump(const std::string& fileName) const {
    ABORT_IF(mmapMem_.is_open(),"No need to dump again");
    LOG(info, "[data] Saving binary shortlist dump to {}", fileName);
    saveBlobToFile(fileName);
  }
  
  void BinaryShortlistGenerator::import(const std::string& filename, double threshold) {
    io::InputFileStream in(filename);
    std::string src, trg;
  
    // Read text file
    std::vector<std::unordered_map<WordIndex, float>> srcTgtProbTable(srcVocab_->size());
    float prob;
  
    while(in >> trg >> src >> prob) {
      if(src == "NULL" || trg == "NULL")
        continue;
  
      auto sId = (*srcVocab_)[src].toWordIndex();
      auto tId = (*trgVocab_)[trg].toWordIndex();
  
      if(srcTgtProbTable[sId][tId] < prob)
        srcTgtProbTable[sId][tId] = prob;
    }
  
    // Create priority queue and count
    std::vector<std::priority_queue<std::pair<float, WordIndex>>> vpq;
    uint64_t shortListsSize = 0;
  
    vpq.resize(srcTgtProbTable.size());
    for(WordIndex sId = 0; sId < srcTgtProbTable.size(); sId++) {
      uint64_t shortListsSizeCurrent = 0;
      for(auto entry : srcTgtProbTable[sId]) {
        if (entry.first>=threshold) {
          vpq[sId].push(std::make_pair(entry.second, entry.first));
          if(shortListsSizeCurrent < bestNum_)
            shortListsSizeCurrent++;
        }
      }
      shortListsSize += shortListsSizeCurrent;
    }
  
    wordToOffsetSize_ = vpq.size() + 1;
    shortListsSize_ = shortListsSize;
  
    // Generate a binary blob
    blob_.resize(sizeof(Header) + wordToOffsetSize_ * sizeof(uint64_t) + shortListsSize_ * sizeof(WordIndex));
    struct Header* pHeader = (struct Header *)blob_.data();
    pHeader->magic = BINARY_SHORTLIST_MAGIC;
    pHeader->firstNum = firstNum_;
    pHeader->bestNum = bestNum_;
    pHeader->wordToOffsetSize = wordToOffsetSize_;
    pHeader->shortListsSize = shortListsSize_;
    uint64_t* wordToOffset = (uint64_t*)((char *)pHeader + sizeof(Header));
    WordIndex* shortLists = (WordIndex*)((char*)wordToOffset + wordToOffsetSize_*sizeof(uint64_t));
  
    uint64_t shortlistIdx = 0;
    for (size_t i = 0; i < wordToOffsetSize_ - 1; i++) {
      wordToOffset[i] = shortlistIdx;
      for(int popcnt = 0; popcnt < bestNum_ && !vpq[i].empty(); popcnt++) {
        shortLists[shortlistIdx] = vpq[i].top().second;
        shortlistIdx++;
        vpq[i].pop();
      }
    }
    wordToOffset[wordToOffsetSize_-1] = shortlistIdx;
  
    // Sort word indices for each shortlist
    for(int i = 1; i < wordToOffsetSize_; i++) {
      std::sort(&shortLists[wordToOffset[i-1]], &shortLists[wordToOffset[i]]);
    }
    pHeader->checksum = (uint64_t)util::hashMem<uint64_t>((uint64_t *)blob_.data()+2,
                                                          blob_.size()/sizeof(uint64_t)-2);
  
    wordToOffset_ = wordToOffset;
    shortLists_ = shortLists;
  }
  
  void BinaryShortlistGenerator::saveBlobToFile(const std::string& fileName) const {
    io::OutputFileStream outTop(fileName);
    outTop.write(blob_.data(), blob_.size());
  }
  
  }  // namespace data
  }  // namespace marian
  