# coding: utf-8

# source /srv/envs/SyntaxNet/bin/activate

import json
import re
import sys

import jieba
import numpy as np
from gensim import models as word2vec

from utils.cache import redis_client5
from utils.db import get_all_business_tags
from utils.es import es_scan

reload(sys)
sys.setdefaultencoding("utf8")


def get_new_user_tractate_info():
    q = {
        "query": {
            "bool": {
                "must": [{
                    "term": {
                        "is_online": "true"
                    }
                }, {
                    "terms": {
                        "content_level": [6, 5, 4, 3.5, 3]
                    }
                }, {
                    "terms": {
                        "operators_add_tags": [3315, 14288]
                    }
                }]
            }
        },
        "_source": {
            "includes": ["id", "portrait_tag_name", "tags_v3", "keynote_sentence"]
        }
    }
    return es_scan("tractate", q)


def get_tractate_vector(sentence_lst, all_keywords_set, model):
    word_list = []
    for s in sentence_lst:
        tmp_lst = []
        for i in jieba.lcut(s):
            s = i.encode("utf-8")
            if s in all_keywords_set:
                tmp_lst.append(s)
        word_list.append(tmp_lst)
    word_list = word_list[:15]
    count = 0
    res = {}
    for lst in word_list:
        vecs = []
        for name in lst:
            try:
                vecs.append(model.wv.get_vector(name))
            except Exception as e:
                pass
        if vecs:
            n = np.average(vecs, axis=0)
            res[count] = n.tolist()
        count += 1
    return res


def save_tractate_vector_to_redis(all_keywords_set, model):
    es_result = get_new_user_tractate_info()
    count = 0
    for i in es_result:
        count += 1
        source = i["_source"]
        sentences = source["keynote_sentence"]
        id = source["id"]
        vec = get_tractate_vector(sentences, all_keywords_set, model)
        redis_key = "rims:tractate:sentense:vector:" + str(id)
        if vec:
            print(count, id, len(vec))
            redis_client5.set(redis_key, json.dumps(vec))
            # TODO remove
            redis_client5.expire(redis_key, 60 * 60 * 24 * 3)


def save_tag_vector_to_redis(all_tags_lst, model):
    res = {}
    for name in all_tags_lst:
        try:
            vec = model.wv.get_vector(name).tolist()
            res[name] = json.dumps(vec)
        except Exception as e:
            pass
    redis_key = "rims:tractate:tags:vector"
    if res:
        redis_client5.hmset(redis_key, json.dumps(res))
        redis_client5.expire(redis_key, 60 * 60 * 24 * 3)
    print("tags vector redis: " + str(len(res)))


if __name__ == "__main__":
    model_path = "/data/log/word2vec/app/ipynb_garbage_files/test_w2v_model_4"
    keyword_txt = "/data/log/word2vec/all_key_word.txt"

    model = word2vec.Word2Vec.load(model_path)

    all_keywords_set = set([])
    with open(keyword_txt, "r") as f:
        words = f.readlines()
        for word in words:
            word = word.rstrip("\n")
            if word:
                all_keywords_set.add(word)
    print("keyword: " + str(len(all_keywords_set)))
    print(list(all_keywords_set)[:5])

    for word in all_keywords_set:
        jieba.add_word(word, freq=1000, tag="user_defined")

    save_tractate_vector_to_redis(all_keywords_set, model)

    all_tags_lst = [i.encode("utf-8") for i in get_all_business_tags()]
    print("all tags: " + str(len(all_tags_lst)))
    print(all_tags_lst[:5])
    save_tag_vector_to_redis(all_tags_lst, model)
