from src.deepeval.base_task import BaseTask
from collections import defaultdict
from src.deepeval.utils import accuracy, accuracy_standard_error
from typing import Any
import os
import ast
import re
from datasets import load_dataset,get_dataset_split_names
HF_TOKEN=os.getenv("HF_TOKEN")

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

    def load_dataset_from_hf(self):
        dataset = super().load_dataset_from_hf()
        return dataset # dataset.select(range(min(10, len(dataset))))

    def evaluate(self) -> dict[str, Any]:
        responses = []
        difficulty_results =  defaultdict(lambda:  defaultdict(lambda: {'correct': 0, 'total': 0}))

        total_count = 0
        true = 0


        for row in self.dataset:
            total_count += 1

            # Get values from row
            category = "hard" if row["level"]== 1 else "easy" if row["level"] == 0 else None
            answer_index = row["answer"]
            correct_answer_letter = chr(65 + answer_index)
            context = row["context"]
            choices = ast.literal_eval(row["choices"])  # Convert string to list
            formatted_choices = "\n".join([f"{chr(65 + i)}: {choice}" for i, choice in enumerate(choices)])
            subset = row["idiom_type"]

            if subset == "atasözü":
                question = "Aşağıda verilen durum hangi atasözü ile en iyi ifade edilebilir?"
            elif subset == "deyim":
                question = """Verilen bağlamda "[MASKED]" ile boş bırakılan yere hangi deyim getirilirse cümlenin akışı anlamlı olur?"""
            else:
                question = "Aşağıda verilen durum hangi atasözü ile en iyi ifade edilebilir?"

            # Construct the prompt/message
            instruction = ""
            prompt = f"Soru: {question}\nBağlam: {context}\nSeçenekler:\n{formatted_choices}\n{instruction}\n"
            message = prompt

            # Get/format answer of the model
            model_answer = self.generate_response_mcqa_multi_token(message, choices=choices, max_new_tokens=2)
            responses.append(model_answer)
            model_answer_cleaned = model_answer.strip().replace('\n', '').replace(' ', '').upper().replace(':','')

            # Check if correct based on metric
            if correct_answer_letter == model_answer_cleaned:
                true += 1
                difficulty_results[subset][category]['correct'] += 1

            difficulty_results[subset][category]['total'] += 1

        # Print results categorized by difficulty
        for subset in difficulty_results.keys():
            subset_results = difficulty_results[subset]
            for category, stats in subset_results.items():
                correct = stats['correct']
                total = stats['total']
                calculatedAccuracy = correct / total if total > 0 else 0
                print(f"{subset.capitalize()} {category.capitalize()} Accuracy: {calculatedAccuracy:.2%} ({correct}/{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}