'''
Compute and save DocumentVectors for training and testing sets.
'''

import numpy as np
import pandas as pd
import os
from gensim.models import Word2Vec as w2v

DATA_DIR = 'DATA'
DATAPATH = os.path.join(DATA_DIR)


def docvec_avg(documents, wordvecs):
    '''
    Create document vector using the average of word vectors.

    :param documents: pandas dataframe containing words from the documents
    :param wordvecs: precomputed word2vec model
    :return: pandas dataframe containing document vectors
    '''
    df = pd.DataFrame(index=documents.index, columns=['DocVec', 'label'])
    for i in list(documents.index):
        print('[] Computing Document Vector (Avg) for {} []'.format(i))
        wvinit = np.zeros((10000,))
        for word in documents.loc[i]['words']:
            wvinit += wordvecs.wv[str(word)]
        dv_avg = wvinit / len(documents.loc[i]['words'])
        df.loc[i]['DocVec'] = dv_avg
        df.loc[i]['label'] = documents.loc[i]['label']

    return df


def main():
    train_ds = pd.read_pickle(DATAPATH + '/train_ds_1500.pickle',
                              compression='gzip')
    wordvecs = w2v.load(DATAPATH + '/total_model_10000f_w5_cbowmean_hs_iter50000')

    train_ds_dv = docvec_avg(train_ds, wordvecs)

    train_ds_dv.to_pickle(DATAPATH + '/train_docvec_avg_1500_d10000.pickle',
                          compression='gzip', protocol=4)

    test_ds = pd.read_pickle(DATAPATH + '/test_ds_1900.pickle', compression='gzip')
    test_ds_dv = docvec_avg(test_ds, wordvecs)
    test_ds_dv.to_pickle(DATAPATH + '/test_docvec_avg_1900_d10000.pickle',
                         compression='gzip', protocol=4)


if __name__ == '__main__':
    main()
