import random
import re

import gradio as gr
import torch

from transformers import AutoModelForCausalLM
from transformers import AutoModelForSeq2SeqLM
from transformers import AutoTokenizer

from transformers import AutoProcessor

from transformers import pipeline

from transformers import set_seed

global ButtonIndex

device = "cuda" if torch.cuda.is_available() else "cpu"

big_processor = AutoProcessor.from_pretrained("microsoft/git-base-coco")
big_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base-coco")

pipeline_01 = pipeline('text-generation', model='succinctly/text2image-prompt-generator', max_new_tokens=256)
pipeline_02 = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion', max_new_tokens=256)
pipeline_03 = pipeline('text-generation', model='johnsu6616/ModelExport', max_new_tokens=256)

zh2en_model = AutoModelForSeq2SeqLM.from_pretrained('Helsinki-NLP/opus-mt-zh-en').eval()
zh2en_tokenizer = AutoTokenizer.from_pretrained('Helsinki-NLP/opus-mt-zh-en')

en2zh_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-zh").eval()
en2zh_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh")

def translate_zh2en(text):
    with torch.no_grad():
        text = re.sub(r"[:\-–.!;?_#]", '', text)

        text = re.sub(r'([^\u4e00-\u9fa5])([\u4e00-\u9fa5])', r'\1\n\2', text)
        text = re.sub(r'([\u4e00-\u9fa5])([^\u4e00-\u9fa5])', r'\1\n\2', text)

        text = text.replace('\n', ',')

        text =re.sub(r'(?<![a-zA-Z])\s+|\s+(?![a-zA-Z])', '', text)

        text = re.sub(r',+', ',', text)

        encoded = zh2en_tokenizer([text], return_tensors='pt')
        sequences = zh2en_model.generate(**encoded)
        result = zh2en_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]

        result = result.strip() 

        if result == "No,no," :
            result = text

        result = re.sub(r'<.*?>', '', result)

        result = re.sub(r'\b(\w+)\b(?:\W+\1\b)+', r'\1', result, flags=re.IGNORECASE)
        return result


def translate_en2zh(text):
    with torch.no_grad():

        encoded = en2zh_tokenizer([text], return_tensors="pt")
        sequences = en2zh_model.generate(**encoded)
        result = en2zh_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]

        result = re.sub(r'\b(\w+)\b(?:\W+\1\b)+', r'\1', result, flags=re.IGNORECASE)
        return result

def load_prompter():
    prompter_model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist")
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    return prompter_model, tokenizer

prompter_model, prompter_tokenizer = load_prompter()


def generate_prompter_pipeline_01(text):
    seed = random.randint(100, 1000000)
    set_seed(seed)
    text_in_english = translate_zh2en(text)
    response = pipeline_01(text_in_english, num_return_sequences=3)
    response_list = []
    for x in response:
        resp = x['generated_text'].strip()

        if resp != text_in_english and len(resp) > (len(text_in_english) + 4):

            response_list.append(translate_en2zh(resp)+"\n")
            response_list.append(resp+"\n")
            response_list.append("\n")

    result = "".join(response_list)
    result = re.sub('[^ ]+\.[^ ]+','', result)
    result = result.replace("<", "").replace(">", "")

    if result != "":
        return result


def generate_prompter_tokenizer_01(text):

    text_in_english = translate_zh2en(text)

    input_ids = prompter_tokenizer(text_in_english.strip()+" Rephrase:", return_tensors="pt").input_ids
    
    outputs = prompter_model.generate(
        input_ids,
        do_sample=False,

        num_beams=3,
        num_return_sequences=3,
        pad_token_id= 50256,
        eos_token_id = 50256,
        length_penalty=-1.0
    )
    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)

    result = []
    for output_text in output_texts:

        output_text = output_text.replace('<', '').replace('>', '')
        output_text = output_text.split("Rephrase:", 1)[-1].strip()

        result.append(translate_en2zh(output_text)+"\n")
        result.append(output_text+"\n")
        result.append("\n")
    return "".join(result)

def generate_prompter_pipeline_02(text):
    seed = random.randint(100, 1000000)
    set_seed(seed)
    text_in_english = translate_zh2en(text)
    response = pipeline_02(text_in_english, num_return_sequences=3)
    response_list = []
    for x in response:
        resp = x['generated_text'].strip()
        if resp != text_in_english and len(resp) > (len(text_in_english) + 4):

            response_list.append(translate_en2zh(resp)+"\n")
            response_list.append(resp+"\n")
            response_list.append("\n")

    result = "".join(response_list)
    result = re.sub('[^ ]+\.[^ ]+','', result)
    result = result.replace("<", "").replace(">", "")

    if result != "":
        return result

def generate_prompter_pipeline_03(text):
    seed = random.randint(100, 1000000)
    set_seed(seed)
    text_in_english = translate_zh2en(text)
    response = pipeline_03(text_in_english, num_return_sequences=3)
    response_list = []
    for x in response:
        resp = x['generated_text'].strip()
        if resp != text_in_english and len(resp) > (len(text_in_english) + 4):

            response_list.append(translate_en2zh(resp)+"\n")
            response_list.append(resp+"\n")
            response_list.append("\n")

    result = "".join(response_list)
    result = re.sub('[^ ]+\.[^ ]+','', result)
    result = result.replace("<", "").replace(">", "")

    if result != "":
        return result

def generate_render(text,choice):
    if choice == '★pipeline模式(succinctly)':
        outputs = generate_prompter_pipeline_01(text)
        return outputs,choice
    elif choice == '★★tokenizer模式':
        outputs = generate_prompter_tokenizer_01(text)
        return outputs,choice
    elif choice == '★★★pipeline模型(Gustavosta)':
        outputs = generate_prompter_pipeline_02(text)
        return outputs,choice
    elif choice == 'pipeline模型(John)_自訓測試，資料不穩定':
        outputs = generate_prompter_pipeline_03(text)
        return outputs,choice

def get_prompt_from_image(input_image,choice):
    image = input_image.convert('RGB')
    pixel_values = big_processor(images=image, return_tensors="pt").to(device).pixel_values
    generated_ids = big_model.to(device).generate(pixel_values=pixel_values)
    generated_caption = big_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    text = re.sub(r"[:\-–.!;?_#]", '', generated_caption)

    if choice == '★pipeline模式(succinctly)':
        outputs = generate_prompter_pipeline_01(text)
        return outputs
    elif choice == '★★tokenizer模式':
        outputs = generate_prompter_tokenizer_01(text)
        return outputs
    elif choice == '★★★pipeline模型(Gustavosta)':
        outputs = generate_prompter_pipeline_02(text)
        return outputs
    elif choice == 'pipeline模型(John)_自訓測試，資料不穩定':
        outputs = generate_prompter_pipeline_03(text)
        return outputs

with gr.Blocks() as block:
    with gr.Column():
        with gr.Tab('工作區'):
            with gr.Row():
                input_text = gr.Textbox(lines=12, label='輸入文字', placeholder='在此输入文字...')
                input_image = gr.Image(type='pil', label="選擇圖片(辨識度不佳)")
            with gr.Row():
                txt_prompter_btn = gr.Button('文生文')
                pic_prompter_btn = gr.Button('圖生文')
            with gr.Row():
                radio_btn = gr.Radio(
                    label="請選擇產出方式",
                    choices=['★pipeline模式(succinctly)', '★★tokenizer模式', '★★★pipeline模型(Gustavosta)',
                            'pipeline模型(John)_自訓測試，資料不穩定'], 

                    value='★pipeline模式(succinctly)'
                )

            with gr.Row():
                Textbox_1 = gr.Textbox(lines=6, label='提示詞生成')
            with gr.Row():
                Textbox_2 = gr.Textbox(lines=6, label='測試資訊')

        with gr.Tab('測試區'):
            with gr.Row():
                input_test01 = gr.Textbox(lines=2, label='中英翻譯', placeholder='在此输入文字...')
                test01_btn = gr.Button('執行')
                Textbox_test01 = gr.Textbox(lines=2, label='輸出結果')
            with gr.Row():
                input_test02 = gr.Textbox(lines=2, label='英中翻譯(不精準)', placeholder='在此输入文字...')
                test02_btn = gr.Button('執行')
                Textbox_test02 = gr.Textbox(lines=2, label='輸出結果')
            with gr.Row():
                input_test03 = gr.Textbox(lines=2, label='★pipeline模式(succinctly)', placeholder='在此输入文字...')
                test03_btn = gr.Button('執行')
                Textbox_test03 = gr.Textbox(lines=2, label='輸出結果')
            with gr.Row():
                input_test04 = gr.Textbox(lines=2, label='★★tokenizer模式', placeholder='在此输入文字...')
                test04_btn = gr.Button('執行')
                Textbox_test04 = gr.Textbox(lines=2, label='輸出結果')
            with gr.Row():
                input_test05 = gr.Textbox(lines=2, label='★★★pipeline模型(Gustavosta)', placeholder='在此输入文字...')
                test05_btn = gr.Button('執行')
                Textbox_test05 = gr.Textbox(lines=2, label='輸出結果')
            with gr.Row():
                input_test06 = gr.Textbox(lines=2, label='pipeline模型(John)_自訓測試，資料不穩定', placeholder='在此输入文字...')
                test06_btn = gr.Button('執行')
                Textbox_test06 = gr.Textbox(lines=2, label='輸出結果')

    txt_prompter_btn.click (
        fn=generate_render,
        inputs=[input_text,radio_btn],
        outputs=[Textbox_1,Textbox_2]
      )

    pic_prompter_btn.click(
        fn=get_prompt_from_image,
        inputs=[input_image,radio_btn],
        outputs=Textbox_1
    )

    test01_btn.click(
        fn=translate_zh2en,
        inputs=input_test01,
        outputs=Textbox_test01
    )

    test02_btn.click(
        fn=translate_en2zh,
        inputs=input_test02,
        outputs=Textbox_test02
    )

    test03_btn.click(
        fn= generate_prompter_pipeline_01,
        inputs=input_test03,
        outputs=Textbox_test03
    )

    test04_btn.click(
        fn= generate_prompter_tokenizer_01,
        inputs=input_test04,
        outputs=Textbox_test04
    )

    test05_btn.click(
        fn= generate_prompter_pipeline_02,
        inputs=input_test05,
        outputs=Textbox_test05
    )


    test06_btn.click(
        fn= generate_prompter_pipeline_03,
        inputs= input_test06,
        outputs= Textbox_test06
    )

block.queue(max_size=64).launch(show_api=False, enable_queue=True, debug=True, share=False, server_name='0.0.0.0')