from src.deepeval.base_task import BaseTask
from collections import defaultdict
from src.deepeval.utils import accuracy, accuracy_standard_error
import ast

class TurkishGeneralKnowledgeTask(BaseTask):
    def __init__(self, model_name):
        super().__init__("metunlp/turkish_general_knowledge", model_name=model_name)

    def load_dataset_from_hf(self):
        dataset = super().load_dataset_from_hf()
        return dataset

    def evaluate(self):
        responses = []
        difficulty_results = defaultdict(lambda: {'correct': 0, 'total': 0})
        total_count = 0
        true = 0

        for row in self.dataset:
            total_count += 1
            question = row["question"]
            choices = ast.literal_eval(row["choices"]) # Convert string to list
            answer_index = row["answer"]  # Assuming it's zero-based index
            difficulty = row["difficulty"]

            # print(f"Choices: {choices}")
            # print("Type of choices:", type(choices))
            # Categorize difficulty
            if difficulty <= 3:
                category = 'easy'
            elif 3 < difficulty <= 6:
                category = 'medium'
            else:
                category = 'hard'

            # Create a multiple-choice prompt to encourage index output
            formatted_choices = "\n".join([f"{chr(65+i)}: {choice}" for i, choice in enumerate(choices)])

            instruction = ""
            message = f"{question}\nChoices:\n{formatted_choices}\n{instruction}\n"
            
            #"""Wrap the result between final_answer tags. For example: <final_answer/> letter <final_answer>.
            #"""
            model_answer = self.generate_response_mcqa_multi_token(message, choices=choices, max_new_tokens=2)
            responses.append(model_answer)
            # print(f"Correct Answer: {choices[answer_index]}")
            # print(f"Model Answer: {model_answer}")
            
            #TODO: Make the cleaning in the mcqa function
            model_answer_cleaned = model_answer.strip().replace('\n', '').replace(' ', '').upper()
            
            # Check if the answer is correct
            correct_answer_letter = chr(65 + answer_index)
            # print("Correct Answer Letter:", correct_answer_letter)

            if correct_answer_letter == model_answer_cleaned:
                true += 1
                difficulty_results[category]['correct'] += 1
            
            difficulty_results[category]['total'] += 1

        # Print results categorized by difficulty
        for category, stats in difficulty_results.items():
            calculatedAccuracy = stats['correct'] / stats['total'] if stats['total'] > 0 else 0
            print(f"{category.capitalize()} Accuracy: {calculatedAccuracy:.2%} ({stats['correct']}/{stats['total']})")

        print("Results:", responses)
        print("Overall Accuracy:", true / total_count)
        acc = accuracy(true, total_count)
        acc_stderr = accuracy_standard_error(acc, total_count)
        return {"acc": acc, "acc_stderr": acc_stderr}