import re
import json
import pickle
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer

import config

# 公式生成
def formula_generate(fstr, formula_list):
    # 对'l'和'1'进行特殊处理
    fstr = fstr.replace('l0', '10') if 'l0' in fstr else fstr
    # 排除"1)"、"10安"
    if re.search(r"[+\-·⋅\*/=<>]", fstr) is None or \
    re.search(r"[不的为是由选在约到和或与时连接于做天表日其]", fstr) or \
    re.search(r"^[\d\-].*?[\u4e00-\u9fff]", fstr) or \
    re.search(r"^\d+\)", fstr) or re.search(r"^[a-z]\)", fstr) or fstr[-1] == '-':
        return '', formula_list
    # "()"内容特殊处理
    bstr_list = re.findall(r"\(.*?\)", fstr)
    for bstr in bstr_list:
        if len(re.findall(r"[\u4e00-\u9fff]", bstr[1:-1])) > 2:
            fstr = fstr.replace(bstr, '')
    # 公式判断处理
    # '.'单独处理
    fstr = fstr[:-1] if fstr[-1] == '.' else fstr
    r_fstr = re.sub(r"[\.·⋅x\*/\)\'\^]", '', fstr)
    if len(r_fstr) > 2:
        # 排除10N/kg, 10km/h之类公式
        if re.sub(r"^\d+[A-Za-z°%μΩθδεπαφ∅∮βωλγνυηωс∁]+\d*$", '', r_fstr) == '':
            pass
        # 排除数字公式
        elif r_fstr.isdigit() or re.sub(r"^-", '', fstr).isdigit():
            pass
        else:
            # 判断左右括号数量进行特殊处理
            lb, rb = 0, 0
            for c in fstr:
                if c == '(': lb += 1
                elif c == ')': rb += 1
            if rb > lb and fstr[-1] == ')':
                fstr = fstr[:-1]
            elif rb > lb and fstr[-1] != ')':
                fstr = '(' + fstr         
            formula_list.append(fstr.lower())
    
    return '', formula_list

# 公式识别
def formula_recognize(formula_string):
    formula_list = []
    fstr = ''
    brackets_flag = False
    formula_length = len(formula_string)
    for i,char in enumerate(formula_string):
        # 获取公式首字符
        if fstr == '' and re.search(r"[A-Za-z\dρμθδεαφ∅∮βωλγυνΩηωс∁]", char):
            fstr += char
        elif fstr != '':
            # 对括号进行特殊处理
            if brackets_flag == False and char == '(':
                bracket_char = fstr[-1]
                brackets_flag = True
            elif brackets_flag == True and char == ')':
                brackets_flag = False
                # "()"内容特殊处理
                bstr_list = re.findall(r"\(.*?\)", fstr+')')
                if len(bstr_list) > 0 and bstr_list[-1][1:-1].isdigit() and \
                   bracket_char.isalpha() is False:
                    fstr = fstr.replace(bstr_list[-1][:-1], '')
                    fstr, formula_list = formula_generate(fstr, formula_list)
                    continue
            # 括号要特殊处理
            if brackets_flag == True:
                fstr += char
                continue
            # 对公式尾字符进行判断
            if i == formula_length - 1:
                if re.search(r"[^\.,;、!?:#\'\"_]", char):
                    fstr += char
                fstr, formula_list = formula_generate(fstr, formula_list)
            # 若当前字符为标点符号, 则判定公式识别结束
            elif re.search(r"[,;、!?:#]", char):
                fstr, formula_list = formula_generate(fstr, formula_list)
            # 对'.'进行特殊处理, 若'.'之前或之后非数字, 则判定公式识别结束
            elif char == '.' and re.sub(r"^\d\.\d$", '', formula_string[i-1:i+2]) != '':
                fstr, formula_list = formula_generate(fstr, formula_list)
            # 若当前字符为汉字后续字符为数字或当前字符为数字后续字符为汉字, 则判定公式识别结束
            elif re.search(r"^([\u4e00-\u9fff]\d|\d[\u4e00-\u9fff])", formula_string[i:i+2]):
                fstr, formula_list = formula_generate(fstr, formula_list)
            # 若当前字符为加减乘除等于符号, 后一个字符为汉字或标点符号, 则判定公式识别结束
            elif re.search(r"[\+\-·⋅x\*/=<>≈∽∼]", char) and \
                 re.search(r"[\u4e00-\u9fff\.,;、!?:#\'\"]", formula_string[i+1:i+2]):
                fstr, formula_list = formula_generate(fstr, formula_list)
            # 若当前字符为汉字后续两个字符为非加减乘除等于符号, 则判定公式识别结束
            # re.sub(r"^[\u4e00-\u9fff][\u4e00-\u9fff\.,;、!?:#\'\"]*", '', formula_string[i:i+3]) == ''
            elif re.sub(r"^[\u4e00-\u9fff][^A-Za-z\+\-·⋅x\*/=<>≈∽∼]*", '', formula_string[i:i+3]) == '':
                if re.search(r"[甲乙丙丁有总额放吸电动浮排液物实]", char):
                    fstr += char
                fstr, formula_list = formula_generate(fstr, formula_list)
            else:
                fstr += char

    return list(set(formula_list))

# mongodb公式处理初始化
def formula_init():
    # 获取mongodb数据
    origin_dataset = config.mongo_coll.find(no_cursor_timeout=True, batch_size=5)

    # 将公式组合成字典形式{"f=ma": [123, 456]}
    formula_dict = dict()
    for data in origin_dataset:
        if "content_clear" not in data:
            continue
        formula_list = formula_recognize(data["content_clear"])
        if len(formula_list) == 0:
            continue
        for formula in formula_list:
            if formula not in formula_dict:
                formula_dict[formula] = [data["id"]]
            else:
                formula_dict[formula].append(data["id"])
    # 将公式字典转化为列表形式用于通过索引获取元素
    formula_id_list = [[formula, id_list] for formula,id_list in formula_dict.items()]

    # 词袋模型初始化
    bow_model = CountVectorizer(token_pattern=r"[^ ]")
    all_formula_list = [ele[0] for ele in formula_id_list]
    bow_vector = bow_model.fit_transform(all_formula_list)
    # 保存词袋模型数据
    with open(config.bow_model_path, "wb") as bm:
        pickle.dump(bow_model, bm)
    np.save(config.bow_vector_path, bow_vector.toarray().astype("float32"))
    with open(config.formula_data_path, 'w', encoding='utf8', errors='ignore') as f:
        json.dump(formula_id_list, f, ensure_ascii=False)

if __name__ == "__main__":
    # text = '已知c水=4.2×103J/(kg·℃),求-10'
    # text = "水的密度：ρ＝1.0×10kg/m3是"
    # text = "'金宜高速'是经过河池市政府所在地的第一条高速公路,王华一家从金城江到刘三姐故里宜州自驾旅行,单程共消耗汽油5*kg.小车总重为1.5*10^4*N,静止在水平地面上时轮子与地面接触的总面积为0.15*m^2,(汽油的热值q=4.6*10^7J/kg).求:(1)小车静止时对水平地面的压强;(2)5*kg汽油完全燃烧放出的热量,"
    # text = "p蜡=0.9*10^3Kq/m^3"
    # text = "在一个案件中,公安人员在海滩案发现场发现了罪犯留下的清晰的双脚站立脚印,立即用蜡浇灌了一只鞋模.测量鞋模的平均厚度为3*cm,质量675*g,又经测试达到脚印同样深度的压强为1.5*10^4*Pa,请你帮助公安人员计算出罪犯的体重为多少?(ρ蜡=0.9*10^3*kg/m^3)"
    # for i in range(1):
    #     print(formula_recognize(text))
    
    # mongodb公式处理初始化
    formula_init()