from biomni.task.base_task import base_task
import pandas as pd
import numpy as np
np.random.seed(42)

def shuffle(x):
    np.random.shuffle(x)
    return x

class humanity_last_exam(base_task):

    def __init__(self, path = './data', category = 'Biology/Medicine', answer_type = 'multipleChoice'):
        if category not in ['Other', 'Humanities/Social Science', 'Math', 'Physics', 'Computer Science/AI', 'Biology/Medicine', 'Chemistry', 'Engineering']:
            raise ValueError("category must be one of ['Other', 'Humanities/Social Science', 'Math', 'Physics', 'Computer Science/AI', 'Biology/Medicine', 'Chemistry', 'Engineering']")
        if answer_type not in ['exactMatch', 'multipleChoice']:
            raise ValueError("answer_type must be one of ['exactMatch' or 'multipleChoice']")
        
        self.dataset = category  # Store dataset type
        self.answer_type = answer_type
        df = pd.read_parquet(path + '/hle/test_sampled_biology_medicine.parquet')

        # Extract answer choices from question text
        def extract_options(question):
            # Find the "Answer Choices:" section
            if "Answer Choices:" not in question:
                return []
            choices = question.split("Answer Choices:")[1].strip()
            
            # Split on A., B., C. etc and clean up
            options = []
            letters = ['A.', 'B.', 'C.', 'D.', 'E.', 'F.', 'G.', 'H.', 'I.', 'J.', 'K.', 'L.', 'M.', 'N.', 'O.', 'P.', 'Q.', 'R.', 'S.', 'T.', 'U.', 'V.', 'W.', 'X.', 'Y.', 'Z.']
            for i, letter in enumerate(letters):
                if letter in choices:
                    # Define the next letter if available
                    next_letter = letters[i+1] if i+1 < len(letters) else None
                    
                    # Split between current letter and next letter
                    parts = choices.split(letter)[1]
                    if next_letter and next_letter in parts:
                        option = parts.split(next_letter)[0].strip()
                    else:
                        option = parts.strip()
                        
                    options.append(option)
            return options
        
        def extract_question(question):
            return question.split("Answer Choices:")[0].strip()

        # Extract options and answers only for multiple choice questions and category is the same as the dataset
        df = df[df['category'] == self.dataset]
        df = df[df['answer_type'] == 'multipleChoice']
        df['question_text'] = df.question
        df['letter_answer'] = df['answer'].apply(lambda x: x[0])

        self.query = df.question_text.values
        # self.options = df.options_letters.values
        self.answer = df.letter_answer.values
        
        self.prompt = """Question: {question}"""
        
    def get_example(self, index = None):
        if index is None:
            index = np.random.randint(len(self.query))


        return {"prompt": self.prompt.format(
            question = self.query[index], 
            # options = self.options[index]
        ), "answer": self.answer[index]}

    def get_iterator(self): 
        for i in range(len(self.query)):
            yield self.get_example(i)

    def evaluate(self, response):
        ## expected a list/array of symbols
        from sklearn.metrics import accuracy_score
        ground_truth = self.answer
        response = np.array(response)

        return {
            'accuracy': accuracy_score(ground_truth, response),
            'coverage': np.mean(response != self.refrain_label),
            'refrain_ratio': np.mean(response == self.refrain_label),
            'precision': accuracy_score(ground_truth[np.where(response != self.refrain_label)], response[np.where(response != self.refrain_label)]),
        }

    def output_class(self):
        from pydantic import BaseModel, Field
        from typing import Optional
        class MultipleChoiceOutput(BaseModel):
            """Multiple choice output."""

            choice: Optional[str] = Field(
                description="Multiple choice answer. For example, if there is <answer>A</answer> in the prompt, the output should be 'A'."
            )
        return MultipleChoiceOutput