# -------------------------------------------------------------------
# OThink-R1
# Copyright (c) 2025 OPPO AgenticIR-Lab and contributors. All rights reserved.
# ------------------------------------------------------------------- 
# Licensed under the MIT license. 
# See LICENSE file in the project root for details.
# -------------------------------------------------------------------

import json
import os
import random
import re
import argparse
from openai import OpenAI
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

from datetime import datetime
import ast


KEY         = "Your-API-URL-For-LLM-Judge"
URL         = "Your-API-URL-For-LLM-Judge"
MODEL       = "The-Chosen-LLM-Judge-Model"
TEMPERATURE = 0.6

# The LLM-Judge-Client
client = OpenAI(api_key=KEY, base_url=URL)



SYSTEM_PROMPT = """
Please act as a classifier for the following task.

**Task Overview**: You will receive a question along with its correct answer. The model's thought process will be provided within `<think></think>` tags. Your job is to evaluate whether the model's thinking is necessary and classify it accordingly.

**Criteria for Classification**:

**Redundant reasoning**:
When the model has already derived the correct answer but continues to think, classify as follows:
1) **Multi-Solution Exploration**: The model has already derived the correct answer. However, the model continues to explore alternative solutions, using phrases like "Is there any chance" or "Another Way."
2) **Repeated Self-Validation**: The model has already derived the correct answer. However, the model continues to check or validates its previous intermediate reasoning steps for several times.
3) **Defensive Assumptions**: The model has already derived the correct answer.  However, the model is overly cautious, continuing to propose extraneous hypotheses to confirm the correctness of its answer, which is irrelevant to problem-solving. 

**Essential Thinking**:
The essential reasoning follows the following princeples:
1) **Key-word Identification**: The model provides concise analysis that helps quickly identify the key words, core elements of the problem.
2) **Misunderstanding Prevention**: The model conduct reasoning which aims to eliminate misunderstandings within the problem statement, thereby preventing errors stemming from incorrect assumptions. 
3) **Premise Omission Avoidance**: The model conduct reasoning which confirms all necessary premises of problems to ensure understanding completeness.


**Response Format**:
- Correct answer: (The correct answer of the question)
- Is the model's thinking essential or redundant? (yes/no)
- If essential, the reason is? 
- If redundant, the reason is?
- Thinking category: `<answer> (Redundant/Essential) </answer>`
- Classification: `<subset> class of the response </subset>`

Please analyze the model's thought process based on these criteria.
"""


def user_prompt_construct(question, gold_answer, r1response):
    return f"""
Plase assistant me to classify the types of responses generated by the model based on the provided questions and answers.

Question: {question}

Asnwer: {gold_answer}

The response of models: {r1response}
"""


def extract_subset(text):
    result_match = re.search(r'<subset>(.*?)</subset>', text, re.DOTALL)
    if result_match:
        result = result_match.group(1).strip()
        return result
    else:
        return ""
        # raise ValueError("No valid string found in the text")


def extract_result(text):
    result_match = re.search(r'<answer>(.*?)</answer>', text, re.DOTALL)
    if result_match:
        result = result_match.group(1).strip()
        return result
    else:
        return ""
        # raise ValueError("No valid string found in the text")

def get_qa(index, question, gold_answer, r1response):
    try:
        response = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt_construct(question, gold_answer, r1response)}
            ],
            temperature=TEMPERATURE,
            stream=False
        )

        final = response.choices[0].message.content

        try:
            answer = extract_result(final)
            subset = extract_subset(final)
            result_dict = {
                "index": index,
                "type": answer,
                "subset": subset,
                "classification_reason": final,
                "original": {
                    "question": question,
                    "gold_answer": gold_answer,
                    "r1response": r1response
                }
            }
            return result_dict
        except json.JSONDecodeError as e:
            print(f"Failed to decode JSON for {question}: {e}")
            return None
        except Exception as e:
            print(f"Unexpected JSON error for {question}: {e}")
            return None

    except Exception as e:
        print(f"API request failed for {question}: {e}")
        return None

def process_batch(datas, max_workers=10):
    inputs = [(data['Index'], data['question'], data['Gold_Answer'], data['R1-Response']) for data in datas]
    results = [None] * len(inputs)   
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
 
        future_to_index = {
            executor.submit(get_qa, index, question, gold_answer, r1response): i 
            for i, (index, question, gold_answer, r1response) in enumerate(inputs)
        }
        
 
        for future in tqdm(as_completed(future_to_index), total=len(inputs), desc="Processing characters"):
            original_index = future_to_index[future]
            try:
                result = future.result()
                if result:
                    results[original_index] = result   
            except Exception as e:
                print(f"Error processing index {original_index}: {e}")
    
    return [r for r in results if r is not None]

def read_data(data_path):
    data_dict = {}
 
    with open(f"{data_path}/question.log", 'r', encoding='utf-8') as file_question:
        question_content = file_question.read()
    question_blocks = question_content.split('------------------------------------------')
    for block in question_blocks:
 
        index_match = re.search(r'Index:\s*(\d+)', block)
        if not index_match:
            continue
        sample_index = int(index_match.group(1))
 
        match = re.search(r'Question: \s*([\s\S]+?)(?=\n--|-$|$)', block)
        question_text = match.group(1).strip() if match else ""

 
        data_dict.update({
            f"{sample_index}": {
                "Index": str(sample_index),
                "question": question_text
            }
        })
    
 
    with open(f"{data_path}/r1.log", 'r', encoding='utf-8') as file_r1response:
        r1response_content = file_r1response.read()
    r1response_blocks = r1response_content.split('------------------------------------------')
    for block in r1response_blocks:
 
        index_match = re.search(r'Index:\s*(\d+)', block)
        if not index_match:
            continue
        sample_index = int(index_match.group(1))
        if str(sample_index) not in data_dict:
            continue
 

        match = re.search(r'Generated text:\s*([\s\S]+?)(?=\n--|-$|$)', block)
        generated_text = match.group(1).strip() if match else ""
 
        gold_match = re.search(r'Gold solution:\s*(.*)', block)
        gold_solution = gold_match.group(1).strip() if gold_match else ""
 
        data_dict[f"{sample_index}"].update({
            "Gold_Answer": gold_solution,
            "R1-Response": generated_text
        })


    return data_dict

def main(output_file, data_path, max_workers=10):
    data_dict = read_data(data_path)
    data_list = list(data_dict.values())
    
 
    batch_size = 50  
    all_results = []
    
    for i in range(0, len(data_list), batch_size):
        batch = data_list[i:i + batch_size]
        print(f"Processing batch {i//batch_size + 1}/{(len(data_list)-1)//batch_size + 1}")
        
        batch_results = process_batch(batch, max_workers)
        all_results.extend(batch_results)
        
 
        with open(output_file, 'w', encoding='utf-8') as g:
            for result in all_results:
                json_line = json.dumps(result, ensure_ascii=False)
                g.write(json_line + '\n')
    
    print(f"Completed! Total {len(all_results)} personas generated.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Classification of Incorrect Answers Caused by Overthinking.')
    parser.add_argument('--max_workers', type=int, default=10, help='Maximum number of concurrent API requests')
    parser.add_argument('--data_path', type=str, default='./data/GSM8K-7B', help='data path')
    parser.add_argument('--output_file', type=str, 
                        default=f"log/log_classification_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log", 
                        help='data path')
    args = parser.parse_args()
    
    main(args.output_file, args.data_path, args.max_workers)