import random

import datasets

from .math_agent import MathAgent

raw_dataset = datasets.load_dataset("BytedTsinghua-SIA/DAPO-Math-17k", split="train")
TRAIN_SIZE = 17917 - 1024
TEST_SIZE = 1024

train_dataset = raw_dataset.select(range(TRAIN_SIZE))
test_dataset = raw_dataset.select(range(TRAIN_SIZE, TRAIN_SIZE + TEST_SIZE))


class DAPOAgent(MathAgent):
    env_id: str = "dapo"

    def reformat_datum(self, datum: dict) -> dict:
        return {
            "problem": datum['prompt'][0]['content']
            .replace(
                'The last line of your response should be of the form Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\n',
                '',
            )
            .replace('\nRemember to put your answer on its own line after "Answer:".', ''),
            "answer": datum["reward_model"]["ground_truth"],
            "problem_id": datum["extra_info"]["index"],
        }

    def get_dataset(self, validation: bool = False):
        return train_dataset if not validation else test_dataset

    async def evaluation_prompts(
        self, num_prompts: int, validation: bool = False
    ) -> list[tuple[str, dict]]:
        dataset = self.get_dataset(validation)
        prompts = []
        for i, golden in [(i, dataset[i]) for i in range(num_prompts)]:
            golden = self.reformat_datum(golden)
            prompts.append((self.make_prefix(**golden), golden))
        return prompts

    async def get_prompt(self, validation=False) -> tuple[str, dict]:
        dataset = self.get_dataset(validation)
        golden = dataset[random.randrange(len(dataset))]
        golden = self.reformat_datum(golden)
        prompt = self.make_prefix(**golden)
        return prompt, golden

    async def get_reward(self, response, golden: dict) -> float:
        return self.compute_score(response, golden, golden_key="answer")