import os
import numpy as np
from matplotlib import pyplot as plt
from bioimageio_chatbot.chatbot import create_customer_service, QuestionWithHistory, UserProfile
from pydantic import BaseModel, Field
from schema_agents.schema import Message
from typing import Any, Dict, List, Optional, Union
from schema_agents.role import Role
from tvalmetrics import RagScoresCalculator
from itertools import cycle
from bs4 import BeautifulSoup
import pandas as pd
import asyncio
import yaml


dir_path = os.path.dirname(os.path.realpath(__file__))
KNOWLEDGE_BASE_PATH = "./bioimageio-knowledge-base"

class EvalInput(BaseModel):
    """Input for evaluating scores of LLM-based system."""
    question: str = Field(description="The question that was asked.")
    reference_answer: str = Field(description="The answer that was expected.")
    llm_answer: str = Field(description="The answer that was generated by the LLM-based system.")
    retrieved_context_list: Optional[List[str]] = Field(description="Retrieved context used by the LLM-based system to make answer.")
    top_k_context_list: Optional[List[str]] = Field(description="Top k contexts that would be retrieved by the LLM-based system. There's an assumption that the retrieved context list is a subset of the top k context list.")
    
class ContextConsistency(BaseModel):
    """Scores of evaluating consistency between LLM answer and retrieved context."""
    main_point_list: List[str] = Field(description="List of main points of `llm_answer`. Each main point is a string, maximum summarize 10 points.")
    main_point_derived_from_context_list: List[bool] = Field(description="List of booleans representing whether each main point in `main_point_list` was derived from context in `retrieved_context_list`.")
    
class ContextScores(BaseModel):
    """Scores of evaluating retrieval based llm answer."""
    retrieval_precision: List[bool] = Field(description="Representing the retrieval precision score. Considering `question` and `retrieved_context_list`, determine whether the context is relevant for answering the question. If the context is relevant for answering the question, respond with true. If the context is not relevant for answering the question, respond with false.")
    augmentation_accuracy: List[bool] = Field(description="Representing the augmentation accuracy score. Considering `llm_answer` and `retrieved_context_list`, determine whether the answer contains information derived from the context. If the answer contains information derived from the context, respond with true. If the answer does not contain information derived from the context, respond with false. ")
    augmentation_consistency: ContextConsistency = Field(description="Whether there is information in the `llm_answer` that does not come from the context. Summarize main points in `llm_answer`, determine whether the statement in main points can be derived from the context. If the statement can be derived from the context response with true. Otherwise response with false.")

class EvalScores(BaseModel):
    """Scores of evaluating llm answer."""
    similarity_score: float = Field(description="Float between 0 and 5 representing the similarity score, where 5 means the same and 0 means not similar, how similar in meaning is the `llm_answer` to the `reference_answer`. It should be 0 if there is factual error detected! ")
    context_scores: Optional[ContextScores] = Field(description="Scores of evaluating retrieval based llm answer.")
    
def extract_original_content(input_content):
    soup = BeautifulSoup(input_content, 'html.parser')

    # Check if the content is a table
    table = soup.find('table')
    if table:
        # Extract content from the table
        rows = table.find_all('tr')[1:]
        original_contents = []

        for row in rows:
            # Find all the cells in the row
            cells = row.find_all('td')

            # Check if there are at least two cells in the row
            if len(cells) >= 2:
                # Extract content from the second cell (index 1)
                content = cells[1].get_text(strip=True)
                original_contents.append(content)
    else:
        # Check if the content is in source code format
        code = soup.find('code')
        if code:
            # Extract content from the source code
            original_contents = [code.get_text(strip=True)]
        else:
            # Content format not recognized
            original_contents = None

    return original_contents

def load_query_answer(eval_file):
    # read Knowledge-Retrieval-Evaluation - Hoja 1 csv file
    query_answer = pd.read_csv(os.path.join(dir_path, eval_file))
    return query_answer

def create_gpt(model="gpt-3.5-turbo-1106"):
    async def respond_direct(question: str, role: Role) -> str:
        """Generate a response to a question."""
        return await role.aask(question)
        
    chatgpt = Role(
        name="GPT",
        model=model,
        profile="Customer Service",
        goal="You are responsible for answering user questions, providing clarifications. Your overarching objective is to make the user experience both educational and enjoyable.",
        constraints=None,
        actions=[respond_direct],
    )
    return chatgpt


async def generate_answers(eval_file, process, anwser_col, question_col='Question', groundtruth_col='GPT-4-turbo Answer (With Context)- GT', target_col='BioImage.IO Chatbot Answer', retrieval_col=None, eval_index=None):
    query_answer = pd.read_csv(eval_file)
    
    async def process_question(i, question):
        print(f"\n==================\nProcessing {i}th question...")
        anwser, reference = await process(question)
        query_answer.loc[i, anwser_col] = anwser
        query_answer.loc[i, anwser_col + "-reference"] = reference
        query_answer.to_csv(os.path.join(dir_path, eval_file))
        print(f"Finish processing {i}th question!")

    if eval_index is None:
        eval_index = range(len(query_answer))

    # Create a list of tasks
    tasks = [process_question(i, query_answer.iloc[i]["Question"]) for i in eval_index]
    # Run tasks concurrently
    await asyncio.gather(*tasks)
    print(f"Update {eval_file} successfully!")
    
async def generate_all_anwsers(eval_file, eval_index):
    customer_service = create_customer_service(KNOWLEDGE_BASE_PATH)
    event_bus = customer_service.get_event_bus()
    event_bus.register_default_events()
    profile = UserProfile(name="Lei", occupation="", background="")
    
    chat_history = []

    async def process_gpt35(question):
        gpt35 = create_gpt()
        anwser = await gpt35.handle(Message(content=question, role="User"))
        return anwser,  "NA"

    async def process_chatbot(question):
        m = QuestionWithHistory(question=question, chat_history=chat_history, user_profile=UserProfile.parse_obj(profile), channel_id=None)
        responses = await customer_service.handle(Message(content=m.json(), data=m, role="User"))
        response = responses[-1]
        return response.data.text,  "\n\n".join([str(step.dict()) for step in response.data.steps])

    async def process_gpt4(question):
        gpt4 = create_gpt(model="chatgpt4")
        anwser = await gpt4.handle(Message(content=question, role="User"))
        return anwser, "NA"
    
    await generate_answers(eval_file, process_chatbot, "BioImage.IO Chatbot Answer", eval_index=eval_index)
    # await generate_answers(eval_file, process_gpt35, "GPT-3.5-turbo Answer (Without Context)", eval_index=eval_index)
    # await generate_answers(eval_file, process_gpt4, "ChatGPT-4 Answer (Without Context)", eval_index=eval_index)

async def get_ground_truth(excel_file, eval_index=None):
    query_answer = pd.read_csv(excel_file)
    prompt_chatgpt = "Here is the contexts I found from the documentation: ```\n{context}\n```\nNow based on the given context, please answer the question: ```\n{question}\n```."
    chatgpt = create_gpt(model="gpt-4-1106-preview")
    event_bus = chatgpt.get_event_bus()
    event_bus.register_default_events()

    async def process_question(i):
        nonlocal query_answer, prompt_chatgpt, chatgpt
        print(f"\n==================\nGenerating GT answer for {i}th question...")
        prompt = prompt_chatgpt.format(context=query_answer.iloc[i]["Documentation"], question=query_answer.iloc[i]["Question"])
        ground_truth_answer = await chatgpt.handle(Message(content=prompt, role="User"))
        query_answer.loc[i, 'GPT-4-turbo Answer (With Context)- GT'] = ground_truth_answer[0].content
        query_answer.to_csv(os.path.join(dir_path, excel_file))
        print(f"Finish processing {i}th question!")

    if eval_index is None:
        eval_index = range(len(query_answer))

    # Create a list of tasks
    tasks = [process_question(i) for i in eval_index]

    # Run tasks concurrently
    await asyncio.gather(*tasks)

    # Save the updated DataFrame to the CSV file after all questions are processed
    query_answer.to_csv(os.path.join(dir_path, excel_file))
    print(f"Update {excel_file} successfully!")
 


async def start_evaluate_paral(eval_file, question_col='Question', groundtruth_col='GPT-4-turbo Answer (With Context)- GT', target_col='BioImage.IO Chatbot Answer', retrieval_col=None, eval_index=None):
    async def bot_answer_evaluate(req: EvalInput, role: Role) -> EvalScores:
        """Return the answer to the question."""
        response = await role.aask(req, EvalScores)
        return response
    
    evalBot = Role(
        name="Thomas",
        profile="Evaluator",
        goal="Evaluate the performance of the LLM-based system.",
        constraints=None,
        actions=[bot_answer_evaluate],
        model="gpt-4"
    )
    event_bus = evalBot.get_event_bus()
    event_bus.register_default_events()
    
    # query_answer = load_query_answer(eval_file)
    question_list = list(query_answer[question_col])
    ground_truth_answer_list = list(query_answer[groundtruth_col])
    # chatgpt_answer_list = list(query_answer['GPT-3.5-turbo Answer (Without Context)'])
    # gpt4_direct_answer_list = list(query_answer['ChatGPT-4 Answer (Without Context)'])
    target_answer_list = list(query_answer[target_col])
    
    if eval_index is None:
        eval_index = range(len(question_list))
    
    async def evaluate_question(i):
        print(f"Evaluating {i}th question...")
        question = question_list[i]
        reference_answer = ground_truth_answer_list[i]
        # gpt_direct_answer = chatgpt_answer_list[i]
        # gpt4_direct_answer = gpt4_direct_answer_list[i]
        answer = target_answer_list[i]
        if retrieval_col:
            reference_table = query_answer.iloc[i][retrieval_col]
            if pd.isna(reference_table) or reference_table == "":
                retrieved_context_list = [""]
            else:
                retrieved_context_list = extract_original_content(reference_table)
        else:
            retrieved_context_list = None
        eval_input = EvalInput(question=question, reference_answer=reference_answer, llm_answer=answer, retrieved_context_list=retrieved_context_list)
        
        try:
            scores_chatbot = await evalBot.handle(Message(content=eval_input.json(), data=eval_input, role="User"))
            
            # results = await asyncio.gather(*tasks)
            
            # scores_gpt3_direct, scores_gpt4_direct, scores_chatbot = results
            # scores_chatbot = results[0]
            
            
            SimilaryScore = scores_chatbot[0].data.similarity_score
            query_answer.loc[i, target_col +  ' - Similarity Score'] = SimilaryScore
            
            if retrieval_col:
                RetrievalPrecision = sum(scores_chatbot[0].data.context_scores.retrieval_precision) / len(scores_chatbot[0].data.context_scores.retrieval_precision)
                AugmentationAccuracy = sum(scores_chatbot[0].data.context_scores.augmentation_accuracy) / len(scores_chatbot[0].data.context_scores.augmentation_accuracy)
                AugmentationPrecision = sum(scores_chatbot[0].data.context_scores.augmentation_accuracy) / len(scores_chatbot[0].data.context_scores.augmentation_accuracy)
                AugmentationConsistency = sum(scores_chatbot[0].data.context_scores.augmentation_consistency.main_point_derived_from_context_list) / len(scores_chatbot[0].data.context_scores.augmentation_consistency.main_point_derived_from_context_list)
                query_answer.loc[i, target_col + ' - Retrieval Precision'] = RetrievalPrecision
                query_answer.loc[i, target_col + ' - Augmentation Accuracy'] = AugmentationAccuracy
                query_answer.loc[i, target_col + ' - Augmentation Precision'] = AugmentationPrecision
                query_answer.loc[i, target_col + ' - Augmentation Consistency'] = AugmentationConsistency 
            
            
                
            query_answer.to_csv(os.path.join(dir_path, eval_file))
        except Exception as e:
            print(f"Error: {e}, skipping the question {i}...") 
            
    # Create a list of tasks
    tasks = [evaluate_question(i) for i in eval_index]

    # Run tasks concurrently
    await asyncio.gather(*tasks)
    print(f"Mean of {target_col} similarity score: {query_answer[target_col + ' - Similarity Score'].mean()}")


    
if __name__ == "__main__":
    # # get the version of chatbot in pyproject.toml
    # with open("pyproject.toml", "r") as f:
    #     pyproject = yaml.safe_load(f)
    # version = pyproject["tool"]["poetry"]["version"]
    
    # Download csv from https://docs.google.com/spreadsheets/d/1E7kLdlkkEwM1Bkhn1BYWxbQf34W-3i3wujuTotiarzY/edit
    # Select sheet: function-call-18-12-2023 (version 18-12-2023)
    file_with_gt = os.path.join(dir_path, "Minimal-Eval-Test-20240111 - Minimal-Eval-Test-20240111.csv")
   
    # load query_answer
    query_answer = load_query_answer(file_with_gt)
    # remove columns named 'Unnamed: 0'
    query_answer = query_answer.loc[:, ~query_answer.columns.str.contains('^Unnamed')]
    # remove rows which 'Question' is NA
    query_answer = query_answer[query_answer['Question'].notna()]
    # remove rows which 'Question' is empty
    query_answer = query_answer[query_answer['Question'] != ""]
    # save query_answer to a new csv file
    query_answer.to_csv(os.path.join(dir_path, file_with_gt))
    print(f"Length of query_answer: {len(query_answer)}")
    asyncio.run(get_ground_truth(file_with_gt))
    # find the index of questions which 'BioImage.IO Chatbot Answer - Similarity Score' is lower than 3
    # eval_index = query_answer[query_answer['BioImage.IO Chatbot Answer - Similarity Score'] <3].index
    # eval_index=range(60)
    # # find the index of questions which 'BioImage.IO Chatbot Answer - Similarity Score' is NA
    # eval_index = query_answer[query_answer['BioImage.IO Chatbot Answer - Similarity Score'].isna()].index

    # asyncio.run(generate_all_anwsers(file_with_gt, eval_index))
    # target_cols = ['BioImage.IO Chatbot Answer', 'GPT-3.5-turbo Answer (Without Context)', 'ChatGPT-4 Answer (Without Context)']
    # asyncio.run(start_evaluate_paral(file_with_gt, 
    #                                  target_col=target_cols[0], 
    #                                  eval_index=eval_index)
    # )
    
    
    # asyncio.run(asyncio.start_evaluate_paral(file_with_gt, 
    #                                          target_col=target_cols[1], 
    #                                          eval_index=eval_index)
    # )
    # asyncio.run(asyncio.start_evaluate_paral(file_with_gt, 
    #                                          target_col=target_cols[2], 
    #                                          eval_index=eval_index)
    # )
    # query_answer = load_query_answer(file_with_gt)
    # # get the mean of similarity score
    # gpt3_5_mean = query_answer['GPT-3.5-turbo Answer (Without Context) - Answer Similarity Score'].mean()
    # chatgpt4_mean = query_answer['ChatGPT-4 Answer (Without Context) - Answer Similarity Score'].mean()
    # chatbot_mean = query_answer['BioImage.IO Chatbot Answer - Similarity Score'].mean()
    # print(f"gpt3_5_mean: {gpt3_5_mean}, chatgpt4_mean: {chatgpt4_mean}, chatbot_mean: {chatbot_mean}")
    
    
    
   