import gradio as gr
import pandas as pd
from dotenv import load_dotenv
from langchain_community.llms import CTransformers, HuggingFacePipeline, HuggingFaceHub
from langchain_core.prompts import PromptTemplate
from langchain.chains import LLMChain
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from sentence_transformers import SentenceTransformer, util
from sklearn.cluster import KMeans
import nltk
import pandas as pd
import smtplib
import os
nltk.download('punkt')
from nltk.tokenize import word_tokenize
from nltk import tokenize
import numpy as np
import scipy.spatial
import csv


load_dotenv()

# Global array of different possible LLM selection options
LLM_OPTIONS = [
    ("Llama-2-7B", "TheBloke/Llama-2-7B-Chat-GGML"),
    ("Falcon-180B", "TheBloke/Falcon-180B-Chat-GGUF"),
    ("Zephyr-7B", "zephyr-quiklang-3b-4k.Q4_K_M.gguf"),
    ("Vicuna-33B", "TheBloke/vicuna-33B-GGUF"),
    ("Claude2", "TheBloke/claude2-alpaca-13B-GGUF"),
    ("Alpaca-7B", "TheBloke/LeoScorpius-GreenNode-Alpaca-7B-v1-GGUF")
]

def generate_prompts(user_input):
    print("User input here")
    print(user_input)
    prompt_template = PromptTemplate(
        input_variables=["Question"],
        template=f"Just list 5 distinct and separate yet relevant question prompts for {user_input} and don't put number before any of the prompts."
    )
    config = {'max_new_tokens': 256, 'temperature': 0.7, 'context_length': 256}
    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 question in questions_list if question.strip())
    questions_list = formatted_questions.split("Question:")[1:]
    return questions_list

def answer_question(prompt, model_name):
    print("inside answer question function")
    print("prompt")
    print(prompt)
    print("")
    print("model name")
    print(model_name)
    print("")

    prompt_template = PromptTemplate(
        input_variables=["Question"],
        template=f"Please provide a concise and relevant answer for {prompt} in three sentences or less and don't put Answer in front of what you return. You are a helpful and factual assistant, do not say thank you or you are happy to assist just answer the question."
    )
    config = {'max_new_tokens': 256, 'temperature': 0.7, 'context_length': 256}
    llm = CTransformers(model=model_name,
                        config=config, 
                        threads=os.cpu_count())
    hub_chain = LLMChain(prompt=prompt_template, llm=llm)
    input_data = {"Answer the question": prompt}
    generated_answer = hub_chain.run(input_data)  
    print("generated answer")
    print(generated_answer)
    return generated_answer

def calculate_sentence_similarities(sentences_list):
    model = SentenceTransformer('all-MiniLM-L6-v2')
    embeddings_list = [model.encode(sentences) for sentences in sentences_list]
    similarity_matrices = []
    
    for i in range(len(embeddings_list)):
        for j in range(i + 1, len(embeddings_list)):
            similarity_matrix = util.pytorch_cos_sim(embeddings_list[i], embeddings_list[j]).numpy()
            similarity_matrices.append((i, j, similarity_matrix))
    
    return similarity_matrices

def highlight_similar_sentences(sentences_list, similarity_threshold):
    similarity_matrices = calculate_sentence_similarities(sentences_list)
    highlighted_sentences = [[] for _ in sentences_list]

    for (i, j, similarity_matrix) in similarity_matrices:
        for idx1 in range(similarity_matrix.shape[0]):
            for idx2 in range(similarity_matrix.shape[1]):
                similarity = similarity_matrix[idx1][idx2]
                print(f"Similarity between sentence {idx1} in paragraph {i} and sentence {idx2} in paragraph {j}: {similarity:.2f}")
                if similarity >= similarity_threshold:
                    print("Greater than sim!")
                    if (idx1, "powderblue", similarity) not in highlighted_sentences[i]:
                        highlighted_sentences[i].append((idx1, "powderblue", similarity))
                    if (idx2, "powderblue", similarity) not in highlighted_sentences[j]:
                        highlighted_sentences[j].append((idx2, "powderblue", similarity))

    for i, sentences in enumerate(sentences_list):
        highlighted = []
        for j, sentence in enumerate(sentences):
            color = "none"
            score = 0
            for idx, col, sim in highlighted_sentences[i]:
                if idx == j:
                    color = col
                    score = sim
                    break
            highlighted.append({"text": sentence, "background-color": color, "score": score})
        highlighted_sentences[i] = highlighted

    print(highlighted_sentences)
    return highlighted_sentences


def setTextVisibility(cbg, model_name_input):
    selected_prompts = cbg
    answers = [answer_question(prompt, model_name_input) for prompt in selected_prompts]

    sentences_list = [tokenize.sent_tokenize(answer) for answer in answers]
    highlighted_sentences_list = highlight_similar_sentences(sentences_list, 0.5)

    result = []
    for idx, (prompt, highlighted_sentences) in enumerate(zip(selected_prompts, highlighted_sentences_list)):
        result.append(f"<p><strong>Prompt: {prompt}</strong></p>")
        for sentence_info in highlighted_sentences:
            color = sentence_info.get('background-color', 'none')  # Read the 'color' parameter
            result.append(f"<p style='background-color: {color};'><strong>{sentence_info['text']}</strong></p>")
    
    blue_scores_list = [[info['score'] for info in highlighted_sentences if info['background-color'] == 'powderblue'] for highlighted_sentences in highlighted_sentences_list]
    blue_scores = [score for scores in blue_scores_list for score in scores]

    if blue_scores:
        overall_score = round(np.mean(blue_scores) * 100)
    else:
        overall_score = 0

    final_html = f"""<div>{''.join(result)}<div style="text-align: center; font-size: 24px; font-weight: bold;">Similarity Score: {overall_score}</div></div>"""

    print("")
    print("final html")
    print(final_html)
    return final_html



def process_inputs(file, relevance, diversity, model_name):
    # Check if file is uploaded
    if file is not None:
        # Read questions from the uploaded Excel file
        try:
            df = pd.read_excel(file, engine='openpyxl')
        except Exception as e:
            return f"Failed to read Excel file: {e}", None

        # Ensure that there is only one column in the file
        if df.shape[1] != 1:
            return "The uploaded file must contain only one column of questions.", None

        # Extract the first column
        questions_list = df.iloc[:, 0]

        # Initialize lists to store the expanded data
        expanded_questions = []
        expanded_prompts = []
        expanded_answers = []
        semantic_similarities = []

        # Generate prompts for each question and expand the data
        for question in questions_list:
            prompts = generate_prompts(question)
            expanded_questions.extend([question] * len(prompts))
            expanded_prompts.extend(prompts)

            # Generate answers for each prompt
            answers = [answer_question(prompt, model_name) for prompt in prompts]
            expanded_answers.extend(answers)

            # Calculate semantic similarity score for each answer
            similarity_scores = []
            for answer in answers:
                sentences_list = tokenize.sent_tokenize(answer)
                highlighted_sentences_list = highlight_similar_sentences([sentences_list], 0.5)
                print("highlighted sentences list")
                print(highlighted_sentences_list)

                blue_scores_list = [[info['score'] for info in highlighted_sentences if info['background-color'] == 'powderblue'] for highlighted_sentences in highlighted_sentences_list]
                blue_scores = [score for scores in blue_scores_list for score in scores]

                if blue_scores:
                    overall_score = round(np.mean(blue_scores) * 100)
                else:
                    overall_score = 0

                similarity_scores.append(overall_score)
                print("overall score")
                print(overall_score)

            # Calculate mean similarity score for each question
            question_similarity_score = np.mean(similarity_scores)
            print("question sim score")
            print(question_similarity_score)

            # Extend the list with the same score for all answers to this question
            semantic_similarities.extend([question_similarity_score] * len(prompts))

        # Combine the expanded data into a DataFrame
        output_df = pd.DataFrame({
            'Questions': expanded_questions,
            'Generated Prompts': expanded_prompts,
            'Answers': expanded_answers,
            'Semantic Similarity': semantic_similarities
        })

        # Save the DataFrame to a new Excel file
        output_file = "processed_questions.xlsx"
        output_df.to_excel(output_file, index=False)
    else:
        return "No questions provided.", None

    return "Processing complete. Download the file below.", output_file

text_list = []

def get_model_name(model_label):
    # Retrieve the model name based on the selected label
    for label, name in LLM_OPTIONS:
        if label == model_label:
            return name
    return None

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

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    
    with gr.Tab("Live Mode"):

        gr.Markdown ("## Live Mode Auditing LLMs")
        gr.Markdown("In Live Auditing Mode, you gain the ability to probe the LLM directly.")
        gr.Markdown("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") 
        with gr.Row():
            model_name_input = gr.Dropdown(choices=LLM_OPTIONS, 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="background-color: powderblue"></div>""")

        btnExec.click(setTextVisibility, inputs=[cbg, model_name_input], outputs=html_result)
        clear = gr.ClearButton(link="http://127.0.0.1:7860")


    with gr.Tab("Batch Mode"):
        
        gr.Markdown("## Batch Mode Auditing LLMs")
        gr.Markdown("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.")
        gr.Markdown("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.")
        gr.Markdown("Upon completion, please provide your email address. We will compile and send the answers to you promptly.")
            
        # llm_dropdown = gr.Dropdown(choices=[
        #         ("Llama-2-7B", "TheBloke/Llama-2-7B-Chat-GGML"), 
        #         ("Falcon-180B", "TheBloke/Falcon-180B-Chat-GGUF"), 
        #         ("Zephyr-7B", "TheBloke/zephyr-quiklang-3b-4K-GGUF"), 
        #         ("Vicuna-33B", "TheBloke/vicuna-33B-GGUF"), 
        #         ("Claude2", "TheBloke/claude2-alpaca-13B-GGUF"), 
        #         ("Alpaca-7B", "TheBloke/LeoScorpius-GreenNode-Alpaca-7B-v1-GGUF")], 
        #         label="Large Language Model")

        with gr.Row():
            model_name_batch_input = gr.Dropdown(choices=LLM_OPTIONS, label="Large Language Model")

        file_upload = gr.File(label="Upload an Excel File with Questions", file_types=[".xlsx"])
        with gr.Row():
            relevance_slider = gr.Slider(1, 100, value=70, label="Relevance", info="Choose between 0 and 100", interactive=True)
            diversity_slider = gr.Slider(1, 100, value=25, label="Diversity", info="Choose between 0 and 100", interactive=True) 

            
        submit_button = gr.Button("Submit", variant="primary", min_width=200)
        output_textbox = gr.Textbox(label="Output")
        download_button = gr.File(label="Download Processed File")

        def on_submit(file, relevance, diversity, model_name_batch_input):
            print("in on submit")
            print(model_name_batch_input)
            result, output_file = process_inputs(file, relevance, diversity, model_name_batch_input)
            return result, output_file

        submit_button.click(fn=on_submit, inputs=[file_upload, relevance_slider, diversity_slider, model_name_batch_input], outputs=[output_textbox, download_button])

# Launch the Gradio app
demo.launch()