# coding: UTF-8
import time
import os
import torch
import numpy as np
from tqdm import tqdm
from train_eval import train, init_network
from importlib import import_module
import argparse
import torch.nn.functional  as F
from utils_fasttext import build_dataset, build_iterator, get_time_dif, build_pre_dataset
import jieba
import pickle as pkl
from trie_tree_match import Trie_tree
import random
import glob

parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str, default='FastText', help='')
parser.add_argument('--model_path', type=str, default='ckpt-v1/saved_dict/FastText.ckpt', help='')
parser.add_argument('--embedding', default='pre_trained', type=str, help='random or pre_trained')
parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')
parser.add_argument('--Mode', default="predict", type=str, help='train/predict')
parser.add_argument('--predict_path', default="/workspace/input_data_path/", type=str, help='使用绝对路径')
parser.add_argument('--predict_result_save_path', default="/workspace/output_data_path/", type=str, help='清洗结果保存路径，使用绝对路径')
parser.add_argument('--device_id', default="0", type=str, help='GPU ID')

args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id#'0'
def predict(config, model, predict_iter):
    predict_all = []
    probility_all = []
    with torch.no_grad():
        for texts, _ in predict_iter:
#             print("text:",len(texts))
            outputs = model(texts)
            probilitys = F.softmax(outputs)
            p_batch = torch.max(probilitys.data, 1)[0].cpu().numpy()
            predict_batch = torch.max(probilitys.data, 1)[1].cpu().numpy()
            predict_all.extend(predict_batch)
            probility_all.extend(p_batch)
#             print("*****batch probility length:",len(p_batch))
    return predict_all, probility_all

def prepare_vocab():
    vocab_dir = "敏感词库"
    text_names = [x for x in os.listdir(vocab_dir) if x.endswith(".txt")]
    vocabs_ = []
    for txt_name in text_names:
        with open(f"{vocab_dir}/{txt_name}",encoding="utf-8") as f:
            vocabs = f.read().splitlines()
            vocabs_.extend(vocabs)
    vocabs_=list(set(vocabs_))
    vocabs_.sort()
    with open("敏感词库/total.txt","w",encoding="utf-8") as w:
        w.write("\n".join(vocabs_))

def load_vocab(path = "敏感词库/色情词库.txt"):
    with open(path,encoding="utf-8") as f:
        vocabs = f.read().splitlines()
        return vocabs

def filter_by_trie(text,trie_model,threhold = 3):
    """
    0:dirty
    1:clean
    """
    res = trie_model.find_one(text)
    if len(res)>threhold:
        return 0
    elif len(res)>=1 and len(text)<200:
        return 0
    else:
        return 1


if __name__ == '__main__':
    # prepare_vocab()
    dirty_vocab = load_vocab("./敏感词库/色情词库.txt")
    trie_model  = Trie_tree()
    trie_model.load_vocab(dirty_vocab)
    # jieba.load_userdict(dirty_vocab)
    dataset = 'filter-v3-length200'  # 数据集
    model_name = args.model  #  FastText
    x = import_module('models.' + model_name)
    config = x.Config(dataset, embedding='random')

    config.save_path = args.model_path
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True  # 保证每次结果一样
    
    ## 设置预测路径
#     text_paths = [x for x in os.listdir(args.predict_path) if x.endswith(".txt")]
    text_paths=[]
    data_path_glob = args.predict_path+'part-*'
    text_paths.extend(list(glob.glob(data_path_glob)))

    #预测结果路径
    res_dir = args.predict_result_save_path
    if not os.path.exists(res_dir):
        os.mkdir(res_dir)
    

    total_path_num=0
    total_time=0
    
    #加载词表
    if os.path.exists(config.vocab_path):
        vocab = pkl.load(open(config.vocab_path, 'rb'))
    else:
        vocab = build_vocab(config.train_path, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)
        pkl.dump(vocab, open(config.vocab_path, 'wb'))
    
    # 加载模型
    config.n_vocab = len(vocab)
    model = x.Model(config).to(config.device)
    print(config.device)
    model.load_state_dict(torch.load(config.save_path))
    model.eval()

    random.shuffle(text_paths) 
    for i in range(len(text_paths)):
        text_path=text_paths[i]
        try:
            start_time=time.clock()
            total_path_num+=1

    #       dirty_file_name = text_path.replace(".txt", "_dirty.txt")
    #       dirty_file_path = f"{res_dir}/{dirty_file_name}"
            clean_file_path=args.predict_result_save_path +"train/"+ text_path.split('/')[-1]+'-clean.txt'
            dirty_file_path=args.predict_result_save_path + "dirty/"+text_path.split('/')[-1] + '-dirty.txt'


            if os.path.exists(clean_file_path):
                print("file has already cleaned:",text_path)
                continue
            if os.path.exists(dirty_file_path):
                continue

            config.predict_path = text_path #f"{args.predict_path}/{text_path}"
            ## 预测
            predict_data, all_text  = build_pre_dataset(config, args.word,vocab)
    #       print("-----predict_data length:",len(predict_data))
            if len(all_text)==0:
                continue

            predict_iter = build_iterator(predict_data, config)
            predict_all,probility_all=predict(config, model, predict_iter)

            label2str = {0: "【垃圾】", 1: "正常文本"}
            index_line = 0

            clean_file=open(clean_file_path,"w",encoding="utf-8")
            dirty_file=open(dirty_file_path,"w",encoding="utf-8")

            for index_line in range(len(all_text)):
                text=all_text[index_line]
                if text:
                    p=probility_all[index_line]
                    label = predict_all[index_line]

                    label_trie = filter_by_trie(text, trie_model,3)
                    if label == 0 and p>0.8 and text!="\n":
                        dirty_file.writelines("<<score:"+str(p)+">>\n"+text)
                        dirty_file.writelines("\n\n")
                    elif label_trie == 0:
                        dirty_file.writelines("by trie-tree match:"+text)
                        dirty_file.writelines("\n\n")
                    else:
                        clean_file.writelines(text)
    #                   clean_file.writelines("<<score:"+str(p)+">>\n"+text)
                        clean_file.writelines("\n\n")            

    #       print("-----------one file spend time:",time.clock()-start_time)
            total_time+=(time.clock()-start_time)
            if total_path_num%5==0:
                print("text path:",text_path)
                print("predict result lengith:",len(probility_all))
                print("file text lengith:",len(all_text))
                print("current file num:",total_path_num,"  all num:",len(text_paths))
                print("-----------one file average time:",total_time/total_path_num)
    #           print("-----------one sample spend time:",(time.clock()-start_time)/len(all_text),"samples:",len(all_text))
        except:
            print("the exception occured in this txt file:",text_path)
            continue
    print("-----------total time:",total_time)
    print("-----------total num:",total_path_num)
