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 re

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

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

    def generate_response_oeqa_multi_token(self, msg,max_new_tokens: int = 128):
        """
        Handles multiple-choice questions where answers might have multiple tokens.
        """
        # Ensure tokenizer has proper special tokens set
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        if self.model.config.pad_token_id is None:
            self.model.config.pad_token_id = self.tokenizer.pad_token_id

        chat = [
            {"role": "user", "content": "You are a question-answering chatbot."},
            {"role": "assistant", "content": "I am ready to answer your questions. Feel free to ask anything.\n"},
            {"role": "user", "content": f"{msg}"},
        ]
        formatted_chat = self.tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)

        inputs = self.tokenizer(formatted_chat, return_tensors="pt", padding=True, truncation=True)
        input_ids = inputs.input_ids.to(self.model.device)
        attention_mask = inputs.attention_mask.to(self.model.device)

        # Generate response with proper token limits
        output = self.model.generate(
            input_ids,
            do_sample=True,
            attention_mask=attention_mask,
            eos_token_id=self.tokenizer.eos_token_id,
            pad_token_id=self.tokenizer.pad_token_id,
            temperature=0.4,
            max_new_tokens=max_new_tokens,
        )

        generated_ids = output[0]  # The generated sequence including the prompt
        generated_tokens = generated_ids[len(input_ids[0]):]  # Exclude the input_ids part
        generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)

        return generated_text


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

        for row in self.dataset:
            total_count += 1

            # Get values from row
            category = str(row["difficulty"])
            answer = row["final_answer"]
            question = row["question"]

            # Construct the prompt/message
            instruction = f"""Aşağıdaki matematik problemini verilen nihai cevap formatına uygun olacak şekilde çözün. Tüm adımları gösterdikten sonra, nihai cevabınızı sadece bir kez ve aşağıdaki kurallara uygun şekilde kutu (\\boxed{{}}) içinde verin.

Nihai Cevap için Uyulması Gereken Format Kuralları:

1. Kesirler her zaman en sade hallerinde verilmeli.
    - Matris içi kesirler: x/y biçiminde.
    - Diğer tüm kesirler: \\frac{{x}}{{y}} biçiminde.
2. Çarpma işareti (*) kullanılmamalı. Örnek: 2x yazın, 2**x* değil.
3. Birden çok değişken varsa alfabetik sıraya uyulmalı ve (x, y, z...), polinomları azalan derece sırasına göre yazılmalı.
4. Her zaman aynı gösterim biçimi kullanılmalı. Ondalık yerine kesir kullanılmalı (ör. 0.5 yerine \\frac{{1}}{{2}} ).
5. Faktörize polinomlar daima aynı faktör sırası ile verilsin; her sorguda aynı cevabı verecek şekilde tutarlılığı koruyun.
6. Nihai cevabı kutu dışında tekrar etmeyin, biçimi değiştirmeyin. Aynı soru tekrarlandığında aynı formatı ve cevabı verin.
7. Nihai cevap, tek seferde \\boxed{{...}} içinde verilmeli. Örnek: Cevap x ise, "\\boxed{{x}}".


Görev: Problemi çözün, son adımda yukarıdaki kurallara tam uyan tek bir kutu içinde nihai cevabı verin.


Çözüm:


Nihai cevap:
                        """
            prompt = f"{instruction}\n\nSoru:\n{question}\n"
            message = prompt

            # Get/format answer of the model
            model_answer = self.generate_response_oeqa_multi_token(message)
            responses.append(model_answer)
            model_answer_cleaned = re.search(r"\\boxed{([^}]*)}", model_answer)

            # Check if correct based on metric
            if answer == 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():
            correct = stats['correct']
            total = stats['total']
            calculatedAccuracy = correct / total if total > 0 else 0
            print(f"{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}