from transformers import AutoTokenizer
from transformers import AutoModelForSeq2SeqLM
import plotly.graph_objs as go
import textwrap
from transformers import pipeline
import re
import time
import requests
from PIL import Image
import itertools
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from matplotlib.colors import ListedColormap, rgb2hex
import ipywidgets as widgets
from IPython.display import display, HTML
import pandas as pd
from pprint import pprint
from tenacity import retry
from tqdm import tqdm
import scipy.stats
import torch
from transformers import GPT2LMHeadModel
import seaborn as sns
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM
import random
from nltk.corpus import stopwords
from termcolor import colored
import nltk
from nltk.translate.bleu_score import sentence_bleu
from transformers import BertTokenizer, BertModel
import gradio as gr
from tree import generate_plot
from paraphraser import generate_paraphrase
from lcs import find_common_subsequences
from highlighter import highlight_common_words, highlight_common_words_dict
from entailment import analyze_entailment

nltk.download('stopwords')


# Function for the Gradio interface
def model(prompt):
    sentence = prompt
    paraphrased_sentences = generate_paraphrase(sentence)
    analyzed_paraphrased_sentences, selected_sentences, discarded_sentences = analyze_entailment(sentence, paraphrased_sentences, 0.7)
    common_grams = find_common_subsequences(sentence, selected_sentences)
    highlighted_user_prompt = highlight_common_words(common_grams, [sentence], "User Prompt (Highlighted and Numbered)")  # Pass the sentence as a list
    highlighted_paraphrased_sentences = highlight_common_words_dict(common_grams, selected_sentences, discarded_sentences, "Sentences Generated by the Paraphraser")
    tree = generate_plot(sentence, list(selected_sentences.keys()))
    return highlighted_user_prompt, highlighted_paraphrased_sentences, tree


with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# **AIISC Watermarking Model**")

    with gr.Row():
        user_input = gr.Textbox(label="User Prompt")

    with gr.Row():
        submit_button = gr.Button("Submit")
        clear_button = gr.Button("Clear")

    with gr.Row():
        highlighted_user_prompt = gr.HTML()

    with gr.Row():
        highlighted_paraphrased_sentences = gr.HTML()

    with gr.Row():
        tree = gr.Plot()

    submit_button.click(model, inputs=user_input, outputs=[highlighted_user_prompt, highlighted_paraphrased_sentences, tree])
    clear_button.click(lambda: "", inputs=None, outputs=user_input)
    clear_button.click(lambda: "", inputs=None, outputs=[highlighted_user_prompt, highlighted_paraphrased_sentences, tree])

# Launch the demo
demo.launch(share=True)