#include "globals.h"

using namespace std;

void Vocab::load(string vocab_file) {

  ifstream voc(vocab_file.data());
  if (!voc)
    cerr << "* WARNING: couldn't read vocabulary file '" << vocab_file << "'\n";

  while (voc) {
    int index = 0;
    string word;
    int count;

    voc >> index >> word >> count;

    word_to_id[word] = index;
    id_to_word[index] = word;

    voc >> ws;
  }

  if (! id_to_word.count(-1))
    id_to_word[-1] = "<s>";

  if (! id_to_word.count(-2))
    id_to_word[-2] = "</s>";

  voc.close();
}

string Vocab::operator[](int id) {
  if (id_to_word.count(id))
    return id_to_word[id];

//   return "<unk>";
  return itos(id);
}

int Vocab::operator[](string str) {
  if (word_to_id.count(str))
    return word_to_id[str];
  
//   return word_to_id["<unk>"];
  return 0;
}
