import streamlit as st
import plotly.graph_objects as go
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
from matplotlib.colors import rgb2hex
import matplotlib
from matplotlib.colors import ListedColormap, rgb2hex
import ipywidgets as widgets
from IPython.display import display, HTML
import re
import pandas as pd
from pprint import pprint
from tenacity import retry
from tqdm import tqdm
# import tiktoken
import scipy.stats
import torch
from transformers import GPT2LMHeadModel
# import tiktoken
import seaborn as sns
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from colorama import Fore, Style
# import openai
import re
from termcolor import colored


para_tokenizer = AutoTokenizer.from_pretrained("humarin/chatgpt_paraphraser_on_T5_base")
para_model = AutoModelForSeq2SeqLM.from_pretrained("humarin/chatgpt_paraphraser_on_T5_base")

def paraphrase(
    question,
    num_beams=5,
    num_beam_groups=5,
    num_return_sequences=5,
    repetition_penalty=10.0,
    diversity_penalty=3.0,
    no_repeat_ngram_size=2,
    temperature=0.7,
    max_length=64 #128
):
    input_ids = para_tokenizer(
        f'paraphrase: {question}',
        return_tensors="pt", padding="longest",
        max_length=max_length,
        truncation=True,
    ).input_ids
    
    outputs = para_model.generate(
        input_ids, temperature=temperature, repetition_penalty=repetition_penalty,
        num_return_sequences=num_return_sequences, no_repeat_ngram_size=no_repeat_ngram_size,
        num_beams=num_beams, num_beam_groups=num_beam_groups,
        max_length=max_length, diversity_penalty=diversity_penalty
    )

    res = para_tokenizer.batch_decode(outputs, skip_special_tokens=True)

    return res



def remove_punctuations(text):
    # Remove punctuations while preserving hyphenated words
    return re.sub(r'(?<!\w)-|-(?!\w)', ' ', re.sub(r'[^\w\s-]', '', text))

def tokenize(sentence):
    # Remove punctuations using the updated function and tokenize the sentence into words
    cleaned_sentence = remove_punctuations(sentence)
    return cleaned_sentence.split()


def generate_bigrams(words):
    # Generate bigrams from a list of words
    return [(words[i], words[i+1]) for i in range(len(words)-1)]

def hash_bigram(bigram):
    # Hash function for bigrams
    return hash(tuple(bigram))

def find_matching_words(sentence1, sentence2):
    # Tokenize the sentences
    words1 = tokenize(sentence1)
    words2 = tokenize(sentence2)

    # Generate bigrams
    bigrams1 = generate_bigrams(words1)
    bigrams2 = generate_bigrams(words2)

    # Hash bigrams of sentence 1 and store them in a set for efficient lookup
    hashed_bigrams_set = set(hash_bigram(bigram) for bigram in bigrams1)

    # Find matching words by comparing hashed bigrams of sentence 2 with the set of hashed bigrams from sentence 1
    matching_words = []
    for i, bigram in enumerate(bigrams2):
        if hash_bigram(bigram) in hashed_bigrams_set:
            word1_idx = sentence2.find(bigram[0], sum(len(word) for word in sentence2.split()[:i]))
            word2_idx = sentence2.find(bigram[1], word1_idx + len(bigram[0]))
            matching_words.append((sentence2[word1_idx:word1_idx+len(bigram[0])], sentence2[word2_idx:word2_idx+len(bigram[1])]))

    return matching_words



matching_bigrams_list = []
combined_words_list = []

for paraphrase in paraphrases:
    # Find matching words
    matching_words = find_matching_words(main_sentence, paraphrase)
    matching_bigrams_list.append(matching_words)

    def combine_matching_bigrams(matching_bigrams):
        combined_words = []
        combined_word = ""

        for i, bigram in enumerate(matching_bigrams):
            if i == 0:
                combined_word += ' '.join(bigram)
            elif bigram[0] == matching_bigrams[i-1][1]:
                combined_word += ' ' + bigram[1]
            else:
                combined_words.append(combined_word)
                combined_word = ' '.join(bigram)

        # Append the last combined word
        combined_words.append(combined_word)

        return combined_words

    # Combine matching bigrams into single words
    combined_words = combine_matching_bigrams(matching_words)
    combined_words_list.append(combined_words)

def remove_overlapping(input_set):
    sorted_set = sorted(input_set, key=len, reverse=True)
    output_set = set()
    
    for word in sorted_set:
        if not any(word in existing_word for existing_word in output_set):
            output_set.add(word)
    
    return output_set

    
def find_longest_match(string1, string2):
    # Initialize variables
    longest_match = ''
    
    # Iterate through all possible substrings of string1
    for i in range(len(string1)):
        for j in range(i + 1, len(string1) + 1):
            substring = string1[i:j]
            if ' ' + substring + ' ' in ' ' + string2 + ' ':
                if len(substring) > len(longest_match):
                    longest_match = substring
    
    return longest_match

common_substrings = set()
highlighted_text = []

for i in combined_words_list[0]:
    for j in combined_words_list[1]:
        for k in combined_words_list[2]:
            for l in combined_words_list[3]:
                for m in combined_words_list[4]:
                    matching_portion = find_longest_match(i, j)
                    matching_portion = find_longest_match(matching_portion, k)
                    matching_portion = find_longest_match(matching_portion, l)
                    matching_portion = find_longest_match(matching_portion, m)
                    if matching_portion:
                        common_substrings.add(matching_portion)




prompt_list=["The official position of the United States on the Russia-Ukraine war has been consistent in supporting Ukraine's sovereignty, territorial integrity, and the peaceful resolution of the conflict."    
,"Joe Biden said we’d not send U.S. troops to fight Russian troops in Ukraine, but we would provide robust military assistance and try to unify the Western world against Russia’s aggression."]

options = [f"Prompt #{i+1}: {prompt_list[i]}" for i in range(len(prompt_list))] + ["Another Prompt..."]
selection = st.selectbox("Choose a prompt from the dropdown below . Click on :blue['Another Prompt...'] , if you want to enter your own custom prompt.", options=options)
check=[]

if selection == "Another Prompt...": 
    check = st.text_input("Enter your custom prompt...")
    check = " " + check
    if check:
        st.caption(f""":white_check_mark: Your input prompt is : {check}""")
        st.caption(':green[Kindly hold on for a few minutes while the AI text is being generated]')
    
else:    
    check = re.split(r'#\d+:', selection, 1)[1]
    if check:
        st.caption(f""":white_check_mark: Your input prompt is : {check}""")
        st.caption(':green[Kindly hold on for a few minutes while the Paraphrase texts are being generated]')


main_sentence = check

st.markdown("**Main Sentence**:")
st.write(main_sentence)

# Generate paraphrases
paraphrases = paraphrase(main_sentence)

# Extracting longest common sequences
longest_common_sequences = find_longest_common_sequences(main_sentence, paraphrases)

color_palette = ["#FF0000", "#008000", "#0000FF", "#FF00FF", "#00FFFF"]
highlighted_sentences = []

                        
highlighted_sentence = main_sentence

for substring in remove_overlapping(common_substrings):
    highlighted_sentence = highlighted_sentence.replace(substring, colored(substring, 'white', 'on_blue'))
    highlighted_text.append(substring)

st.markdown("Common substrings that occur in all five lists:")

for substring in highlighted_text:
    st.write(substring)

st.markdown("\nHighlighted Main Sentence:")
st.markdown(highlighted_sentence)






# # Highlighting sequences in main sentence and paraphrases
# for sentence in [main_sentence] + paraphrases:
#     highlighted_sentence = sentence
#     for i, sequence in enumerate(longest_common_sequences):
#         color = color_palette[i % len(color_palette)]
#         highlighted_sentence = highlighted_sentence.replace(sequence, f"<span style='color:{color}'>{sequence}</span>")
#     highlighted_sentences.append(highlighted_sentence)

# # Display paraphrases with numbers
# st.markdown("**Paraphrases**:")
# for i, para in enumerate(paraphrases, 1):
#     st.write(f"Paraphrase {i}:")
#     st.write(para)


# # Displaying the main sentence with highlighted longest common sequences
# st.markdown("**Main sentence with highlighted longest common sequences**:")
# st.markdown(highlighted_sentences[0], unsafe_allow_html=True)


# st.markdown("**Paraphrases with highlighted longest common sequences**:")
# for paraphrase in highlighted_sentences[1:]:
#     st.markdown(paraphrase, unsafe_allow_html=True)