import multiprocessing
import os
import time

from gensim.models import Word2Vec, word2vec
from gm_rpcd.all import bind
from utils.es import es_scan

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")

model_output_name = "w2v_model"
model_path = os.path.join(model_dir, model_output_name)
WORD2VEC_MODEL = word2vec.Word2Vec.load(model_path)


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)


@bind("strategy_embedding/word_vector/word_similarity")
def word_similarity(word):
    return WORD2VEC_MODEL.wv.most_similar(word)


def get_user_portrait_projects(score_limit=5):
    """
    return: {
      '6231F098-9E72-448E-B8D2-19FCB9687005': ['鼻综合', '玻尿酸填充面部'],
      '862538030266882': ['吸脂瘦脸', '吸脂瘦全身']
    }
    """
    es_res = es_scan("device", {}, rw=None)
    count = 0
    res = {}
    for i in es_res:
        count += 1
        print(count)
        source = i["_source"]
        device_id = source.get("device_id", "")
        projects = [i["name"] for i in source.get("projects", []) if i["score"] >= score_limit]
        if projects:
            res[device_id] = projects
    return res


def projects_item2vec(score_limit=5):
    user_dict = get_user_portrait_projects(score_limit=score_limit)
    # TODO if not redis.get user_dict:
    projects = list(user_dict.values())
    model = Word2Vec(projects, hs=0, min_count=3, workers=multiprocessing.cpu_count(), iter=10)
    print(model)
    print(len(projects))
    for word in ["鼻综合", "吸脂瘦脸"]:
        print(model.wv.most_similar(word, topn=5))
    return model


def clicked_tractate_ids_item2vec():
    click_ids = []
    with open(os.path.join(data_dir, "click_tractate_ids.csv"), "r") as f:
        data = f.readlines()
        data = data[:100]
        for i in data:
            tmp = i.split("|")
            # device_id = tmp[0]
            ids = tmp[1].rstrip("\n").split(",")
            click_ids.append(ids)
    model = Word2Vec(click_ids, hs=0, min_count=3, workers=multiprocessing.cpu_count(), iter=10)
    print(model)
    print(len(click_ids))
    for id in ["396948", "395803", "394988"]:
        print(model.wv.most_similar(id, topn=5))
    return model


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

    # w2v_train("dispose_problem.txt", model_output_name)

    for i in ["双眼皮", "隆鼻"]:
        print(word_similarity(i))

    clicked_tractate_ids_item2vec()

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