# Path: data/common.py
# MODEL_DIR = '/path/to/model/'
# 
# Path: data/wiki9.py
# def split_wiki9_articles(exp_id=0):
#   all_docs = list(os.listdir(WIKI9_DIR))
# 
#   s = gen_seed(exp_id)
#   random.seed(s)
#   random.shuffle(all_docs)
#   random.seed()
# 
#   n = len(all_docs) // 2
#   return all_docs[:n], all_docs[n:]
# 
# class WIKI9Articles(object):
#   def __init__(self, docs, dirname=WIKI9_DIR, verbose=0):
#     self.docs = docs
#     self.dirname = dirname
#     self.verbose = verbose
# 
#   def __iter__(self):
#     for fname in tqdm.tqdm(self.docs) if self.verbose else self.docs:
#       for line in smart_open.open(os.path.join(self.dirname, fname),
#                                   'r', encoding='utf-8'):
#         yield line.split()


from absl import flags
from absl import app
from gensim.models import Word2Vec, FastText
from glove import Corpus, Glove
from data.common import MODEL_DIR
from data.wiki9 import split_wiki9_articles, WIKI9Articles
import logging
import os


def train_word_embedding(exp_id=0, emb_model='w2v'):
  logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s',
                      level=logging.INFO)

  params = {
    'sg': 1,
    'negative': 25,
    'alpha': 0.05,
    'sample': 1e-4,
    'workers': 48,
    'iter': 5,
    'window': 5,
  }

  train_docs, test_docs = split_wiki9_articles(exp_id)
  print(len(train_docs), len(test_docs))
  wiki9_articles = WIKI9Articles(train_docs)
  save_dir = FLAGS.save_dir
  if not os.path.exists(save_dir):
    os.makedirs(save_dir)

  if emb_model == 'w2v':
    model = Word2Vec(wiki9_articles, **params)
  elif emb_model == 'ft':
    model = FastText(wiki9_articles, **params)
  elif emb_model == 'glove':
    model = train_glove(wiki9_articles, params, exp_id,
                        save_dir, save_dict=True)

# Path: data/common.py
# MODEL_DIR = '/path/to/model/'
# 
# Path: data/wiki9.py
# def split_wiki9_articles(exp_id=0):
#   all_docs = list(os.listdir(WIKI9_DIR))
# 
#   s = gen_seed(exp_id)
#   random.seed(s)
#   random.shuffle(all_docs)
#   random.seed()
# 
#   n = len(all_docs) // 2
#   return all_docs[:n], all_docs[n:]
# 
# class WIKI9Articles(object):
#   def __init__(self, docs, dirname=WIKI9_DIR, verbose=0):
#     self.docs = docs
#     self.dirname = dirname
#     self.verbose = verbose
# 
#   def __iter__(self):
#     for fname in tqdm.tqdm(self.docs) if self.verbose else self.docs:
#       for line in smart_open.open(os.path.join(self.dirname, fname),
#                                   'r', encoding='utf-8'):
#         yield line.split()


# Path: train_word_embedding.py
from absl import flags
from absl import app
from gensim.models import Word2Vec, FastText
from glove import Corpus, Glove
from data.common import MODEL_DIR
from data.wiki9 import split_wiki9_articles, WIKI9Articles
import logging
import os



def train_word_embedding(exp_id=0, emb_model='w2v'):
  logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s',
                      level=logging.INFO)

  params = {
    'sg': 1,
    'negative': 25,
    'alpha': 0.05,
    'sample': 1e-4,
    'workers': 48,
    'iter': 5,
    'window': 5,
  }

  train_docs, test_docs = split_wiki9_articles(exp_id)
  print(len(train_docs), len(test_docs))
  wiki9_articles = WIKI9Articles(train_docs)
  save_dir = FLAGS.save_dir
  if not os.path.exists(save_dir):
    os.makedirs(save_dir)

  if emb_model == 'w2v':
    model = Word2Vec(wiki9_articles, **params)
  elif emb_model == 'ft':
    model = FastText(wiki9_articles, **params)
  elif emb_model == 'glove':
    model = train_glove(wiki9_articles, params, exp_id,
                        save_dir, save_dict=True)

  else: