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_config_names
HF_TOKEN=os.getenv("HF_TOKEN")

class MMLUTask(BaseTask):
    def __init__(self, model_name):
        self.subsets = get_dataset_config_names("metunlp/mmlu_tr")
        print(self.subsets)
        super().__init__("metunlp/mmlu_tr", model_name=model_name)

    def load_dataset_from_hf(self):
        evaluate_count = 50
        dataset_dict = {}
        for subset in self.subsets:
            subset_data = load_dataset(self.dataset_repo, subset, token=HF_TOKEN, split="train")
            dataset_dict[subset] = subset_data.select(range(min(evaluate_count, len(subset_data))))
        return dataset_dict


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

        total_count = 0
        true = 0

        for subset in self.subsets:
            curr_dataset = self.dataset[subset]
            print(curr_dataset[0])

            for row in curr_dataset:
                total_count += 1

                # Get values from row
                question = row["question"]
                answer_index = row["answer"]
                subject = row["subject"]
                correct_answer_letter = chr(65 + answer_index)
                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)])


                # Construct the prompt/message
                instruction = f"Aşağıda {subject} konusunda çoktan seçmeli bir soru verilmiştir."
                prompt = f"{instruction}\n\nSoru: {question}\nSeçenekler:\n{formatted_choices}\n\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]['correct'] += 1

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

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