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
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`?")
    context_scores: Optional[ContextScores] = Field(description="Scores of evaluating retrieval based llm answer.")
    
    
def load_query_answer():
    # read Knowledge-Retrieval-Evaluation - Hoja 1 csv file
    query_answer = pd.read_csv(os.path.join(dir_path, "Knowledge-Retrieval-Evaluation - Hoja 1.csv"))
    return query_answer

def create_chatgpt():
    async def respond_direct(question: str, role: Role) -> str:
        """Generate a response to a question."""
        return await role.aask(question)
        
    chatgpt = Role(
        name="ChatGPT",
        model="gpt-3.5-turbo-1106",
        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 get_answers():
    query_answer = pd.read_csv(os.path.join(dir_path, "Knowledge-Retrieval-Evaluation-Results.csv"))
    customer_service = create_customer_service(KNOWLEDGE_BASE_PATH)
    chat_history = []
    # get query_answer from yaml file
    for i in range(len(query_answer)):
        question = query_answer.iloc[i]["Question"]
        # chatgpt answer
        chatgpt = create_chatgpt()
        chatgpt_answer = await chatgpt.handle(Message(content=question, role="User"))
        query_answer.loc[i, 'ChatGPT3.5 Answer (Without Context)'] = chatgpt_answer[0].content
        # BMZ chatbot
        profile = UserProfile(name="", occupation="", background="")
        m = QuestionWithHistory(question=question, chat_history=chat_history, user_profile=UserProfile.parse_obj(profile), channel_id=None)
        llm_answer = await customer_service.handle(Message(content=m.json(), data=m , role="User"))
        # add a column to query_answer, content is llm_answer[0].content
        query_answer.loc[i, 'BMZ Chatbot Answer'] = llm_answer[0].content
    # save query_answer to original csv file
    query_answer.to_csv(os.path.join(dir_path, "Knowledge-Retrieval-Evaluation-Results.csv"))
    print("Updated Knowledge-Retrieval-Evaluation-Results.csv")
    
async def start_evaluate():
    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()
    question_list = list(query_answer['Question'])
    reference_answer_list = list(query_answer['ChatGPT3.5 Answer (With Context)'])
    chatgpt_answer_list = list(query_answer['ChatGPT3.5 Answer (Without Context)'])
    BMZ_chatbot_answer_list = list(query_answer['BMZ Chatbot Answer'])
    
    
    for i in range(len(question_list)):
        question = question_list[i]
        reference_answer = reference_answer_list[i]
        gpt_direct_answer = chatgpt_answer_list[i]
        chatbot_answer = BMZ_chatbot_answer_list[i]
        retrieved_context= query_answer.iloc[i]["Documentation"]
        # read text in retrieved_context, and split it into a list of contexts by []
        retrieved_context_list = retrieved_context.split('[')
        
        
        # gpt3direct 
        eval_input_gpt3_direct = EvalInput(question=question, reference_answer=reference_answer, llm_answer=gpt_direct_answer)
        scores_gpt3_direct = await evalBot.handle(Message(content= eval_input_gpt3_direct.json(), data= eval_input_gpt3_direct, role="User"))
        
        # chatbot answer
        eval_input_chatbot = EvalInput(question=question, reference_answer=reference_answer, llm_answer=chatbot_answer, retrieved_context_list=retrieved_context_list)
        scores_chatbot = await evalBot.handle(Message(content= eval_input_chatbot.json(), data= eval_input_chatbot, role="User"))
        SimilaryScore = scores_chatbot[0].data.similarity_score
        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)
        # save scores to a new dataframe
        query_answer.loc[i, 'ChatGPT-3.5 Direct Answer Similarity Score'] = scores_gpt3_direct[0].data.similarity_score
        query_answer.loc[i, 'Chatbot Answer Similarity Score'] = SimilaryScore
        query_answer.loc[i, 'Chatbot Answer Retrieval Precision'] = RetrievalPrecision
        query_answer.loc[i, 'Chatbot Answer Augmentation Accuracy'] = AugmentationAccuracy
        query_answer.loc[i, 'Chatbot Answer Augmentation Precision'] = AugmentationPrecision
        query_answer.loc[i, 'Chatbot Answer Augmentation Consistency'] = AugmentationConsistency
        
    # save query_answer to original csv file
    query_answer.to_csv(os.path.join(dir_path, "Knowledge-Retrieval-Evaluation-Results.csv"))
       
if __name__ == "__main__":
    
    asyncio.run(start_evaluate())
    