from openai import OpenAI
from tqdm import tqdm
import os
import json
import random

random.seed(42)
client = OpenAI(api_key="[API-KEY]", base_url="https://api.deepseek.com")
client_beta = OpenAI(api_key="[API-KEY]", base_url="https://api.deepseek.com/beta")

def dpsk(text):
    messages = [{"role": "user", "content": text}]
    response = client.chat.completions.create(
        model="deepseek-reasoner",
        messages=messages
    )

    reasoning_content = response.choices[0].message.reasoning_content
    content = response.choices[0].message.content
    return reasoning_content, content


def dpsk_prefix(text, prefix):
    messages = [
        {"role": "user", "content": text},
        {"role": "assistant", "content": prefix, "prefix": True}
    ]
    response = client_beta.chat.completions.create(
        model="deepseek-reasoner",
        messages=messages,
        max_tokens=10,
    )
    reasoning_content = response.choices[0].message.reasoning_content
    content = response.choices[0].message.content
    return reasoning_content, content

direct_prompts = [
    "Let me answer him without thinking more. Answer: ",
    "I will answer directly. I won't output any thinking process. Answer: ",
    "I will answer with a single number. The answer is: ",
    "I should not think, but should answer directly. The answer is: "
]

input_directorys = [
    "path/to/LongCoT/AIME/aime2025-I.jsonl",
    "path/to/LongCoT/AIME/aime2025-II.jsonl"
]
output_dir = "path/to/LongCoT/AIME/dpsk_aime_2025_answers"

os.makedirs(output_dir, exist_ok=True)

for input_directory in input_directorys:
    idx = 0
    with open(input_directory, 'r', encoding='utf-8') as infile:
        for line in tqdm(infile):
            idx += 1
            data = json.loads(line.strip())
            question_id = input_directory.split("/")[-1].split(".")[0] + f"-{idx}"

            print("SOLVING", question_id)
            if f"{question_id}.json" in os.listdir(output_dir):
                continue

            question = data.get("question")
            reasoning, answer = dpsk(question)

            data['reasoning'] = reasoning
            data['model_answer'] = answer
            direct_answers = []
            for prompt in direct_prompts:
                for _ in range(2):
                    _, content = dpsk_prefix(question, prompt)
                    direct_answers.append(content)
            data["direct_answers"] = direct_answers
            
            output_path = os.path.join(output_dir, f"{question_id}.json")
            with open(output_path, 'w', encoding='utf-8') as outfile:
                json.dump(data, outfile, ensure_ascii=False, indent=4)

