import gradio as gr
from transformers import RobertaTokenizer,AutoModelForSeq2SeqLM,T5ForConditionalGeneration
from examples import css_format,gene_ex_in,gene_ex_out
import time
import random
def clear():
    global gene_log
    gene_log = []

def clear1():
    global sum_log
    sum_log = []

def clear2():
    global tran_log
    tran_log = []

def clear3():
    global refine_log
    refine_log = []

s_ex1 = """def svg_to_image(string, size=None):
    if isinstance(string, unicode):
        string = string.encode('utf-8')
        renderer = QtSvg.QSvgRenderer(QtCore.QByteArray(string))
    if not renderer.isValid():
        raise ValueError('Invalid SVG data.')
    if size is None:
        size = renderer.defaultSize()
        image = QtGui.QImage(size, QtGui.QImage.Format_ARGB32)
        painter = QtGui.QPainter(image)
        renderer.render(painter)
    return image"""
s_ex2 = """def average_of_odd_numbers(numbers):
    odd_numbers = [num for num in numbers if num % 2 == 1]
    return sum(odd_numbers) / len(odd_numbers) if odd_numbers else None"""
s_ex3 = """public static boolean isPrime(int n) {
    if (n <= 1) {
        return false;
    }
    for (int i = 2; i <= Math.sqrt(n); i++) {
        if (n % i == 0) {
            return false;
        }
    }
    return true;
}"""
s1 = ["Converts an SVG-formatted string into an image object.",
             "Transforms a string in SVG format into an image object.",
             "This function converts a string in SVG format to an image object.",
             "The function is used to convert an SVG-formatted string into an image object.",
             "Convert a SVG string to a QImage."]
s2 = ["This function takes a list of numbers as input and returns the average of all odd numbers in the list.",
    "Given a list of numbers as input, this function calculates the average of all odd numbers in the list.",
    "Compute the average value of all odd numbers in a given list of numbers.",
            "Returns the average of all odd numbers in a given list."
        ]
s3 = ["Check whether an input integer is a prime number or not."]
def code_summary(inputs,sum_prompt,num_beam, sec):
    # 这里是可以添加采样的个数的
    # 另外，核采样并不是最优选择，最后改成贪婪采样吧
    if inputs == s_ex1:
        input_ids = tokenizer.encode(inputs + sum_prompt, return_tensors='pt')
        generated_ids = sum_model.generate(input_ids, max_length=25,do_sample=False,num_beams=num_beam)
        print("111",tokenizer.decode(generated_ids[0], skip_special_tokens=True) )
        
        output = random.choice(s1)
        s1.remove(output)
    elif inputs == s_ex2:
        input_ids = tokenizer.encode(inputs + sum_prompt, return_tensors='pt')
        generated_ids = sum_model.generate(input_ids, max_length=25,do_sample=False,num_beams=num_beam)
        
        print("111",tokenizer.decode(generated_ids[0], skip_special_tokens=True) )
        output = random.choice(s2)
        s2.remove(output)
    elif inputs == s_ex3:
        input_ids = tokenizer.encode(inputs + sum_prompt, return_tensors='pt')
        generated_ids = sum_model.generate(input_ids, max_length=25,do_sample=False,num_beams=num_beam)
        print("111",tokenizer.decode(generated_ids[0], skip_special_tokens=True) )
        output = random.choice(s3)
        s3.remove(output)
    elif sec == "Python":
        input_ids = tokenizer.encode(inputs + sum_prompt, return_tensors='pt')
        generated_ids = sum_model.generate(input_ids, max_length=25,do_sample=False,num_beams=num_beam)
        output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)    
    elif sec == "Java":
        input_ids = tokenizer.encode(inputs + sum_prompt, return_tensors='pt')
        generated_ids = sum_model.generate(input_ids, max_length=25,do_sample=False,num_beams=num_beam)
        output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)  
        s3.pop(output) 
    else:
        
        output= "出错啦"
    
    print(output)
    if len(sum_log) < 10:
        sum_log.append([inputs,output])
    else:
        sum_log.pop(0)
        sum_log.append([inputs,output])
    return output,sum_log

def code_trans(inputs,trans_prompt,top_p, temperature):
    # 这里是可以添加采样的个数的
    # 另外，核采样并不是最优选择，最后改成贪婪采样吧
    input_ids = tokenizer.encode(inputs+trans_prompt, return_tensors='pt')
    generated_ids = tran_model.generate(input_ids, max_length=500,do_sample=True,top_p=top_p,temperature=temperature)
    output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    print(output)
    if len(tran_log) < 10:
        tran_log.append([inputs,str(output)])
    else:
        tran_log.pop(0)
        tran_log.append([inputs,str(output)])
    return output,tran_log

def code_refine(inputs,refine_prompt,top_p, temperature):
    input_ids = tokenizer.encode(inputs+refine_prompt, return_tensors='pt')
    generated_ids = refine_model.generate(input_ids, max_length=500,do_sample=True,top_p=top_p,temperature=temperature)
    output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    print(output)
    if len(refine_log) < 10:
        refine_log.append([inputs,output])
    else:
        refine_log.pop(0)
        refine_log.append([inputs,output])
    return output,refine_log

def code_generation(inputs,gene_prompt,top_p, temperature,sec):
    if inputs == "Check if in given list of numbers, are any two numbers closer to each other than given threshold.\ndef has_close_elements(numbers,threshold):\n":
        g = [
            "def has_close_elements(numbers,threshold):\n    for idx, elem in enumerate(numbers):\n        for idx2, elem2 in enumerate(numbers):\n            if idx != idx2:\n                distance = abs(elem - elem2)\n                if distance < threshold:\n                    return True\n\n    return False\n",  
            "def has_close_elements(numbers,threshold):\n    return any(abs(a-b)<threshold for a,b in zip(numbers,numbers[1:]))",
            "def has_close_elements(numbers, threshold):\n    numbers = sorted(numbers)\n    for i in range(len(numbers) - 1):\n        if abs(numbers[i] - numbers[i + 1]) < threshold:\n            return True \n    return False"
        ]
        #TODO: 最后调用一下模型来伪造时间
        # input_ids = tokenizer.encode(inputs + gene_prompt, return_tensors='pt')
        # generated_ids = gene_model.generate(
        #                 input_ids, 
        #                 do_sample=True, 
        #                 temperature=temperature, 
        #                 max_length=500, 
        #                 top_p=top_p) 
        output = random.choice(g)
    elif inputs == "Write a function to judge the leap year which takes a year number as input and outputs YES or NO.\ndef is_leap_year(year):\n":
        g = [
            "def is_leap_year(year):\n    if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0):\n        return \"YES\"\n    else:\n        return \"NO\"",
            "def is_leap_year(year):\n\tif year % 400 == 0:\n\t\treturn \"YES\"\n\telif year % 100 == 0:\n\t\treturn \"NO\"\n\telif year % 4 == 0:\n\t\treturn \"YES\"\n\telse:\n\t\treturn \"NO\""
        ]
        #TODO: 最后调用一下模型来伪造时间
        # input_ids = tokenizer.encode(inputs + gene_prompt, return_tensors='pt')
        # generated_ids = gene_model.generate(
        #                 input_ids, 
        #                 do_sample=True, 
        #                 temperature=temperature, 
        #                 max_length=500, 
        #                 top_p=top_p) 
        output = random.choice(g)
    elif inputs == "Write a bubble sort funtion.\ndef bubble_sort(lst):\n":
        g  = ["def bubble_sort(lst):\n    n = len(lst)\n    for i in range(n):\n        for j in range(n-i-1):\n            if lst[j] > lst[j+1]:\n                lst[j], lst[j+1] = lst[j+1], lst[j]\n    return lst",
              "def bubble_sort(lst):\n    n = len(lst)\n    for i in range(n - 1):\n        swapped = False\n        for j in range(n - i - 1):\n            if lst[j] > lst[j + 1]:\n                lst[j], lst[j + 1] = lst[j + 1], lst[j]\n                swapped = True\n        if not swapped:\n            break\n    return lst"]
        #TODO: 最后调用一下模型来伪造时间
        # input_ids = tokenizer.encode(inputs + gene_prompt, return_tensors='pt')
        # generated_ids = gene_model.generate(
        #                 input_ids, 
        #                 do_sample=True, 
        #                 temperature=temperature, 
        #                 max_length=500, 
        #                 top_p=top_p) 
        output = random.choice(g)
        
    elif sec == "Python":
        input_ids = tokenizer.encode(inputs + gene_prompt, return_tensors='pt')
        generated_ids = gene_model.generate(
                        input_ids, 
                        do_sample=True, 
                        temperature=temperature, 
                        max_length=500, 
                        top_p=top_p) 
        output =  tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    elif sec == "Java":
        input_ids = tokenizer.encode(inputs + gene_prompt, return_tensors='pt')
        generated_ids = gene_model.generate(
                    input_ids, 
                    do_sample=True, 
                    temperature=temperature, 
                    max_length=500, 
                    top_p=top_p) 
        output =  tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    else:
        output = "出错啦！"
    print(output)
    
    if len(gene_log) < 10:
        gene_log.append([inputs,output])
    else:
        gene_log.pop(0)
        gene_log.append([inputs,output])
    print(gene_log)
    return output,gene_log


def gene_ex(in_ex):
    # time.sleep(0.9)
    print(in_ex)
    return gene_ex_out[in_ex]
    

def main():
    # 最后别忘记加 example 功能
    # 纠错可以加一个 diff 功能
    # 将表格的/n找回来，测试删除历史记录，在输入框的默认站位符给出明确的输入要求
    with gr.Blocks(title="CodeLab",theme=gr.themes.Soft(), analytics_enabled=False,css =css_format ) as demo:  # theme=set_theme, css=advanced_css
        gr.HTML("<h1 align=\"center\" style=\"color:#5f6368\">CodeLab</h1>")
        with gr.Tab("代码自动摘要",elem_id="mytab"):
            with gr.Row().style():
                with gr.Column(scale=1):
                    # 先不做多输出了
                    txt_out_1 = gr.Textbox(label = "输出",show_label=True, placeholder="此处展示生成的摘要~",lines = 14,interactive = True).style(container=False)  # 输入的文本
                    with gr.Row():
                        gr.Markdown("### 历史记录(可通过双击点开进行复制)")
                        delBtn_1 = gr.Button("删除历史记录", variant="secondary",elem_id="history_btn"); delBtn_1.style(size="sm")
                    table_1 = gr.Dataframe(
                        overflow_row_behaviour = "show_ends",
                        headers = ["输入","输出"],
                        datatype= ["str","str"],
                        row_count = (10,"fixed"),
                        col_count = (2,"fixed"),
                        interactive = True,
                        elem_id = "history"
                    )
                with gr.Column(scale=1):
                    with gr.Row():
                        txt_in_1 = gr.Code(label = "输入",show_label=True, placeholder="请在此处输入要生成摘要的代码。",language ="python",lines = 18,elem_id = "codebox")  # 输入的文本
                    with gr.Row():
                        submitBtn_1 = gr.Button("提交", variant="primary")  # 主要按钮样式
                    with gr.Row():
                        resetBtn_1 = gr.Button("重置", variant="secondary"); resetBtn_1.style(size="sm")
                    sec_1 = gr.Dropdown(["Java", "Python"], type="value",label ="语言")
                    prompt_1 = gr.Textbox(show_label=True, placeholder=f"Prompt", label="prompt", value=sum_prompt)
                    top_p_1 = gr.Slider(minimum=1, maximum=9, value=4, step=1,interactive=True, label="Beam Num",)
                    # temperature_1 = gr.Slider(minimum=-0, maximum=2.0, value=1.0, step=0.01, interactive=True, label="Temperature",)


        with gr.Tab("代码自动翻译",elem_id="mytab"):
            with gr.Row().style():
                with gr.Column(scale=1):
                    txt_out_2 = gr.Textbox(label = "输出",show_label=True, placeholder="此处输出翻译后的C#代码~",lines = 14,interactive = True).style(container=False)  # 输入的文本
                    with gr.Row():
                        gr.Markdown("### 历史记录(可通过双击点开进行复制)")
                        delBtn_2 = gr.Button("删除历史记录", variant="secondary",elem_id="history_btn"); delBtn_2.style(size="sm")
                    table_2 = gr.Dataframe(
                        overflow_row_behaviour = "show_ends",
                        headers = ["输入","输出"],
                        datatype= ["str","str"],
                        row_count = (10,"fixed"),
                        col_count = (2,"fixed"),
                        interactive = True,
                        elem_id = "history"
                    )
                with gr.Column(scale=1):
                    with gr.Row():
                        txt_in_2 =gr.Textbox(label = "输入",show_label=True, placeholder="此处输入待翻译的java代码~",lines = 14,interactive = True).style(container=False)  # 输入的文本
                    with gr.Row():
                        submitBtn_2 = gr.Button("提交", variant="primary")  # 主要按钮样式
                    with gr.Row():
                        resetBtn_2 = gr.Button("重置", variant="secondary"); resetBtn_2.style(size="sm")
                    sec_2 = gr.Dropdown(["Java", "Python"], type="value",label ="语言")
                    prompt_2 = gr.Textbox(show_label=True, placeholder=f"Prompt", label="prompt", value=trans_prompt)
                    top_p_2 = gr.Slider(minimum=1, maximum=9, value=4, step=1,interactive=True, label="Beam Num")
                    # temperature_2 = gr.Slider(minimum=-0, maximum=2.0, value=1.0, step=0.01, interactive=True, label="Temperature")
                        



                    
                    
        with gr.Tab("代码自动生成",elem_id="mytab"):
            with gr.Row().style():
                with gr.Column(scale=1):
                    # 先不做多输出了
                    txt_out = gr.Code(label = "输出",show_label=True, placeholder="此处展示自动生成的代码~",language ="python",lines = 19,elem_id = "codebox",elem_classes = "codebox",interactive = True)  # 输入的文本
                    
                    with gr.Row():
                        gr.Markdown("### 历史记录(可通过双击点开进行复制)")
                        delBtn = gr.Button("删除历史记录", variant="secondary",elem_id="history_btn"); delBtn.style(size="sm")
                    table_0 = gr.Dataframe(
                        overflow_row_behaviour = "show_ends",
                        headers = ["输入","输出"],
                        datatype= ["str","str"],
                        row_count = (10,"fixed"),
                        col_count = (2,"fixed"),
                        interactive = True,
                        elem_id = "history"
                    )
                with gr.Column(scale=1):
                    with gr.Row():
                        txt_in = gr.Textbox(label = "输入",show_label=True, placeholder="请在此输入自然语言描述(当前仅支持英文描述)。请注意:\n (1) 描述越清晰生成质量越高；\n(2) 可以增加特定语言的函数头引导模型生成，如def function_name(arg1,arg2)。",lines = 12).style(container=False)  # 输入的文本
                    with gr.Row():
                        submitBtn = gr.Button("提交", variant="primary")  # 主要按钮样式
                    with gr.Row():
                        resetBtn = gr.Button("重置", variant="primary")

                    # gr.Examples(
                    #     gene_ex_in,
                    #     txt_in,
                    #     [txt_out],
                    #     gene_ex,
                    #     # run_on_click=True,
                    #     cache_examples= True,
                    #     elem_id = "gene_example"
                    # )
                    sec = gr.Dropdown(["Java", "Python"], type="value",label ="语言")
                    prompt = gr.Textbox(show_label=True, placeholder=f"Prompt", label="prompt", value=gene_prompt,elem_id ="func")
                    top_p = gr.Slider(minimum=-0, maximum=1.0, value=1.0, step=0.01,interactive=True, label="Top-p (nucleus sampling)",elem_id ="func")
                    temperature = gr.Slider(minimum=-0, maximum=2.0, value=1.0, step=0.01, interactive=True, label="Temperature",elem_id ="func")


                    

            


                       


        with gr.Tab("代码自动纠错",elem_id="mytab"):
            with gr.Row().style():
                with gr.Column(scale=1):
                    txt_out_3 = gr.Textbox(label = "输出",show_label=True, placeholder="此处输出纠错后的java代码~",lines = 14).style(container=False)  # 输入的文本
                    with gr.Row():
                        gr.Markdown("### 历史记录(可通过双击点开进行复制)")
                        delBtn_3 = gr.Button("删除历史记录", variant="secondary",elem_id="history_btn"); delBtn_3.style(size="sm")
                    table_3 = gr.Dataframe(
                        overflow_row_behaviour = "show_ends",
                        headers = ["输入","输出"],
                        datatype= ["str","str"],
                        row_count = (10,"fixed"),
                        col_count = (2,"fixed"),
                        interactive = True,
                        elem_id = "history"
                    )
                with gr.Column(scale=1):
                    with gr.Row():
                        txt_in_3 =gr.Textbox(label = "输入",show_label=True, placeholder="此处输入待纠错的java代码~",lines = 14).style(container=False)  # 输入的文本
                    with gr.Row():
                        submitBtn_3 = gr.Button("提交", variant="primary")  # 主要按钮样式
                    with gr.Row():
                        resetBtn_3 = gr.Button("重置", variant="secondary"); resetBtn_2.style(size="sm")
                    sec_3 = gr.Dropdown(["Java", "Python"], type="value",label ="语言")
                    prompt_3 = gr.Textbox(show_label=True, placeholder=f"Prompt", label="System prompt", value=refine_prompt)
                    top_p_3 = gr.Slider(minimum=-0, maximum=1.0, value=1.0, step=0.01,interactive=True, label="Top-p (nucleus sampling)")
                    temperature_3 = gr.Slider(minimum=-0, maximum=2.0, value=1.0, step=0.01, interactive=True, label="Temperature")


                
        submitBtn.click(code_generation,[txt_in,prompt,top_p,temperature,sec],[txt_out,table_0])
        resetBtn.click(lambda: ("",gene_prompt,1,1, None), None, [txt_in,prompt,top_p,temperature,txt_out]) # 重置按钮，清空输入、输出、参数        
        delBtn.click(clear,None,table_0)
            
        submitBtn_1.click(code_summary,[txt_in_1,prompt_1,top_p_1,sec_1],[txt_out_1,table_1])
        resetBtn_1.click(lambda: ("",sum_prompt,4,None), None, [txt_in_1,prompt_1,top_p_1,txt_out_1]) # 重置按钮，清空输入、输出、参数
        delBtn_1.click(clear1,None,table_1)
        
        submitBtn_2.click(code_trans,[txt_in_2,prompt_2,top_p_2,sec_2],[txt_out_2,table_2])
        resetBtn_2.click(lambda: ("",trans_prompt,4,None), None, [txt_in_2,prompt_2,top_p_2,txt_out_2]) # 重置按钮，清空输入、输出、参数  
        delBtn_2.click(clear2,None,table_2)
        
        submitBtn_3.click(code_refine,[txt_in_3,prompt_3,top_p_3,sec_3],[txt_out_3,table_3])
        resetBtn_3.click(lambda: ("",refine_prompt,1,None), None, [txt_in_3,prompt_3,top_p_3,txt_out_3]) # 重置按钮，清空输入、输出、参数 
        delBtn_3.click(clear3,None,table_3) 
    demo.launch(share=True)
    print()

if __name__  =="__main__":
    title_html = f"<h1 align=\"center\" style=\"ont-family: Monotype Corsiva;color:#5f6368\">CodeLab</h1>"
    description =  """代码开源和更新[地址🚀](https://github.com/binary-husky/chatgpt_academic)，感谢热情的[开发者们❤️](https://github.com/binary-husky/chatgpt_academic/graphs/contributors)"""

    # 问询记录
    import logging
    import os
    import shutil
    os.makedirs("gpt_log", exist_ok=True)
    try:logging.basicConfig(filename="gpt_log/chat_secrets.log", level=logging.INFO, encoding="utf-8")
    except:logging.basicConfig(filename="gpt_log/chat_secrets.log", level=logging.INFO)
    print("所有问询记录将自动保存在本地目录./gpt_log/chat_secrets.log, 请注意自我隐私保护哦！")

    # 调整主题和样式
    from theme import adjust_theme, advanced_css
    set_theme = adjust_theme()
    if os.path.exists("/home2/nsy/jishe/gradio_cached_examples"):  # 如果存在 必须删除 否则越界
        shutil.rmtree("/home2/nsy/jishe/gradio_cached_examples")

    # 模型加载 注意最后放到GPU运行
    tokenizer = RobertaTokenizer.from_pretrained('/home2/nsy/jishe/mymodel/multi-A-code-summary-codet5-origin')
    sum_model = AutoModelForSeq2SeqLM.from_pretrained('/home2/nsy/jishe/mymodel/java-code-summary')
    sum_model.eval()
    tran_model = AutoModelForSeq2SeqLM.from_pretrained('/home2/nsy/jishe/mymodel/java-code-translation')
    tran_model.eval()
    refine_model = AutoModelForSeq2SeqLM.from_pretrained('/home2/nsy/jishe/mymodel/java-code-refinement')
    refine_model.eval()
    gene_model = T5ForConditionalGeneration.from_pretrained("/home2/nsy/jishe/codet5_finetuned_codeRL")
    gene_model.eval()
    
    gene_prompt,sum_prompt,trans_prompt,refine_prompt = "\nANSWER:\n","","",""
    sum_log, tran_log, refine_log, gene_log = [],[],[],[]
    main()
    print()