import asyncio
import json
import time

from datasets import load_dataset

from lagent.agents import AsyncMathCoder
from lagent.agents.aggregator import InternLMToolAggregator
from lagent.llms import AsyncGPTAPI
from lagent.prompts.parsers import ToolParser

loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

interpreter_prompt = (
    'Below is a math problem. Please solve it step by step with the assistance of Python programming. Consider using Sympy or Numpy library '
    'to facilitate your derivation, calculation and equation solving. Utilize the "pi" symbol and "Rational" from Sympy '
    'for $$\pi$$ and fractions, and simplify all fractions and square roots without converting them to decimal values. '
    'Please encapsulate each generated Jupyter Python code block with tags "{begin}" and "{end}". Conclude the '
    r'final answer when observations are sufficient and encapsulate the numerical result with LaTeX syntax "\boxed{{}}" '
    'without any unit, and end your conclusion with the special token "[END]" to denote the completion of your response. '
    'Keep the following points in mind:\n'
    '- You must alternately use human and programming languages in the chain of thought;\n'
    '- The number of your reasoning steps should not exceed **three**, which means you may merge some intermediate steps when the original answer is tedious.'
)

async_llm = dict(
    type=AsyncGPTAPI,
    model='gpt-4o-2024-05-13',
    retry=50,
    key='',
    max_new_tokens=2048,
    stop_words=['</python'],
    proxies=dict(),
)
async_agent = AsyncMathCoder(
    llm=async_llm,
    output_format=ToolParser(
        tool_type='interpreter',
        template=interpreter_prompt,
        begin='<python>\n',
        end='\n</python>'),
    aggregator=InternLMToolAggregator(
        environment_role='system',
        environment_begin='<output>\n',
        environment_end='\n</output>'),
    finish_condition=lambda m: '[END]' in m.content,
)

ds = load_dataset('lighteval/MATH', split='train')
problems = [item['problem'] for item in ds.select(range(30))]

tic = time.time()
coros = [async_agent(q, session_id=i) for i, q in enumerate(problems)]
res = loop.run_until_complete(asyncio.gather(*coros))
print(time.time() - tic)
with open('tmp_6.json', 'w') as f:
    json.dump([async_agent.get_steps(i) for i in range(len(problems))],
              f,
              ensure_ascii=False,
              indent=4)