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
# from colorama import colored
import string

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, commas, and full stops
    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)
    # Also split on punctuation marks to handle cases where words are adjacent to punctuation
    return re.findall(r"[\w'-]+|[.,;!?]", cleaned_sentence)

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:
            # Check if the entire bigram exists in the sentence
            if bigram[0] in words2 and bigram[1] in words2:
                # Find the start index of the first word of the bigram in the sentence
                word1_idx = words2.index(bigram[0])
                # Find the start index of the second word of the bigram in the sentence, starting from the index following the first word
                word2_idx = words2.index(bigram[1], word1_idx + 1)
                # Append the matching words to the list
                matching_words.append((words2[word1_idx], words2[word2_idx]))

    return matching_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

def remove_spaces_before_punctuation(text):
    import string
    punctuation = string.punctuation
    result = ""
    for i, char in enumerate(text):
        if i == 0:
            result += char
        else:
            if char in punctuation and text[i-1] == " ":
                result = result[:-1] + char
            else:
                result += char
    return result


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)

st.markdown("**Paraphrase Sentences**:")
for i in paraphrases:
    st.write(i)

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 += bigram[0] + ('' if bigram[1] in string.punctuation else ' ') + bigram[1]
            elif bigram[0] == matching_bigrams[i-1][1]:
                combined_word += bigram[1] if bigram[1] in string.punctuation else ' ' + bigram[1]
            else:
                combined_words.append(combined_word.strip())
                combined_word = bigram[0] + ('' if bigram[1] in string.punctuation else ' ') + bigram[1]

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

        return combined_words


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

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)


                        

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

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


# def highlight_text(text, substrings):
#     highlighted_text = text
#     for i, substring in enumerate(substrings):
#         highlighted_text = highlighted_text.replace(substring, f'<span style="background-color: {colors[i]}; color: white;">{substring}</span>')
#     return highlighted_text

# # Assuming you have main_sentence, paraphrases, and common_substrings defined

# colors = ['blue', 'green', 'orange', 'purple', 'red']  # Different colors for each paraphrase

# # Highlight main_sentence
# highlighted_main_sentence = highlight_text(main_sentence, common_substrings[0])

# st.markdown("\nHighlighted Main Sentence:")
# st.write(highlighted_main_sentence, unsafe_allow_html=True)

# # Highlight paraphrases
# for i, (paraphrase, common_substring) in enumerate(zip(paraphrases, common_substrings[1:])):
#     highlighted_paraphrase = highlight_text(paraphrase, common_substring)
#     st.markdown(f"\nHighlighted Paraphrase {i+1}:")
#     st.write(highlighted_paraphrase, unsafe_allow_html=True)

    

# Assuming you have defined common_substrings and remove_overlapping functions

highlighted_sentence = remove_spaces_before_punctuation(" ".join(tokenize(main_sentence)))
highlighted_text = []

for substring in remove_overlapping(common_substrings):
    highlighted_sentence = highlighted_sentence.replace(substring, f'<span style="background-color: blue; color: white;">{substring}</span>')
    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 with LCS:**")
st.write(highlighted_sentence, unsafe_allow_html=True)


highlighted_sentence_list = []

# Define colors for highlighting
colors = ['blue', 'green', 'orange', 'purple', 'red']

for i in range(0, 5):
    highlighted_sentence = remove_spaces_before_punctuation(" ".join(tokenize(paraphrases[i])))
    highlighted_text = []

    # Assign a unique color to each paraphrase
    color = colors[i % len(colors)]

    # Iterate over substrings and apply the color
    for substring in combined_words_list[i]:
        highlighted_sentence = highlighted_sentence.replace(substring, f'<span style="background-color: {color}; color: white;">{substring}</span>')
        highlighted_text.append(substring)
    
    highlighted_sentence_list.append(highlighted_sentence)

st.markdown("\nHighlighted Paraphrase Sentences with LCS:")
for sentence in highlighted_sentence_list:
    st.write(sentence, unsafe_allow_html=True)



# highlighted_main_sentence = main_sentence

# # Iterate through each paraphrase and apply different colors to combined words
# for i, combined_words in enumerate(combined_words_list):
#     color = colors[i % len(colors)]  # Get color for this paraphrase
#     # Highlight combined words from this paraphrase with the corresponding color
#     for substring in combined_words:
#         highlighted_main_sentence = highlighted_main_sentence.replace(substring, f'<span style="background-color: {color}; color: white;">{substring}</span>')

# st.markdown("\nHighlighted Main Sentence with LCS from All Paraphrases:")
# st.write(highlighted_main_sentence, unsafe_allow_html=True)


# colors = ['blue', 'green', 'orange', 'purple', 'red']
# highlighted_main_sentence = main_sentence

# # Iterate through each paraphrase and apply different colors to combined words
# for i, combined_words in enumerate(combined_words_list):
#     color = colors[i % len(colors)]  # Get color for this paraphrase
#     # Highlight combined words from this paraphrase with the corresponding color
#     for substring in combined_words:
#         highlighted_main_sentence = highlighted_main_sentence.replace(substring, f'<span style="background-color: {color}; color: white;">{substring}</span>', 1)  # Add a limit of 1 to only replace the first occurrence

# st.markdown("\nHighlighted Main Sentence with LCS from All Paraphrases:")
# st.write(highlighted_main_sentence, unsafe_allow_html=True)






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