'''
Compute and save Word2Vec models.
'''

import pandas as pd
from gensim.models import word2vec as w2v
import os
import itertools as it

DATA_DIR = 'DATA'


def read_pickle(filename):
    '''Read pickle and return pandas DataFrame.

    :param filename: Pickle file
    :return: Pandas Dataframe
    '''

    df = pd.read_pickle(filename, compression='gzip')
    return df


def main():
    train_data = read_pickle(os.path.join(DATA_DIR) + '/total_data.pickle')

    # Getting Sentences
    print("Getting Sentences")

    sentences = it.chain([str(item) for item in sublist]
                         for sublist in train_data.loc[:]['words'])

    # Computing vectors
    print("Computing Vectors")
    total_model = w2v.Word2Vec(sentences=sentences, sg=1, cbow_mean=1, size=10000,
                               hs=1, min_count=1, workers=2, iter=50000,
                               compute_loss=True, window=5)
    print("Latest training loss: ", total_model.get_latest_training_loss())

    print("Saving model")
    total_model.save(os.path.join(DATA_DIR) +
                     '/total_model_10000f_w5_cbowmean_hs_iter50000')

    return 0


if __name__ == '__main__':
    import time
    start_time = time.time()
    main()
    print("Time Taken: ", str(time.time() - start_time), ' seconds')
