import os
import time

from gensim.models import word2vec

base_dir = os.getcwd()
print("base_dir: " + base_dir)
model_dir = os.path.join(base_dir, "_models")
data_dir = os.path.join(base_dir, "_data")


class W2vSentences:
    def __init__(self, f_name):
        self.f_name = f_name

    def __iter__(self):
        with open(self.f_name, mode="r", encoding="utf-8", errors="ignore") as f:
            for line in f.readlines():
                yield line.split()


def w2v_train(f_name, model_output_name):
    input_file = os.path.join(data_dir, f_name)
    print("input: " + input_file)
    sentences = W2vSentences(input_file)
    w2v_model = word2vec.Word2Vec(sentences, min_count=2, workers=2, size=100, window=10)
    model_path = os.path.join(model_dir, model_output_name)
    print("output: " + model_path)
    w2v_model.save(model_path)


if __name__ == "__main__":
    begin_time = time.time()

    w2v_train("dispose_problem.txt", "w2v_model")

    # input_file = os.path.join(data_dir, "dispose_problem.txt")
    # for line in open(input_file, "r"):
    #     print(line)

    print("total cost: {:.2f}mins".format((time.time() - begin_time) / 60))
