# Path: data/common.py
# WORD_EMB_PATH = '/path/to/word/embedding/'
# 
# Path: data/bookcorpus.py
# VOCAB_PATH = os.path.join(BOOKCORPUS_DATA_DIR, 'word_count{}')


from gensim.models import KeyedVectors
from data.common import WORD_EMB_PATH
from data.bookcorpus import VOCAB_PATH
import collections
import os
import tqdm
import numpy as np
import sklearn.linear_model
import tensorflow as tf
  num_batches = num_words // batch_size + (num_words % batch_size != 0)

  for batch_idx in tqdm.trange(num_batches):
    ws = all_words[batch_idx * batch_size: (batch_idx + 1) * batch_size]
    w_embs = model.predict(word2vec[ws])
    for w, w_emb in zip(ws, w_embs):
      combined_emb[w] = w_emb

  for w in vocab:
    combined_emb[w] = thought_emb[vocab[w]]

  print("Created expanded vocabulary of %d words" % len(combined_emb))
  return combined_emb


def load_bookcorpus_vocab():
  vocab = dict()
  with open(VOCAB_PATH) as f:
    i = 1
    for line in f:
      vocab[line.split("\t")[0]] = i
      i += 1
  return vocab


def load_pretrained_word_embedding(glove=False):
  print("loading pretrained embedding from disk...")
  word_emb_model = KeyedVectors.load_word2vec_format(
      GLOVE_EMBEDDING_PATH if glove else W2V_EMBEDDING_PATH,
      binary=not glove, limit=1500000)

# Path: data/common.py
# WORD_EMB_PATH = '/path/to/word/embedding/'
# 
# Path: data/bookcorpus.py
# VOCAB_PATH = os.path.join(BOOKCORPUS_DATA_DIR, 'word_count{}')


# Path: thought/vocabulary_expansion.py
from gensim.models import KeyedVectors
from data.common import WORD_EMB_PATH
from data.bookcorpus import VOCAB_PATH
import collections
import os
import tqdm
import numpy as np
import sklearn.linear_model
import tensorflow as tf

  num_batches = num_words // batch_size + (num_words % batch_size != 0)

  for batch_idx in tqdm.trange(num_batches):
    ws = all_words[batch_idx * batch_size: (batch_idx + 1) * batch_size]
    w_embs = model.predict(word2vec[ws])
    for w, w_emb in zip(ws, w_embs):
      combined_emb[w] = w_emb

  for w in vocab:
    combined_emb[w] = thought_emb[vocab[w]]

  print("Created expanded vocabulary of %d words" % len(combined_emb))
  return combined_emb


def load_bookcorpus_vocab():
  vocab = dict()
  with open(VOCAB_PATH) as f:
    i = 1
    for line in f:
      vocab[line.split("\t")[0]] = i
      i += 1
  return vocab


def load_pretrained_word_embedding(glove=False):
  print("loading pretrained embedding from disk...")
  word_emb_model = KeyedVectors.load_word2vec_format(
      GLOVE_EMBEDDING_PATH if glove else W2V_EMBEDDING_PATH,
      binary=not glove, limit=1500000)

  return word_emb_model