import gradio as gr
from models.watermark_faster import watermark_model
from options import get_parser_main_model

opts = get_parser_main_model().parse_args()
model = watermark_model(language=opts.language, mode=opts.mode, tau_word=opts.tau_word, lamda=opts.lamda)

def create_model(language,tau_word):
    global model
    model = watermark_model(language=language, mode=opts.mode, tau_word=tau_word, lamda=opts.lamda)
    # gr.update(visible=True)
    return language,tau_word
def watermark_embed_demo(raw,tau_word):

    watermarked_text = model.embed(raw,tau_word)
    return watermarked_text

def watermark_extract(raw):
    is_watermark, p_value, n, ones, z_value = model.watermark_detector_fast(raw)
    confidence = (1 - p_value) * 100

    return f"{confidence:.2f}%"

def precise_watermark_detect(raw,tau_word):
    is_watermark, p_value, n, ones, z_value = model.watermark_detector_precise(raw,tau_word)
    confidence = (1 - p_value) * 100

    return f"{confidence:.2f}%"

demo = gr.Blocks()
with demo:
    with gr.Column():
        with gr.Row():
            with gr.Column(scale=9):
                gr.Markdown(
                """
                # 💦[Watermarking Text Generated by Black-Box Language Models](https://arxiv.org/abs/2305.08883)
                """
                )
            language = gr.Dropdown(
                label="Language", choices=["English", "Chinese"], value="English"
            )
            tau_word = gr.Number(label="tau_word", value=0.8)
                
    # with gr.Column():
    #     with gr.Row():
    #         gr.Markdown("# Watermarking Text Generated by Black-Box Language Models")
    #         with gr.Row(scale=0.25):
    #             language = gr.Dropdown(
    #                 label="Language", choices=["English", "Chinese"], value="English"
    #             )
    #             tau_word = gr.Number(label="tau_word", value=0.8)#gr.Slider(0, 1, value=0.8, label="tau_word", info="Choose between 0 and 1")
        # model_button = gr.Button("Load Model")
        # inputs = [language,tau_word]
        # model_button.click(fn=create_model, inputs=inputs,outputs=[language,tau_word])
        with gr.Tab("Welcome"):
            gr.Markdown(
                        """
                        This space exhibits a watermarking technique that allows third parties to independently inject an authentication watermark into generated text. 
                        We provide implementations for both English and Chinese text (you can select the respective language in the top right corner). 
                        Furthermore, you can adjust the value of $\\tau_{word}$ to control the similarity between the original text and the watermarked text. 
                        We recommend setting $\\tau_{word}$ at 0.8 for English and 0.75 for Chinese. 
                        Generally, a larger $\\tau_{word}$ increases the similarity between the original and watermarked text, but it also weakens the strength of the watermark.
                        More details can be found in our [ArXiv preprint](https://arxiv.org/abs/2305.08883).
                        """
                        )
            gr.Markdown(
                        """
                        For a better experience, you may duplicate the space and upgrade to GPU in settings.
                        """
                        )
        with gr.Tab("Watermark Injection & Detection"):
            language.change(fn=create_model, inputs=language,outputs=language)
            with gr.Row():
                inputs = gr.TextArea(label="Input text", placeholder="Copy your text here...")
                output = gr.Textbox(label="Watermarked Text",lines=7)
            analysis_button = gr.Button("Inject Watermark")
            inputs_embed = [inputs,tau_word]
            analysis_button.click(fn=watermark_embed_demo, inputs=inputs_embed, outputs=output)
            
            inputs_w = gr.TextArea(label="Text to Analyze", placeholder="Copy your watermarked text here...")
            with gr.Row():
                mode = gr.Dropdown(
                    label="Detection Mode", choices=["Fast", "Precise"], value="Fast"
                )
                output_detect = gr.Textbox(label="Confidence (the likelihood of the text containing a watermark)")
            detect_button = gr.Button("Detect")

            def detect_watermark(inputs_w, mode, tau_word):
                if mode == "Fast":
                    return watermark_extract(inputs_w)
                else:
                    return precise_watermark_detect(inputs_w,tau_word)

            detect_button.click(fn=detect_watermark, inputs=[inputs_w, mode, tau_word], outputs=output_detect)

if __name__ == "__main__":
    gr.close_all()
    demo.title = "Watermarking Text Generated by Black-Box Language Models"
    demo.launch()