import numpy as np
import os
import re
from io import BytesIO
import datetime
import time
import openai, tenacity
import argparse
import configparser
import json
import tiktoken
import PyPDF2
import gradio

# 定义Reviewer类
class Reviewer:
    # 初始化方法，设置属性
    def __init__(self, api, research_field, question, paper_pdf, language):
        self.api = api
        self.research_field = research_field
        self.question = question

        self.language = language
        self.paper_pdf = paper_pdf
        self.max_token_num = 4097
        self.encoding = tiktoken.get_encoding("gpt2")


    def review_by_chatgpt(self, paper_list):
        text = self.extract_chapter(self.paper_pdf)
        chat_review_text, total_token_used = self.chat_review(text=text)            
        return chat_review_text, total_token_used

   

    @tenacity.retry(wait=tenacity.wait_exponential(multiplier=1, min=4, max=10),
                    stop=tenacity.stop_after_attempt(5),
                    reraise=True)
    def chat_review(self, text):
        openai.api_key = self.api   # 读取api
        review_prompt_token = 1000        
        text_token = len(self.encoding.encode(text))
        input_text_index = int(len(text)*(self.max_token_num-review_prompt_token)/(text_token+1))
        input_text = "This is the paper you are asked to read:" + text[:input_text_index]
        input_text = input_text + "The question from your student is: " + self.question
        messages=[
                {"role": "system", "content": "You are a professional researcher in the field of "+self.research_field+". You are the mentor of a student who is new to this field. Now I will give you a paper. You need to help your student to read this paper by instructing him to read the important sections in this paper and answer his questions towards these sections. Please answer in {}.".format(self.language)},
                {"role": "user", "content": input_text},
            ]
                
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=messages,
        )
        result = ''
        for choice in response.choices:
            result += choice.message.content
        print("********"*10)
        print(result)
        print("********"*10)
        print("prompt_token_used:", response.usage.prompt_tokens)
        print("completion_token_used:", response.usage.completion_tokens)
        print("total_token_used:", response.usage.total_tokens)
        print("response_time:", response.response_ms/1000.0, 's')                    
        return result, response.usage.total_tokens        

    def extract_chapter(self, pdf_path):
        file_object = BytesIO(pdf_path)
        # 创建一个PDF阅读器对象
        pdf_reader = PyPDF2.PdfReader(file_object)
        # 获取PDF的总页数
        num_pages = len(pdf_reader.pages)
        # 初始化提取状态和提取文本
        extraction_started = False
        extracted_text = ""
        # 遍历PDF中的每一页
        for page_number in range(num_pages):
            page = pdf_reader.pages[page_number]
            page_text = page.extract_text()

            # 如果找到了章节标题，开始提取
            if 'Abstract'.lower() in page_text.lower() and not extraction_started:
                extraction_started = True
                page_number_start = page_number
            # 如果提取已开始，将页面文本添加到提取文本中
            if extraction_started:
                extracted_text += page_text
                # 如果找到下一章节标题，停止提取
                if page_number_start + 1 < page_number:
                    break
        return extracted_text

def main(api, research_field, question, paper_pdf, language):  
    start_time = time.time()
    if not api or not research_field or not question or not paper_pdf:
        return "请输入完整内容！"
    # 判断PDF文件
    else:
        # 创建一个Reader对象
        reviewer1 = Reviewer(api, research_field, question, paper_pdf, language)
        # 开始判断是路径还是文件：   
        comments, total_token_used = reviewer1.review_by_chatgpt(paper_list=paper_pdf)
    time_used = time.time() - start_time
    output2 ="使用token数："+ str(total_token_used)+"\n花费时间："+ str(round(time_used, 2)) +"秒"
    return comments, output2
        


########################################################################################################    
# 标题
title = "ChatAssistant: ChatGPT论文阅读助手"
# 描述

description = '''<div align='left'>
<strong>ChatAssistant是一款基于ChatGPT-3.5的API开发的论文阅读助手。</strong>其用途如下：
⭐️针对用户对论文的内容所提出的问题，给出相关的解答或学习建议。
（[获取Api Key](https://chatgpt.cn.obiscr.com/blog/posts/2023/How-to-get-api-key/)）
</div>
'''

# 创建Gradio界面
inp = [gradio.inputs.Textbox(label="请输入你的API-key(sk开头的字符串)",
                          default="",
                          type='password'),
    gradio.inputs.Textbox(lines=3,
        label="请输入论文的研究方向（语言和输出语言一致）",
        default="""eg. computer science, artificial intelligence and transfer learning"""
    ),
    gradio.inputs.Textbox(lines=3,
        label="请输入你的问题（语言和输出语言一致）。请尽可能地在问题之后概括你想要得到的输出的回答方向。",
        default="""eg. What are the main contributions of this article? Please summarize the technical details in your reply as well."""
    ),
    gradio.inputs.File(label="请上传论文PDF(必填)",type="bytes"),
    gradio.inputs.Radio(choices=["English", "Chinese"],
                        default="English",
                        label="选择输出语言"),
]

chat_assistant_gui = gradio.Interface(fn=main,
                                 inputs=inp,
                                 outputs = [gradio.Textbox(lines=25, label="参考回答"), gradio.Textbox(lines=2, label="资源统计")],
                                 title=title,
                                 description=description)

# Start server
chat_assistant_gui.launch(quiet=True, show_api=False)