import sys

import os

import pandas as pd
from gensim.models import KeyedVectors
from gensim.models import Word2Vec
from loguru import logger

curr_path = os.path.abspath(os.path.dirname(__file__))
project_root_path = curr_path[:curr_path.find("video_to_word") + len("video_to_word")]
sys.path.append(project_root_path)

from config.url_and_db import ai_word_product_engine

logger.remove()  # 删去 import logger之后自动产生的handler，不删除的话会出现重复输出的现象
logger.add("/data/pythonProject/video_to_word/logs/train_word2vec_model.{time:YYYY-MM-DD}.log",
           rotation="00:00",
           format="{time:YYYY-MM-DD HH:mm:ss,SSS} [{process}] [{thread}] {level} {file} {line} - {message}",
           level="INFO")


@logger.catch()
def train_word2vec_model():
    # 1 get corpus
    sql = f"select word_split from tb_asr_result where word_split is not null"
    doc_df = pd.read_sql(sql, ai_word_product_engine)
    doc_df['word_split_lst'] = doc_df['word_split'].apply(lambda x: [word for word in x.split(' ')])
    logger.info(f"获取到 {doc_df['word_split_lst'].shape[0]} 个脚本文件.")

    # 2 train the gensim word2vec model with our own corpus
    model = Word2Vec(doc_df['word_split_lst'].values, min_count=5, vector_size=50, workers=3, window=3, sg=1)

    # 3 store just the words and their trained embeddings
    word_vectors = model.wv
    # 定时任务按天执行，直接覆盖历史模型文件
    word_vectors.save("/data/pythonProject/video_to_word/models/word2vec.wordvectors")
    logger.info(f"word2vec model 训练完成!")


if __name__ == '__main__':
    # 0 train model
    train_word2vec_model()

    # 1 load back with memory-mapping = read-only,
    wv = KeyedVectors.load("/data/pythonProject/video_to_word/models/word2vec.wordvectors", mmap='r')

    # 2 get numpy vector of a word (for test)
    logger.info(f"the vector of  红包 is {wv['红包']}")
