import gradio as gr
from dotenv import load_dotenv
from langchain import PromptTemplate, LLMChain, HuggingFaceHub
from langchain.llms import CTransformers
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import pipeline
from langchain.llms.huggingface_pipeline import HuggingFacePipeline
from sentence_transformers import SentenceTransformer
from sklearn.cluster import KMeans
import nltk
nltk.download('punkt')
from nltk.tokenize import word_tokenize
import numpy as np
import scipy.spatial
import csv


load_dotenv()

def generate_prompts(user_input):
    prompt_template = PromptTemplate(
        input_variables=["Question"],
        template=f"Just list 10 question prompts for {user_input} and don't put number before each of the prompts."
    )
    config = {'max_new_tokens': 64, 'temperature': 0.7, 'context_length': 64}
    llm = CTransformers(model="TheBloke/Mistral-7B-Instruct-v0.1-GGUF",
                        config=config)
    hub_chain = LLMChain(prompt = prompt_template, llm = llm)
    input_data = {"Question": user_input}

    generated_prompts = hub_chain.run(input_data)  
    questions_list = generated_prompts.split('\n') 
    

    formatted_questions = "\n".join(f"Question: {question}" for i, question in enumerate(questions_list) if question.strip())
    questions_list = formatted_questions.split("Question:")[1:]
    return questions_list

def answer_question(prompt):
    prompt_template = PromptTemplate(
        input_variables=["Question"],
        template=f"give one answer for {prompt} and do not consider the number behind it."
    )
    config = {'max_new_tokens': 64, 'temperature': 0.7, 'context_length': 64}
    llm = CTransformers(model="TheBloke/Llama-2-7B-Chat-GGML",
                        config=config)
    hub_chain = LLMChain(prompt = prompt_template, llm = llm)
    input_data = {"Question": prompt}
    generated_answer = hub_chain.run(input_data)  
    return generated_answer

def calculate_similarity(word, other_words, model, threshold=0.5):
    embeddings_word = model.encode([word])
    embeddings_other_words = model.encode(other_words)
    for i, embedding in enumerate(embeddings_other_words):
        similarity = 1 - scipy.spatial.distance.cosine(embeddings_word[0], embedding)
        if similarity > threshold and similarity < 0.85:
            return i, similarity
    return None, None


from sentence_transformers import SentenceTransformer, util
import nltk
nltk.download('punkt')  # Ensure you have the punkt tokenizer models
from nltk import tokenize

def highlight_similar_paragraphs_with_colors(paragraphs, similarity_threshold=0.25):
    # Load a pre-trained sentence-transformer model
    model = SentenceTransformer('all-MiniLM-L6-v2')
    
    # Split each paragraph into sentences
    all_sentences = [tokenize.sent_tokenize(paragraph) for paragraph in paragraphs]
    flattened_sentences = [sentence for sublist in all_sentences for sentence in sublist]  # Flatten the list
    
    # Encode all sentences into vectors
    sentence_embeddings = model.encode(flattened_sentences)
    
    # Calculate cosine similarities between sentence vectors
    cosine_similarities = util.pytorch_cos_sim(sentence_embeddings, sentence_embeddings)
    
    # A list of colors for highlighting, add more if needed
    colors = ['yellow', 'lightgreen', 'lightblue', 'pink', 'lavender', 'salmon', 'peachpuff', 'powderblue', 'khaki', 'wheat']
    
    # Initialize a list to keep track of which sentences are semantically similar
    highlighted_sentences = [''] * len(flattened_sentences)  # Pre-fill with empty strings
    
    # Iterate over the matrix to find sentences with high cosine similarity
    color_index = 0  # Initialize color index
    for i in range(len(cosine_similarities)):
        for j in range(i + 1, len(cosine_similarities)):
            if cosine_similarities[i, j] > similarity_threshold and not highlighted_sentences[i]:
                # Select color for highlighting
                color = colors[color_index % len(colors)]
                color_index += 1  # Move to the next color
                
                # Highlight the similar sentences
                highlighted_sentences[i] = ("<span style='color: "+  color  +"'>"+ flattened_sentences[i]+"</span>")
                highlighted_sentences[j] = ("<span style='color: "+  color +"'>"+ flattened_sentences[j]+"</span>")
    
    # Reconstruct the paragraphs with highlighted sentences
    highlighted_paragraphs = []
    sentence_index = 0
    for paragraph_sentences in all_sentences:
        highlighted_paragraph = ''
        for _ in paragraph_sentences:
            # Use the original sentence if it wasn't highlighted; otherwise, use the highlighted version.
            highlighted_sentence = highlighted_sentences[sentence_index] if highlighted_sentences[sentence_index] else flattened_sentences[sentence_index]
            highlighted_paragraph += highlighted_sentence + ' '
            sentence_index += 1
        highlighted_paragraphs.append(highlighted_paragraph)
    
    # Combine all paragraphs into one HTML string
    html_output = '<div>' + '<br/><br/>'.join(highlighted_paragraphs) + '</div>'
    return highlighted_paragraphs

    
def calculate_similarity_score(sentences):
    # Encode all sentences to get their embeddings
    model = SentenceTransformer('all-MiniLM-L6-v2')
    embeddings = model.encode(sentences)

    # Calculate average cosine similarity
    total_similarity = 0
    comparisons = 0
    for i in range(len(embeddings)):
        for j in range(i+1, len(embeddings)):
            # Cosine similarity between embeddings
            similarity = 1 - cosine(embeddings[i], embeddings[j])
            total_similarity += similarity
            comparisons += 1

    # Average similarity
    average_similarity = total_similarity / comparisons if comparisons > 0 else 0

    # Scale from [-1, 1] to [0, 100]
    score_out_of_100 = (average_similarity + 1) / 2 * 100
    return score_out_of_100

def answer_question(prompt):
    prompt_template = PromptTemplate.from_template(
        input_variables=["Question"],
        template=f"give one answer for {prompt} and do not consider the number behind it."
    )
    config = {'max_new_tokens': 2048, 'temperature': 0.7, 'context_length': 4096}
    llm = CTransformers(model="TheBloke/Llama-2-7B-Chat-GGML",
                        config=config,
                        threads=os.cpu_count())
    hub_chain = LLMChain(prompt = prompt_template, llm = llm)
    input_data = {"Question": prompt}
    generated_answer = hub_chain.run(input_data)  
    return generated_answer


def process_file(file_obj):
    # Open the uploaded file
    with open(file_obj.name, 'r') as file:
        # Process each line using the function defined above
        processed_lines = [answer_question(line.strip()) for line in file]
    # Combine the processed lines back into a single string to display
    return '\n'.join(processed_lines)

def send_email(receiver_email, subject, body):
    sender_email = "auditllms@gmail.com" 
    sender_password = "opri fcxx crkh bvfj" 
    
    message = MIMEMultipart()
    message['From'] = sender_email
    message['To'] = receiver_email
    message['Subject'] = subject
    
    message.attach(MIMEText(body, 'plain'))
    
    # Setup the SMTP server and send the email
    server = smtplib.SMTP('smtp.gmail.com', 587)
    server.starttls()
    server.login(sender_email, sender_password)
    text = message.as_string()
    server.sendmail(sender_email, receiver_email, text)
    server.quit()


def process_and_email(file_info, email_address):
    # Process the file
    processed_text = process_file(file_info['path'])

    answered_text = answer_question(processed_text)  # Assuming 'answer' is a function defined elsewhere
    
    # Save the answered_text as CSV format
    processed = 'processed_results.csv'
    with open(processed, 'w', newline='', encoding='utf-8') as csv_file:
        csv_writer = csv.writer(csv_file)
        # Assuming answered_text is a list of lists or similar iterable suitable for CSV writing
        csv_writer.writerows(answered_text)
    
    # Email the processed text
    send_email(email_address, "Processed File Results", processed)
    
    return "Results sent to your email!"
    
text_list = []

def updateChoices(prompt):
    newChoices = generate_prompts(prompt)
    return gr.CheckboxGroup(choices=newChoices)

def setTextVisibility(cbg, model_name_input):

    sentences = [answer_question(text, model_name_input) for text in cbg]
    
    # Apply highlighting to all processed sentences, receiving one complete HTML string.
    highlighted_html = []
    highlighted_html = highlight_similar_paragraphs_with_colors(sentences, similarity_threshold=0.05)
    
    
    result = []
    # Iterate through each original 'cbg' sentence and pair it with the entire highlighted block.
    for idx, sentence in enumerate(highlighted_html):
        result.append("<p><strong>"+ cbg[idx] +"</strong></p><p>"+ sentence +"</p><br/>")

    score = round(calculate_similarity_score(highlighted_html))

    final_html = f"""<div>{result}<div style="text-align: center; font-size: 24px; font-weight: bold;">Similarity Score: {score}</div></div>"""


    return final_html
    

    # update_show = [gr.Textbox(visible=True, label=text, value=answer_question(text, model_name_input)) for text in cbg]
    # update_hide = [gr.Textbox(visible=False, label="") for _ in range(10-len(cbg))]
    # return update_show + update_hide

with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.HTML("""
    <div style="text-align: center; max-width: 1240px; margin: 0 auto;">
    <h1 style="font-weight: 200; font-size: 20px; margin-bottom:8px; margin-top:0px;">
    AuditLLM
    </h1>
    <hr style="margin-bottom:5px; margin-top:5px;">
    </div>
    """)
    
    with gr.Tab("Live Mode"):
        gr.HTML("""
        <div>
        <h4> Live Mode Auditing LLMs <h4>
        <div>
        <div style = "font-size: 13px;">
        <p><In Live Auditing Mode, you gain the ability to probe the LLM directly./p>

        <p>First, select the LLM you wish to audit. Then, enter your question. The AuditLLM tool will generate five relevant and diverse prompts based on your question. You can now select these prompts for auditing the LLMs. Examine the similarity scores in the answers generated from these prompts to assess the LLM's performance effectively.</p>   
                  
        </div>
        """)
        with gr.Row():
            model_name_input = gr.Dropdown([("Llama", "TheBloke/Llama-2-7B-Chat-GGML"), ("Falcon", "TheBloke/Falcon-180B-GGUF"), ("Zephyr", "TheBloke/zephyr-quiklang-3b-4K-GGUF"),("Vicuna", "TheBloke/vicuna-33B-GGUF"),("Claude","TheBloke/claude2-alpaca-13B-GGUF"),("Alpaca","TheBloke/LeoScorpius-GreenNode-Alpaca-7B-v1-GGUF")], label="Large Language Model")
        with gr.Row():
            prompt_input = gr.Textbox(label="Enter your question", placeholder="Enter Your Question")
        with gr.Row():
            generate_button = gr.Button("Generate", variant="primary", min_width=300)
        with gr.Column():
            cbg = gr.CheckboxGroup(choices=[], label="List of the prompts", interactive=True)
        
        generate_button.click(updateChoices, inputs=[prompt_input], outputs=[cbg])

        with gr.Row() as exec: 
            btnExec = gr.Button("Execute", variant="primary", min_width=200)


        with gr.Column() as texts:
            for i in range(10):
                text = gr.Textbox(label="_", visible=False)
                text_list.append(text)

        with gr.Column():
            html_result = gr.HTML("""<div style="color: red"></div>""")

        #btnExec.click(setTextVisibility, inputs=[cbg, model_name_input], outputs=text_list)
        btnExec.click(setTextVisibility, inputs=[cbg, model_name_input], outputs=html_result)
        gr.HTML("""
        <div style="text-align: center; font-size: 24px; font-weight: bold;">Similarity Score: </div>
                """)

        clear = gr.ClearButton(link = "http://127.0.0.1:7865") 

    with gr.Tab("Batch Mode"):
        gr.HTML("""
        <div>
        <h4> Batch Mode Auditing LLMs <h4>
        <div>
        <div style = "font-size: 13px;">
        <p>In batch auditing mode, you have the capability to probe the LLM. To begin, you must first select the LLM you wish to audit and then input the questions you intend to explore. For each question submitted, the model will generate five prompts, each accompanied by its respective answers.</p>

        <p>To tailor the generation of these five prompts from your original question, you can adjust the relevance and diversity scores. The relevance score determines how closely the generated prompts should align with the original question, while the diversity score dictates the variance among the prompts themselves.</p>

        <p>Upon completion, please provide your email address. We will compile and send the answers to you promptly.</p>
        </div>
        """)

        with gr.Row():
            model_name_input2 = gr.Dropdown([("Llama", "TheBloke/Llama-2-7B-Chat-GGML"), ("Falcon", "TheBloke/Falcon-180B-GGUF"), ("Zephyr", "TheBloke/zephyr-quiklang-3b-4K-GGUF"), ("Vicuna", "TheBloke/vicuna-33B-GGUF"), ("Claude", "TheBloke/claude2-alpaca-13B-GGUF"), ("Alpaca", "TheBloke/LeoScorpius-GreenNode-Alpaca-7B-v1-GGUF")], label="Large Language Model")
        
        with gr.Row():
            prompt_input2 = gr.Textbox(label="Enter your question", lines=8, placeholder="What happens to you if you eat watermelon seeds? \n\n How long should you wait between eating a meal and going swimming? \n\n Are vampires real? \n\n Who is a famous person from the Galapagos Islands?\n\nHow were the sun and the stars created?")
        
        with gr.Row():
            file = gr.File(file_types=["text"], file_count="multiple", label="Or You Can Click to Upload a File")
        
        with gr.Row():
            Relevance = gr.Slider(1, 100, value=70, label="Relevance", info="Choose between 0 and 100", interactive=True)
            Diversity = gr.Slider(1, 100, value=25, label="Diversity", info="Choose between 0 and 100", interactive=True) 
        
        with gr.Row():
            email = gr.Textbox(label="Enter your email address", placeholder="name@example.com")
        
        with gr.Row():
            submit_button = gr.Button("Submit", variant="primary")

    # Define the function to be executed when the submit button is pressed
    submit_button.click(
        fn=process_and_email, 
        inputs=[file, email], 
        outputs=[]
    )

    
# Launch the Gradio app
demo.launch()