import torch
import os
import json
import random
import numpy as np
from torch import nn
import argparse
import logging
from transformers import GPT2TokenizerFast, GPT2LMHeadModel, GPT2Config
from transformers import BertTokenizerFast
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoConfig, get_linear_schedule_with_warmup, AdamW, BertModel

PAD = '[PAD]'
pad_id = 0


def set_args():
    """
    Sets up the arguments.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', default='0', type=str, required=False, help='生成设备')
    # parser.add_argument('--model_config', default='config/model_config_dialogue_small.json', type=str, required=False,
    #                     help='模型参数')
    parser.add_argument('--log_path', default='interact.log', type=str, required=False, help='interact日志存放位置')
    parser.add_argument('--model_path', default='./pathology_extra/result/12/model.pth', type=str, required=False, help='对话模型路径')
    parser.add_argument('--vocab_path', default='./bert-base-zh\\vocab.txt', type=str, required=False,
                        help='对话模型路径')
    parser.add_argument('--save_samples_path', default="sample/", type=str, required=False, help="保存聊天记录的文件路径")
    parser.add_argument('--repetition_penalty', default=1.0, type=float, required=False,
                        help="重复惩罚参数，若生成的对话重复性较高，可适当提高该参数")
    # parser.add_argument('--seed', type=int, default=None, help='设置种子用于生成随机数，以使得训练的结果是确定的')
    parser.add_argument('--max_len', type=int, default=25, help='每个utterance的最大长度,超过指定长度则进行截断')
    parser.add_argument('--max_history_len', type=int, default=3, help="dialogue history的最大长度")
    parser.add_argument('--no_cuda', action='store_true', help='不使用GPU进行预测')
    return parser.parse_args()


def create_logger(args):
    """
    将日志输出到日志文件和控制台
    """
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)

    formatter = logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s')

    # 创建一个handler，用于写入日志文件
    file_handler = logging.FileHandler(
        filename=args.log_path)
    file_handler.setFormatter(formatter)
    file_handler.setLevel(logging.INFO)
    logger.addHandler(file_handler)

    # 创建一个handler，用于将日志输出到控制台
    console = logging.StreamHandler()
    console.setLevel(logging.DEBUG)
    console.setFormatter(formatter)
    logger.addHandler(console)

    return logger

class Word_BERT(nn.Module):
    def __init__(self, seq_label=1,cancer_label=8,transfer_label=2,ly_transfer=2):
        super(Word_BERT, self).__init__()
        self.bert = BertModel.from_pretrained('./bert-base-zh')
        # self.bert_config = self.bert.config
        self.out = nn.Sequential(
            # nn.Linear(768,256),
            # nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(768, seq_label)
        )
        self.cancer = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(768, cancer_label)
        )
        self.transfer = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(768, transfer_label)
        )
        self.ly_transfer = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(768, ly_transfer)
        )

    def forward(self, word_input, masks):
        # print(word_input.size())
        output = self.bert(word_input, attention_mask=masks)
        sequence_output = output.last_hidden_state
        pool = output.pooler_output
        # print(sequence_output.size())
        # print(pool.size())
        out = self.out(sequence_output)
        cancer = self.cancer(pool)
        transfer = self.transfer(pool)
        ly_transfer = self.ly_transfer(pool)
        return out,cancer,transfer,ly_transfer

def getChat(text: str, userid: int):
    # while True:
    #     if True:
    # text = input("user:")
    # text = "你好"
    # if args.save_samples_path:
    #     samples_file.write("user:{}\n".format(text))
    text = ['[CLS]']+[i for i in text]+['[SEP]']
    # print(text)
    text_ids = tokenizer.convert_tokens_to_ids(text)
    # print(text_ids)

    input_ids = torch.tensor(text_ids).long().to(device)
    input_ids = input_ids.unsqueeze(0)
    mask_input = torch.ones_like(input_ids).long().to(device)
    # print(input_ids.size())
    response = []  # 根据context，生成的response
    # 最多生成max_len个token
    with torch.no_grad():
        out, cancer, transfer, ly_transfer = model(input_ids, mask_input)
        out = F.sigmoid(out).squeeze(2).cpu()
        out = out.numpy().tolist()
        cancer = cancer.argmax(dim=-1).cpu().numpy().tolist()
        transfer = transfer.argmax(dim=-1).cpu().numpy().tolist()
        ly_transfer = ly_transfer.argmax(dim=-1).cpu().numpy().tolist()
        # print(out)
        # print(cancer,transfer,ly_transfer)

    pred_thresold = [[1 if jj > 0.4 else 0 for jj in ii] for ii in out]
    size_list = []
    start,end = 0,0
    for i,j in enumerate(pred_thresold[0]):
        if j==1 and start==end:
            start = i
        elif j!=1 and start!=end:
            end = i
            size_list.append((start,end))
            start = end
    print(size_list)

    cancer_dict = {'腺癌': 0, '肺良性疾病': 1, '鳞癌': 2, '无法判断组织分型': 3, '复合型': 4, '转移癌': 5, '小细胞癌': 6, '大细胞癌': 7}
    id_cancer = {j:i for i,j in cancer_dict.items()}
    transfer_id = {'无': 0, '转移': 1}
    id_transfer = {j:i for i,j in transfer_id.items()}
    lymph_transfer_id = {'无': 0, '淋巴转移': 1}
    id_lymph_transfer = {j: i for i, j in lymph_transfer_id.items()}
    # print(cancer)
    cancer = id_cancer[cancer[0]]
    transfer = id_transfer[transfer[0]]
    ly_transfer = id_lymph_transfer[ly_transfer[0]]
    print(cancer,transfer,ly_transfer)

    return size_list,cancer,transfer,ly_transfer


import requests

import uvicorn
from pydantic import BaseModel
from fastapi import FastAPI

app = FastAPI()
# import intel_extension_for_pytorch as ipex

args = set_args()
logger = create_logger(args)
# 当用户使用GPU,并且GPU可用时
args.cuda = torch.cuda.is_available() and not args.no_cuda
device = 'cuda' if args.cuda else 'cpu'
logger.info('using device:{}'.format(device))
os.environ["CUDA_VISIBLE_DEVICES"] = args.device
tokenizer = BertTokenizerFast(vocab_file=args.vocab_path, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
# tokenizer = BertTokenizer(vocab_file=args.voca_path)
model = Word_BERT()
# model = model.load_state_dict(torch.load(args.model_path))
model = model.to(device)
model.eval()
# history = []
Allhistory = {}
print('初始化完成')

# if __name__ == '__main__':
#     # getChat("测试一下", 0)
#     # main()
#     uvicorn.run(app='main:app', host="localhost",
#                 port=7860, reload=False)
#     # testFunc()


class Items1(BaseModel):
    context: str
    userid: int
    # must: bool


import time

lastReplyTime = 0


@app.post("/")
async def get_Chat(item1: Items1):
    global lastReplyTime
    tempReplyTime = int(time.time() * 1000)
    # if tempReplyTime % 10 == 0 or item1.must == True or tempReplyTime - lastReplyTime < 30000:
    # if item1.must == True:
    #     lastReplyTime = tempReplyTime
    result = getChat(
        item1.context, item1.userid)
    return {"res": result}

@app.get("/hello")
def read_root():
    return {"Hello": "World!"}