# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Compute an expanded vocabulary of embeddings using a word2vec model.

This script loads the word embeddings from a trained skip-thoughts model and
from a trained word2vec model (typically with a larger vocabulary). It trains a
linear regression model without regularization to learn a linear mapping from
the word2vec embedding space to the skip-thoughts embedding space. The model is
then applied to all words in the word2vec vocabulary, yielding vectors in the
skip-thoughts word embedding space for the union of the two vocabularies.

The linear regression task is to learn a parameter matrix W to minimize
  || X - Y * W ||^2,
where X is a matrix of skip-thoughts embeddings of shape [num_words, dim1],
Y is a matrix of word2vec embeddings of shape [num_words, dim2], and W is a
matrix of shape [dim2, dim1].

This is based on the "Translation Matrix" method from the paper:

  "Exploiting Similarities among Languages for Machine Translation"
  Tomas Mikolov, Quoc V. Le, Ilya Sutskever
  https://arxiv.org/abs/1309.4168
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import os.path


import gensim.models
import numpy as np
import sklearn.linear_model
import tensorflow as tf

FLAGS = tf.flags.FLAGS

tf.flags.DEFINE_string("model", None,
                       "Checkpoint file or directory containing a checkpoint "
                       "file.")
tf.flags.DEFINE_string("model_vocab", None,
                       "Path to vocabulary file containing a list of newline-"
                       "separated words where the word id is the "
                       "corresponding 0-based index in the file.")
tf.flags.DEFINE_string("word2vec_model", None,
                       "File containing a word2vec model in binary format.")
tf.flags.DEFINE_string("output_dir", None, "Output directory.")
tf.flags.DEFINE_string("inout", "", "Input or output word embedding matrix")
tf.flags.DEFINE_boolean("transfer_embs", False, "Directly transfer embeddings")
tf.flags.DEFINE_string("transfer_embs_path", "", "Embeddings path for direct transfer")
tf.flags.DEFINE_string("word_emb_tensor", "word_embedding", "Name of word embedding tensor")
tf.flags.DEFINE_integer("vocab_size", 50001, "Vocabulary size")

tf.logging.set_verbosity(tf.logging.INFO)


def _load_skip_thoughts_embeddings(checkpoint_path):
  """Loads the embedding matrix from a skip-thoughts model checkpoint.

  Args:
    checkpoint_path: Model checkpoint file or directory containing a checkpoint
        file.

  Returns:
    word_embedding: A numpy array of shape [vocab_size, embedding_dim].

  Raises:
    ValueError: If no checkpoint file matches checkpoint_path.
  """
  if tf.gfile.IsDirectory(checkpoint_path):
    checkpoint_file = tf.train.latest_checkpoint(checkpoint_path)
    if not checkpoint_file:
      raise ValueError("No checkpoint file found in %s" % checkpoint_path)
  else:
    checkpoint_file = checkpoint_path

  tf.logging.info("Loading skip-thoughts embedding matrix from %s",
                  checkpoint_file)
  reader = tf.train.NewCheckpointReader(checkpoint_file)

  word_embedding = reader.get_tensor(FLAGS.word_emb_tensor + FLAGS.inout)
  tf.logging.info("Loaded skip-thoughts embedding matrix of shape %s",
                  word_embedding.shape)

  return word_embedding


def _load_vocabulary(filename):
  """Loads a vocabulary file.

  Args:
    filename: Path to text file containing newline-separated words.

  Returns:
    vocab: A dictionary mapping word to word id.
  """
  tf.logging.info("Reading vocabulary from %s", filename)
  vocab = collections.OrderedDict()
  with tf.gfile.GFile(filename, mode="r") as f:
    for i, line in enumerate(f):
      word = line.decode("utf-8").strip()
      assert word not in vocab, "Attempting to add word twice: %s" % word
      vocab[word] = i
  tf.logging.info("Read vocabulary of size %d", len(vocab))
  return vocab


def _expand_vocabulary(skip_thoughts_emb, model_vocab, word2vec):
  """Runs vocabulary expansion on a skip-thoughts model using a word2vec model.

  Args:
    skip_thoughts_emb: A numpy array of shape [model_vocab_size,
        skip_thoughts_embedding_dim].
    model_vocab: A dictionary of word to id.
    word2vec: An instance of gensim.models.Word2Vec.

  Returns:
    combined_emb: A dictionary mapping words to embedding vectors.
  """
  # Find words shared between the two vocabularies.
  tf.logging.info("Finding shared words")
  shared_words = [w for w in word2vec.vocab if w in model_vocab]

  # Select embedding vectors for shared words.
  tf.logging.info("Selecting embeddings for %d shared words", len(shared_words))
  shared_st_emb = skip_thoughts_emb[[
      model_vocab[w] for w in shared_words
  ]]
  shared_w2v_emb = word2vec[shared_words]

  # Train a linear regression model on the shared embedding vectors.
  tf.logging.info("Training linear regression model")
  model = sklearn.linear_model.LinearRegression()
  model.fit(shared_w2v_emb, shared_st_emb)

  # Create the expanded vocabulary.
  tf.logging.info("Creating embeddings for expanded vocabuary")
  combined_emb = collections.OrderedDict()
  for w in word2vec.vocab:
    # Ignore words with underscores (spaces).
    if "_" not in w:
      w_emb = model.predict(word2vec[w].reshape(1, -1))
      combined_emb[w] = w_emb.reshape(-1)

  for w in model_vocab:
    combined_emb[w] = skip_thoughts_emb[model_vocab[w]]

  tf.logging.info("Created expanded vocabulary of %d words", len(combined_emb))

  return combined_emb


def main(unused_argv):

  if not tf.gfile.IsDirectory(FLAGS.output_dir):
    tf.gfile.MakeDirs(FLAGS.output_dir)

  # Load the skip-thoughts embeddings and vocabulary.
  skip_thoughts_emb = _load_skip_thoughts_embeddings(FLAGS.model)

  if FLAGS.transfer_embs:
    unk_vec = skip_thoughts_emb[[0]]
    word_vecs = np.load(FLAGS.transfer_embs_path)
    embeddings = np.vstack((unk_vec, word_vecs))

    embeddings_file = os.path.join(FLAGS.output_dir, FLAGS.word_emb_tensor + FLAGS.inout + ".npy")
    np.save(embeddings_file, embeddings)
    tf.logging.info("Wrote embeddings file to %s", embeddings_file)

    with open(os.path.join(FLAGS.output_dir, FLAGS.word_emb_tensor + FLAGS.inout + ".txt"), 'w') as f:
      with open(FLAGS.model_vocab, 'r') as g:
        f.write(g.read().decode("utf-8"))

  else:
    model_vocab = _load_vocabulary(FLAGS.model_vocab)

    # Load the Word2Vec model.
    word2vec = gensim.models.KeyedVectors.load_word2vec_format(
        FLAGS.word2vec_model, binary=True)

    # Run vocabulary expansion.
    embedding_map = _expand_vocabulary(skip_thoughts_emb, model_vocab,
                                       word2vec)

    # Save the output.
    vocab = embedding_map.keys()
    vocab_file = os.path.join(FLAGS.output_dir, FLAGS.word_emb_tensor + FLAGS.inout + ".txt")
    with tf.gfile.GFile(vocab_file, "w") as f:
      f.write("\n".join(vocab))
    tf.logging.info("Wrote vocabulary file to %s", vocab_file)

    embeddings = np.array(embedding_map.values())
    embeddings_file = os.path.join(FLAGS.output_dir, FLAGS.word_emb_tensor + FLAGS.inout + ".npy")
    np.save(embeddings_file, embeddings)
    tf.logging.info("Wrote embeddings file to %s", embeddings_file)


if __name__ == "__main__":
  tf.app.run()
