from importlib import import_module
import torch
import pandas as pd
import pickle as pkl
from tqdm import tqdm
from numpy.linalg import norm
from utils import build_vocab, build_iterator
import numpy as np
import os
from sklearn.feature_extraction.text import CountVectorizer
import jieba
import joblib
from sklearn.model_selection import train_test_split
import fasttext
import gradio as gr


MAX_VOCAB_SIZE = 10000
UNK, PAD = '<UNK>', '<PAD>'
module = import_module("models.TextCNN")
dataset = 'THUCNews'
embedding = 'embedding_SougouNews.npz'  # 预训练的词向量表
config = module.Config(dataset, embedding)
model = module.Model(config).to(config.device)
model.load_state_dict(torch.load(config.save_path,map_location='cpu'))
model.eval()


def build_dataset(content, config, ues_word):
    if ues_word:
        tokenizer = lambda x: x.split(' ')  # 以空格隔开，word-level
    else:
        tokenizer = lambda x: [y for y in x]  # char-level
    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'))
    def load_dataset(pad_size=32):
        contents = []
        words_line = []
        token = tokenizer(content)
        seq_len = len(token)
        if pad_size:
            if len(token) < pad_size:
                token.extend([PAD] * (pad_size - len(token)))
            else:
                token = token[:pad_size]
                seq_len = pad_size
        for word in token:
            words_line.append(vocab.get(word, vocab.get(UNK)))
        contents.append((words_line, -1, seq_len))
        return contents  # [([...], 0), ([...], 1), ...]
    train = load_dataset(config.pad_size)
    return vocab, train

def wordvector(config,df,pad_size=32):
    wordvectorlist = []
    answerlist = []
    sentences = []
    department = []
    tokenizer = lambda x: [y for y in x]  # char-level
    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'))
    for row in df.index:
        content = df.loc[row]['ask']
        answer = df.loc[row]['answer']
        department.append(df.loc[row]['department'])
        answerlist.append(answer)
        sentences.append(content)
        words_line = []
        token = tokenizer(content)
        if pad_size:
            if len(token) < pad_size:
                token.extend([PAD] * (pad_size - len(token)))
            else:
                token = token[:pad_size]
                seq_len = pad_size
        for word in token:
            words_line.append(vocab.get(word, vocab.get(UNK)))
        vector = words_line
        wordvectorlist.append(vector)
    return wordvectorlist, answerlist
def cos_similarity_test(vector,word_vector_list):
    best_cosine = -10
    for i in range(len(word_vector_list)):
        best_index = 0
        B = word_vector_list[i]
        cosine = np.dot(vector, B) / (norm(vector) * norm(B))
        if cosine > best_cosine:
            best_cosine = cosine
            best_index = i
    return best_index

def predict_single_text(text,flag):
    # 预测单条文本的分类
    ft_model = None
    if flag == 4:
        ft_model = fasttext.load_model("./fasttext_model_erke.bin")
    if flag == 5:
        ft_model = fasttext.load_model("./fasttext_model_waike.bin")
    prediction = ft_model.predict([text], k=1)  # k=1表示只返回概率最高的标签
    if prediction:
        label = prediction[0][0][len('__label__'):]  # 获取标签
        str = prediction[0][0][0]
        label = str[9:-1]
        # print(str)
        confidence = prediction[1][0]  # 获取置信度
        return label, confidence
    else:
        return None, None
def proceed(content,NumOfChamber,content_vector):
    with open('./stop_words.txt', 'r', encoding="utf-8") as f:
        stop_words = list(l.strip() for l in f.readlines())
    stop_words.extend(['\n', '（', '）', ' '])  # 由于停用词中没有'\n'和中文的左右括号和空格，所以单独再加上去
    if NumOfChamber == 0:
        print("大类为：男科")
        file = pd.read_csv("./THUCNews/data/Data/nanke/nanke1.csv", encoding='utf-8')
        df = file[["department", "ask", "answer"]]
        word_vector_list, answer_list = wordvector(config, df)
        most_similar_index = cos_similarity_test(content_vector,word_vector_list)
        most_similar_answer = answer_list[most_similar_index]
        return None, most_similar_answer
    if NumOfChamber == 2:
        print("大类为：妇产科")
        file = pd.read_csv("./THUCNews/data/Data/fuchanke/fuchanke1.csv", encoding='utf-8')
        df = file[["department", "ask", "answer"]]
        word_vector_list, answer_list = wordvector(config, df)
        most_similar_index = cos_similarity_test(content_vector,word_vector_list)
        most_similar_answer = answer_list[most_similar_index]
        return None, most_similar_answer
    if NumOfChamber == 3:
        print("大类为：肿瘤科")
        file = pd.read_csv("./THUCNews/data/Data/zhongliuke/zhongliuke1.csv", encoding='utf-8')
        df = file[["department", "ask", "answer"]]
        word_vector_list, answer_list = wordvector(config, df)
        most_similar_index = cos_similarity_test(content_vector,word_vector_list)
        most_similar_answer = answer_list[most_similar_index]
        return None, most_similar_answer
    if NumOfChamber == 1:
        print("大类为：内科")
        data2 = pd.read_excel(io="data2_pro.xlsx", engine="openpyxl")
        department_counts = data2['department'].value_counts()
        underrepresented_categories = department_counts[department_counts < 20000].index
        data2 = data2[~data2['department'].isin(underrepresented_categories)]
        data2['Pre_Text'].fillna('', inplace=True)
        X = data2['Pre_Text']  # 文本数据
        y = data2['department']  # 标签
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        count_vectorizer = CountVectorizer(stop_words=stop_words, max_features=5000)
        X_train_count = count_vectorizer.fit_transform(X_train)
        X_test_count = count_vectorizer.transform(X_test)
        input_text = content
        nb_classifier = joblib.load('naive_bayes_model.pkl')
        input_text_processed = " ".join([w for w in list(jieba.cut(str(input_text))) if w not in stop_words])
        input_text_processed = count_vectorizer.transform([input_text_processed])
        predicted_department = nb_classifier.predict(input_text_processed)

        file = pd.read_csv("./THUCNews/data/Data/neike/neike1.csv", encoding='utf-8')
        df = file[file["department"] == predicted_department[0]]
        df = df[["department", "ask", "answer"]]
        word_vector_list, answer_list = wordvector(config, df)
        most_similar_index = cos_similarity_test(content_vector,word_vector_list)
        most_similar_answer = answer_list[most_similar_index]
        return predicted_department[0], most_similar_answer
    if NumOfChamber == 4:
        print("大类为：儿科")
        input_text = ' '.join([word for word in jieba.lcut(content) if word not in stop_words])
        predicted_department, confidence_score = predict_single_text(input_text,4)
        file = pd.read_csv("./THUCNews/data/Data/erke/erke1.csv", encoding='utf-8')
        df = file[file["department"] == predicted_department]
        df = df[["department", "ask", "answer"]]
        word_vector_list, answer_list = wordvector(config, df)
        most_similar_index = cos_similarity_test(content_vector,word_vector_list)
        most_similar_answer = answer_list[most_similar_index]
        return predicted_department, most_similar_answer
    if NumOfChamber == 5:
        print("大类为：外科")
        input_text = ' '.join([word for word in jieba.lcut(content) if word not in stop_words])
        predicted_department, confidence_score = predict_single_text(input_text,5)
        file = pd.read_csv("./THUCNews/data/Data/waike/waike1.csv", encoding='utf-8')
        print(type(predicted_department))
        print(predicted_department)
        df = file[file["department"] == predicted_department]
        df = df[["department", "ask", "answer"]]
        word_vector_list, answer_list = wordvector(config, df)
        most_similar_index = cos_similarity_test(content_vector,word_vector_list)
        most_similar_answer = answer_list[most_similar_index]
        return predicted_department, most_similar_answer



def working_func(content):
    # content = input("请输入您的病情描述：")
    vocab, train_data = build_dataset(content,config, False)
    content_vector = train_data[0][0]
    config.n_vocab = len(vocab)
    train_iter = build_iterator(train_data, config)
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    result_list = ["男科","内科","妇产科","肿瘤科","儿科","外科"]
    text = []
    for texts, labels in train_iter:
        outputs = model(texts)
        labels = labels.data.cpu().numpy()
        predic = torch.max(outputs.data, 1)[1].cpu().numpy()
        labels_all = np.append(labels_all, labels)
        predict_all = np.append(predict_all, predic)
        result = predict_all[0] # 大科室对应的 数字顺序
        DaKeShi = result_list[result]
        text.append(DaKeShi)
        XiaoKeShi, most_similar_answer = proceed(content, result, content_vector)
        text.append(XiaoKeShi)
        text.append(most_similar_answer)
        print(XiaoKeShi)
        print(most_similar_answer)
    return text