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 ast


class ComplexReasoningTask(BaseTask):
    def __init__(self, model_name):
        super().__init__("metunlp/complex-ales", model_name=model_name)

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


    def evaluate(self) -> dict[str, Any]:
        responses = []
        correct_answers = []

        total_count = 0
        true = 0

        for row in self.dataset:
            total_count += 1

            # Get values from row
            choices = ast.literal_eval(row["choices"]) # Convert string to list
            narrative = row["narrative"]
            question = row["question"]
            formatted_choices = "\n".join([f"{chr(65+i)}: {choice}" for i, choice in enumerate(choices)])
            correct_answer_letter = row["answer_choice"]
            correct_answers.append(correct_answer_letter)

            # Prints for debugging
            # print(f"Choices: {choices}")
            # print("Type of choices:", type(choices))


            # Construct the prompt/message
            instruction = ""
            prompt = f"Soru:\n{narrative}\n{question}\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(':','')

            if correct_answer_letter == model_answer_cleaned:
                true += 1
            # Print answers
            # print(f"Correct Answer: {correct_answer_letter}")
            # print(f"Model Answer: {model_answer}")
            # print(f"Model Answer Cleaned: {model_answer_cleaned}")

        print("Answers:", correct_answers)
        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}