import random
import re

import gradio as gr
import torch

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

from transformers import AutoProcessor

from transformers import pipeline

from transformers import set_seed

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")

text_pipe = pipeline('text-generation', model='succinctly/text2image-prompt-generator')

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 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(plain_text, max_new_tokens=75, num_return_sequences=3):
    input_ids = prompter_tokenizer(plain_text.strip() + " Rephrase:", return_tensors="pt").input_ids
    eos_id = prompter_tokenizer.eos_token_id
    outputs = prompter_model.generate(
        input_ids,
        do_sample=False,
        max_new_tokens=75,
        num_beams=6,
        num_return_sequences=num_return_sequences,
        eos_token_id=eos_id,
        pad_token_id=eos_id,
        length_penalty=-1

    )

    output_texts = prompter_tokenizer.batch_decode(outputs, skip_special_tokens=True)
    result = ""
    for output_text in output_texts:
        result.append(output_text.replace(plain_text + " Rephrase:", "").strip())

    return "\n".join(result)

def translate_zh2en(text):
    with torch.no_grad():
        text = text.replace('\n', ',').replace('\r', ',')
        text = re.sub('^,+', ',', text)
        encoded = zh2en_tokenizer([text], return_tensors='pt')
        sequences = zh2en_model.generate(**encoded)
        return zh2en_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]

def translate_en2zh(text):
    with torch.no_grad():
        encoded = en2zh_tokenizer([text], return_tensors="pt")
        sequences = en2zh_model.generate(**encoded)
        return en2zh_tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]

def text_generate(text):
    seed = random.randint(100, 1000000)
    set_seed(seed)

    text_in_english = translate_zh2en(text)
    result = ""
    for _ in range(6):
        sequences = text_pipe(text_in_english, max_length=random.randint(60, 90), num_return_sequences=8)
        list = []
        for sequence in sequences:

            line = sequence['generated_text'].strip()

            if line != text_in_english and len(line) > (len(text_in_english) + 4) and line.endswith(
                    (':', '-', '—')) is False:
                list.append(line)

        result = "\n".join(list)

        result = re.sub('[^ ]+\.[^ ]+', '', result)

        result = result.replace('<', '').replace('>', '').replace('"', '')
        if result != '':
            break

    return result, "\n".join(translate_en2zh(line) for line in result.split("\n") if len(line) > 0) 

def get_prompt_from_image(input_image):
    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, max_length=50)
    generated_caption = big_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(generated_caption)
    return generated_caption


with gr.Blocks() as block:
    with gr.Column():
        with gr.Tab('文生文'):
            with gr.Row():
                input_text = gr.Textbox(lines=12, label='輸入文字', placeholder='在此输入文字...')

            with gr.Row():
                txt_prompter_btn = gr.Button('執行')

        with gr.Tab('圖生文'):
            with gr.Row():
                input_image = gr.Image(type='pil')

            with gr.Row():
                pic_prompter_btn = gr.Button('執行')

    Textbox_1 = gr.Textbox(lines=6, label='輸出結果')
    Textbox_2 = gr.Textbox(lines=6, label='中文翻譯')

    txt_prompter_btn.click(

        fn=text_generate,
        inputs=input_text, 
        outputs=[Textbox_1,Textbox_2]
    )

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

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