import argparse
import json
import os
import logging
import re
import time
import math
import numpy as np
from typing import List, Dict, Tuple, Optional, Union, Counter
from datetime import datetime
from openai import OpenAI
from datasets import load_dataset
from tqdm import tqdm
import statistics
from collections import defaultdict

# Add sys path to import optillm modules
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from optillm.utils.answer_extraction import extract_answer as unified_extract_answer

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize OpenAI client
# client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"), base_url="https://openrouter.ai/api/v1")

client = OpenAI(api_key="optillm", base_url="http://localhost:8001/v1")

SYSTEM_PROMPT = '''You are solving AIME (American Invitational Mathematics Examination) problems.

Important: Always end your solution with the final answer in one of these two formats:

1. \\[
   \\boxed{X}.
   \\]

2. $n=\\boxed{X}$

where X is your integer answer between 0 and 999.'''

# Define the thought transition phrases to track
THOUGHT_TRANSITIONS = [
    "Wait,", 
    "Alternatively,", 
    "However,", 
    "Additionally,"
]

def load_2024_dataset() -> list[dict]:
    """
    Load the 2024 dataset of problems.
    Returns:
        list[dict]: The dataset of problems.
    """
    dataset_original = load_dataset("AI-MO/aimo-validation-aime")
    # Filter out problems that are not from 2024
    dataset = dataset_original["train"].filter(lambda example: "2024" in example["url"])
    logging.debug(f"Filtered dataset size: {len(dataset)}.")
    assert len(dataset) == 30, f"Expected 30 problems after filtering by 2024, but found {len(dataset)}"
    return dataset

def load_2025_dataset() -> list[dict]:
    """
    Load the 2025 dataset of problems from math-ai/aime25.
    Returns:
        list[dict]: The dataset of problems.
    """
    dataset = load_dataset("math-ai/aime25")
    # The AIME 2025 dataset has 30 problems in the "test" split
    dataset = dataset["test"]
    logging.debug(f"Loaded AIME 2025 dataset size: {len(dataset)}.")
    assert len(dataset) == 30, f"Expected 30 problems in AIME 2025, but found {len(dataset)}"
    return dataset

def load_dataset_by_year(year: int) -> list[dict]:
    """
    Load dataset by year (2024 or 2025).
    Returns:
        list[dict]: The dataset of problems.
    """
    if year == 2024:
        return load_2024_dataset()
    elif year == 2025:
        return load_2025_dataset()
    else:
        raise ValueError(f"Unsupported year: {year}. Only 2024 and 2025 are supported.")

def extract_answer(response: str) -> Optional[int]:
    """
    Extract the numerical answer from a math solution response using unified extraction.
    AIME problems expect integer answers between 0 and 999.
    """
    if not response:
        return None

    # Use unified answer extraction with AIME problem context
    extracted_answer = unified_extract_answer(
        response,
        problem_type="aime",
        problem_id=None
    )

    if extracted_answer is None:
        return None

    # Math-verify returns a list of all possible matches
    # Check if extracted_answer is a list and find first valid integer
    if isinstance(extracted_answer, list):
        for item in extracted_answer:
            if isinstance(item, (int, float)):
                answer = int(item)
                if 0 <= answer <= 999:
                    return answer
            elif isinstance(item, str) and item.isdigit():
                answer = int(item)
                if 0 <= answer <= 999:
                    return answer
        return None

    # Convert to integer if needed - AIME answers are always integers
    if isinstance(extracted_answer, (int, float)):
        answer = int(extracted_answer)
        # AIME answers are typically 0-999
        if 0 <= answer <= 999:
            return answer
    elif isinstance(extracted_answer, str) and extracted_answer.isdigit():
        answer = int(extracted_answer)
        if 0 <= answer <= 999:
            return answer

    return None

def analyze_thinking(response: str) -> Dict:
    """
    Analyze thinking patterns in the response.
    Extract tokens between <think> and </think> tags and count thought transitions.
    
    Args:
        response (str): The model's response text
        
    Returns:
        Dict: Analysis metrics including thinking tokens and thought transitions
    """
    # Default result with zero values
    result = {
        "has_think_tags": False,
        "thinking_tokens": 0,
        "thinking_tokens_text": "",
        "total_tokens": len(response.split()),
        "thought_transitions": 0,
        "transition_counts": {phrase: 0 for phrase in THOUGHT_TRANSITIONS},
        "transition_positions": []
    }
    
    # Extract content between <think> and </think> tags
    think_pattern = re.compile(r'<think>(.*?)</think>', re.DOTALL)
    think_match = think_pattern.search(response)
    
    if think_match:
        thinking_text = think_match.group(1)
        result["has_think_tags"] = True
        result["thinking_tokens"] = len(thinking_text.split())
        result["thinking_tokens_text"] = thinking_text
        
        # Count thought transitions
        position = 0
        for phrase in THOUGHT_TRANSITIONS:
            # Find all occurrences of each transition phrase
            for match in re.finditer(r'\b' + re.escape(phrase) + r'\b', thinking_text):
                result["transition_counts"][phrase] += 1
                # Record the approximate token position of the transition
                token_position = len(thinking_text[:match.start()].split())
                result["transition_positions"].append((phrase, token_position))
                
        # Sort transition positions by token position
        result["transition_positions"].sort(key=lambda x: x[1])
        
        # Calculate total transitions
        result["thought_transitions"] = sum(result["transition_counts"].values())
    
    return result

def analyze_logits_probs(logprobs_data: List[Dict]) -> Dict:
    """
    Analyze token probability distributions and entropy patterns.
    
    Args:
        logprobs_data: List of dictionaries containing token and logprob information
        
    Returns:
        Dict: Analysis metrics including entropy statistics
    """
    if not logprobs_data:
        return {
            "entropy_stats": None,
            "transition_entropy": None,
            "token_count": 0
        }
    
    token_entropies = []
    token_probs = []
    token_texts = []
    
    # Process each token's logprobs
    for token_info in logprobs_data:
        if not token_info.get("top_logprobs"):
            continue
        
        # Extract probabilities from logprobs
        probs = []
        for token, logprob in token_info["top_logprobs"].items():
            probs.append(math.exp(logprob))
        
        # Normalize probabilities to sum to 1
        total_prob = sum(probs)
        if total_prob > 0:
            probs = [p/total_prob for p in probs]
        
        # Calculate entropy: -sum(p_i * log(p_i))
        entropy = -sum(p * math.log2(p) if p > 0 else 0 for p in probs)
        token_entropies.append(entropy)
        token_probs.append(probs[0] if probs else 0)  # Store top token probability
        token_texts.append(token_info["token"])
    
    # Analyze entropy changes around thought transitions
    transition_entropy = {}
    
    for phrase in THOUGHT_TRANSITIONS:
        # Find indices where this transition phrase begins
        transition_indices = []
        
        # Simple approach: find where token texts match the start of the phrase
        for i, token in enumerate(token_texts):
            if phrase.startswith(token) and i < len(token_texts) - 1:
                # Check if this could be the start of the transition phrase
                # This is a simplification; more complex matching would require full tokenization
                transition_indices.append(i)
        
        # Analyze entropy changes around transitions
        if transition_indices:
            before_entropy = []
            after_entropy = []
            
            for idx in transition_indices:
                # Look at 5 tokens before and after transition
                before_window = max(0, idx-5)
                after_window = min(len(token_entropies), idx+5)
                
                if idx > before_window:
                    before_entropy.extend(token_entropies[before_window:idx])
                if after_window > idx:
                    after_entropy.extend(token_entropies[idx:after_window])
            
            transition_entropy[phrase] = {
                "before_mean": statistics.mean(before_entropy) if before_entropy else 0,
                "after_mean": statistics.mean(after_entropy) if after_entropy else 0,
                "count": len(transition_indices)
            }
    
    # Calculate overall entropy statistics
    entropy_stats = {
        "mean": statistics.mean(token_entropies) if token_entropies else 0,
        "median": statistics.median(token_entropies) if token_entropies else 0,
        "max": max(token_entropies) if token_entropies else 0,
        "min": min(token_entropies) if token_entropies else 0,
        "std": statistics.stdev(token_entropies) if len(token_entropies) > 1 else 0
    }
    
    # Calculate entropy per quartile of generation
    if token_entropies:
        quartile_size = max(1, len(token_entropies) // 4)
        entropy_stats["quartiles"] = [
            statistics.mean(token_entropies[i:i+quartile_size])
            for i in range(0, len(token_entropies), quartile_size)
            if i < len(token_entropies)
        ]
    else:
        entropy_stats["quartiles"] = []
    
    return {
        "entropy_stats": entropy_stats,
        "transition_entropy": transition_entropy,
        "token_count": len(token_entropies)
    }

def get_llm_response(problem: str, model: str, analyze_logits: bool = False, extra_body: dict = None) -> Union[str, List[Dict]]:
    """
    Get response from the LLM for a given problem.
    If multiple choices are returned, formats them as attempt dictionaries.
    
    Args:
        problem (str): The problem text
        model (str): The model identifier
        analyze_logits (bool): Whether to request logprobs
        
    Returns:
        Union[str, List[Dict]]: Either a string response or list of attempt dictionaries
    """
    try:
        # Add logprobs parameters if requested
        kwargs = {}
        if analyze_logits:
            kwargs["logprobs"] = True
            kwargs["top_logprobs"] = 3
        
        # Add extra_body if provided
        if extra_body:
            kwargs["extra_body"] = extra_body
        
        response = client.with_options(timeout=6000.0).chat.completions.create(
            model=model,
            messages=[
                {"role": "user", "content": SYSTEM_PROMPT + problem}
            ],
            max_tokens=64000,
            **kwargs
        )
        
        # Save raw response if logprobs are requested
        if analyze_logits:
            raw_filename = f"results/raw_responses_{model.replace('/', '_')}.json"
            problem_id = hash(problem) % 10000  # Simple hash to identify the problem
            save_raw_response(raw_filename, problem_id, response.model_dump())
        
        # If there's more than one choice, format as attempts
        if len(response.choices) > 1:
            attempts = []
            for i, choice in enumerate(response.choices):
                response_text = choice.message.content.strip()
                predicted_answer = extract_answer(response_text)
                attempt_data = {
                    "attempt_number": i + 1,
                    "response": response_text,
                    "predicted_answer": predicted_answer
                }
                
                # Add logprobs if available
                if analyze_logits and hasattr(choice.message, 'logprobs') and choice.message.logprobs:
                    attempt_data["logprobs"] = choice.message.logprobs
                
                attempts.append(attempt_data)
            return attempts
            
        # If single choice, return as before
        response_text = response.choices[0].message.content.strip()
        
        # If analyzing logits, return as a dictionary with logprobs
        if analyze_logits and hasattr(response.choices[0].message, 'logprobs') and response.choices[0].message.logprobs:
            return {
                "response": response_text,
                "logprobs": response.choices[0].message.logprobs
            }
        
        # Otherwise return just the text
        return response_text
        
    except Exception as e:
        logger.error(f"Error getting LLM response: {e}")
        logger.error(f"Error type: {type(e).__name__}")
        if "timeout" in str(e).lower():
            logger.error("API call timed out - consider increasing timeout for complex approaches like MARS")
        raise e  # Re-raise instead of silently returning empty string

def make_n_attempts(problem: str, model: str, n: int, analyze_thoughts: bool = False, analyze_logits: bool = False, extra_body: dict = None) -> List[Dict]:
    """
    Make n attempts to solve a problem and return all responses and predictions.
    
    Args:
        problem (str): The problem text
        model (str): The model identifier
        n (int): Number of attempts to make
        analyze_thoughts (bool): Whether to analyze thinking patterns
        analyze_logits (bool): Whether to analyze token probabilities
        
    Returns:
        List[Dict]: List of dictionaries containing response and predicted answer for each attempt
    """
    attempts = []
    remaining_attempts = n
    
    while remaining_attempts > 0:
        try:
            response = get_llm_response(problem, model, analyze_logits, extra_body)
        except Exception as e:
            logger.error(f"Failed to get response for attempt {n - remaining_attempts + 1}: {e}")
            # Create a failed attempt record
            attempt_data = {
                "attempt_number": len(attempts) + 1,
                "response": f"ERROR: {str(e)}",
                "predicted_answer": None,
                "error": str(e)
            }
            attempts.append(attempt_data)
            remaining_attempts -= 1
            continue
        
        # If response is already formatted as attempts
        if isinstance(response, list):
            for attempt in response:
                if analyze_thoughts:
                    attempt["thought_analysis"] = analyze_thinking(attempt["response"])
                if analyze_logits and "logprobs" in attempt:
                    attempt["logit_analysis"] = analyze_logits_probs(attempt["logprobs"]["content"])
            attempts.extend(response)
            remaining_attempts = n - len(attempts)
        elif isinstance(response, dict) and "response" in response:
            # Process dict response with logprobs
            response_text = response["response"]
            predicted_answer = extract_answer(response_text)
            attempt_data = {
                "attempt_number": len(attempts) + 1,
                "response": response_text,
                "predicted_answer": predicted_answer
            }
            if analyze_thoughts:
                attempt_data["thought_analysis"] = analyze_thinking(response_text)
            if analyze_logits and "logprobs" in response:
                attempt_data["logit_analysis"] = analyze_logits_probs(response["logprobs"]["content"])
            attempts.append(attempt_data)
            remaining_attempts -= 1
        else:
            # Process simple string response
            predicted_answer = extract_answer(response)
            attempt_data = {
                "attempt_number": len(attempts) + 1,
                "response": response,
                "predicted_answer": predicted_answer
            }
            if analyze_thoughts:
                attempt_data["thought_analysis"] = analyze_thinking(response)
            attempts.append(attempt_data)
            remaining_attempts -= 1
    
    return attempts

def evaluate_pass_at_n(attempts: List[Dict], correct_answer: int) -> Tuple[bool, Optional[int]]:
    """
    Evaluate if any of the n attempts got the correct answer.
    
    Args:
        attempts (List[Dict]): List of attempt results
        correct_answer (int): The correct answer
        
    Returns:
        Tuple[bool, Optional[int]]: (whether any attempt was correct, first correct attempt number)
    """
    for attempt in attempts:
        if attempt["predicted_answer"] == correct_answer:
            return True, attempt["attempt_number"]
    return False, None

def load_existing_results(filename: str) -> List[Dict]:
    """Load existing results from file if it exists."""
    try:
        with open(filename, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        return []

def save_result(filename: str, result: Dict):
    """Save a single result to the results file."""
    results = load_existing_results(filename)
    results.append(result)
    with open(filename, 'w') as f:
        json.dump(results, f, indent=2)

def get_last_processed_index(results: List[Dict]) -> int:
    """Get the index of the last processed problem."""
    if not results:
        return -1
    return max(int(r.get('index', -1)) for r in results)

def analyze_results(results: List[Dict], n: int, analyze_thoughts: bool = False, analyze_logits: bool = False):
    """
    Analyze and print summary statistics of the results.
    
    Args:
        results (List[Dict]): List of evaluation results
        n (int): Number of attempts per problem
        analyze_thoughts (bool): Whether to analyze thinking patterns
        analyze_logits (bool): Whether to analyze token probabilities
    """
    total = len(results)
    correct = sum(1 for r in results if r['is_correct'])
    accuracy = correct / total if total > 0 else 0
    
    print("\n=== Results Summary ===")
    print(f"Evaluation mode: pass@{n}")
    print(f"Total problems: {total}")
    print(f"Correct answers: {correct}")
    print(f"Accuracy: {accuracy:.2%}")
    
    # Calculate attempt statistics
    successful_attempts = [r['first_correct_attempt'] for r in results if r['is_correct']]
    if successful_attempts:
        avg_attempts = sum(successful_attempts) / len(successful_attempts)
        print(f"\nFor correct solutions:")
        print(f"Average attempts needed: {avg_attempts:.2f}")
        print(f"Attempt distribution:")
        for i in range(1, n + 1):
            count = sum(1 for x in successful_attempts if x == i)
            print(f"  Attempt {i}: {count} problems")
    
    if analyze_thoughts:
        print("\n=== Thinking Pattern Analysis ===")
        
        # Collect metrics about thinking patterns for correct vs incorrect attempts
        correct_attempts = []
        incorrect_attempts = []
        
        for result in results:
            for attempt in result['attempts']:
                if 'thought_analysis' in attempt:
                    if result['is_correct'] and attempt['predicted_answer'] == result['correct_answer']:
                        correct_attempts.append(attempt)
                    else:
                        incorrect_attempts.append(attempt)
        
        # Function to calculate statistics for a group of attempts
        def calc_stats(attempts):
            if not attempts:
                return {
                    "count": 0,
                    "avg_thinking_tokens": 0,
                    "avg_thought_transitions": 0,
                    "transition_usage": {phrase: 0 for phrase in THOUGHT_TRANSITIONS},
                    "has_think_tags_pct": 0
                }
            
            thinking_tokens = [a['thought_analysis']['thinking_tokens'] for a in attempts]
            thought_transitions = [a['thought_analysis']['thought_transitions'] for a in attempts]
            has_think_tags = sum(1 for a in attempts if a['thought_analysis']['has_think_tags'])
            
            # Count total transition usage
            transition_usage = defaultdict(int)
            for attempt in attempts:
                for phrase, count in attempt['thought_analysis']['transition_counts'].items():
                    transition_usage[phrase] += count
            
            return {
                "count": len(attempts),
                "avg_thinking_tokens": statistics.mean(thinking_tokens) if thinking_tokens else 0,
                "median_thinking_tokens": statistics.median(thinking_tokens) if thinking_tokens else 0,
                "min_thinking_tokens": min(thinking_tokens) if thinking_tokens else 0,
                "max_thinking_tokens": max(thinking_tokens) if thinking_tokens else 0,
                "avg_thought_transitions": statistics.mean(thought_transitions) if thought_transitions else 0,
                "median_thought_transitions": statistics.median(thought_transitions) if thought_transitions else 0,
                "transition_usage": dict(transition_usage),
                "has_think_tags_pct": (has_think_tags / len(attempts)) * 100 if attempts else 0
            }
        
        # Calculate statistics
        correct_stats = calc_stats(correct_attempts)
        incorrect_stats = calc_stats(incorrect_attempts)
        all_stats = calc_stats(correct_attempts + incorrect_attempts)
        
        # Print statistics
        print(f"\nOverall Thinking Statistics (All {all_stats['count']} Attempts):")
        print(f"- Average thinking tokens: {all_stats['avg_thinking_tokens']:.2f}")
        print(f"- Median thinking tokens: {all_stats['median_thinking_tokens']}")
        print(f"- Range: {all_stats['min_thinking_tokens']} - {all_stats['max_thinking_tokens']} tokens")
        print(f"- Average thought transitions: {all_stats['avg_thought_transitions']:.2f}")
        print(f"- Median thought transitions: {all_stats['median_thought_transitions']}")
        print(f"- Percentage with <think> tags: {all_stats['has_think_tags_pct']:.2f}%")
        print(f"- Transition phrase usage:")
        for phrase, count in all_stats['transition_usage'].items():
            print(f"  - {phrase}: {count} occurrences")
        
        print(f"\nCorrect Attempts ({correct_stats['count']}):")
        print(f"- Average thinking tokens: {correct_stats['avg_thinking_tokens']:.2f}")
        print(f"- Median thinking tokens: {correct_stats['median_thinking_tokens']}")
        print(f"- Average thought transitions: {correct_stats['avg_thought_transitions']:.2f}")
        print(f"- Median thought transitions: {correct_stats['median_thought_transitions']}")
        print(f"- Percentage with <think> tags: {correct_stats['has_think_tags_pct']:.2f}%")
        print(f"- Transition phrase usage:")
        for phrase, count in correct_stats['transition_usage'].items():
            print(f"  - {phrase}: {count} occurrences")
        
        print(f"\nIncorrect Attempts ({incorrect_stats['count']}):")
        print(f"- Average thinking tokens: {incorrect_stats['avg_thinking_tokens']:.2f}")
        print(f"- Median thinking tokens: {incorrect_stats['median_thinking_tokens']}")
        print(f"- Average thought transitions: {incorrect_stats['avg_thought_transitions']:.2f}")
        print(f"- Median thought transitions: {incorrect_stats['median_thought_transitions']}")
        print(f"- Percentage with <think> tags: {incorrect_stats['has_think_tags_pct']:.2f}%")
        print(f"- Transition phrase usage:")
        for phrase, count in incorrect_stats['transition_usage'].items():
            print(f"  - {phrase}: {count} occurrences")
        
        # Calculate correlation between thinking tokens and correctness
        if correct_attempts and incorrect_attempts:
            print("\nCorrelation Analysis:")
            
            # Find problems with both correct and incorrect attempts for comparison
            problems_with_both = defaultdict(lambda: {"correct": [], "incorrect": []})
            
            for result in results:
                problem_id = result['index']
                for attempt in result['attempts']:
                    if 'thought_analysis' in attempt:
                        category = "correct" if attempt['predicted_answer'] == result['correct_answer'] else "incorrect"
                        problems_with_both[problem_id][category].append(attempt)
            
            # Filter to problems that have both correct and incorrect attempts
            valid_problems = {
                k: v for k, v in problems_with_both.items() 
                if v["correct"] and v["incorrect"]
            }
            
            if valid_problems:
                print(f"Found {len(valid_problems)} problems with both correct and incorrect attempts")
                
                # For each problem, compare thinking patterns
                avg_token_diff = []
                avg_transition_diff = []
                
                for problem_id, attempts in valid_problems.items():
                    correct_tokens = [a['thought_analysis']['thinking_tokens'] for a in attempts['correct']]
                    incorrect_tokens = [a['thought_analysis']['thinking_tokens'] for a in attempts['incorrect']]
                    
                    correct_transitions = [a['thought_analysis']['thought_transitions'] for a in attempts['correct']]
                    incorrect_transitions = [a['thought_analysis']['thought_transitions'] for a in attempts['incorrect']]
                    
                    avg_correct_tokens = statistics.mean(correct_tokens) if correct_tokens else 0
                    avg_incorrect_tokens = statistics.mean(incorrect_tokens) if incorrect_tokens else 0
                    
                    avg_correct_transitions = statistics.mean(correct_transitions) if correct_transitions else 0
                    avg_incorrect_transitions = statistics.mean(incorrect_transitions) if incorrect_transitions else 0
                    
                    avg_token_diff.append(avg_correct_tokens - avg_incorrect_tokens)
                    avg_transition_diff.append(avg_correct_transitions - avg_incorrect_transitions)
                
                print(f"Average token difference (correct - incorrect): {statistics.mean(avg_token_diff):.2f}")
                print(f"Average transition difference (correct - incorrect): {statistics.mean(avg_transition_diff):.2f}")
    
    if analyze_logits:
        print("\n=== Logit Analysis ===")
        
        # Collect metrics about logit patterns for correct vs incorrect attempts
        correct_attempts = []
        incorrect_attempts = []
        
        for result in results:
            for attempt in result['attempts']:
                if 'logit_analysis' in attempt:
                    if result['is_correct'] and attempt['predicted_answer'] == result['correct_answer']:
                        correct_attempts.append(attempt)
                    else:
                        incorrect_attempts.append(attempt)
        
        # Function to calculate logit statistics for a group of attempts
        def calc_logit_stats(attempts):
            if not attempts:
                return {
                    "count": 0,
                    "entropy": None,
                    "transitions": None
                }
            
            # Collect all entropy stats
            entropy_means = []
            entropy_stds = []
            entropy_quartiles = []
            transition_entropies = defaultdict(lambda: {"before": [], "after": []})
            
            for attempt in attempts:
                if attempt['logit_analysis'].get('entropy_stats') and attempt['logit_analysis']['entropy_stats'].get('mean'):
                    entropy_means.append(attempt['logit_analysis']['entropy_stats']['mean'])
                    entropy_stds.append(attempt['logit_analysis']['entropy_stats']['std'])
                    
                    if attempt['logit_analysis']['entropy_stats'].get('quartiles'):
                        entropy_quartiles.append(attempt['logit_analysis']['entropy_stats']['quartiles'])
                    
                    # Collect transition entropy data
                    if attempt['logit_analysis'].get('transition_entropy'):
                        for phrase, stats in attempt['logit_analysis']['transition_entropy'].items():
                            if stats.get('before_mean') is not None:
                                transition_entropies[phrase]["before"].append(stats['before_mean'])
                            if stats.get('after_mean') is not None:
                                transition_entropies[phrase]["after"].append(stats['after_mean'])
            
            # Calculate average entropy quartiles
            avg_quartiles = []
            if entropy_quartiles:
                # Ensure all quartile lists have the same length
                max_quartiles = max(len(q) for q in entropy_quartiles)
                padded_quartiles = [q + [0] * (max_quartiles - len(q)) for q in entropy_quartiles]
                
                # Calculate average for each quartile position
                for i in range(max_quartiles):
                    quartile_values = [q[i] for q in padded_quartiles if i < len(q)]
                    avg_quartiles.append(statistics.mean(quartile_values) if quartile_values else 0)
            
            # Calculate statistics for transitions
            transition_stats = {}
            for phrase, values in transition_entropies.items():
                if values["before"] and values["after"]:
                    before_mean = statistics.mean(values["before"])
                    after_mean = statistics.mean(values["after"])
                    transition_stats[phrase] = {
                        "before_mean": before_mean,
                        "after_mean": after_mean,
                        "entropy_change": after_mean - before_mean,
                        "count": len(values["before"])
                    }
            
            return {
                "count": len(attempts),
                "entropy": {
                    "mean": statistics.mean(entropy_means) if entropy_means else 0,
                    "std": statistics.mean(entropy_stds) if entropy_stds else 0,
                    "quartiles": avg_quartiles
                },
                "transitions": transition_stats
            }
        
        # Calculate statistics
        correct_stats = calc_logit_stats(correct_attempts)
        incorrect_stats = calc_logit_stats(incorrect_attempts)
        all_stats = calc_logit_stats(correct_attempts + incorrect_attempts)
        
        # Print statistics
        print(f"\nOverall Logit Statistics (All {all_stats['count']} Attempts):")
        if all_stats['entropy'] and all_stats['entropy']['mean']:
            print(f"- Average entropy: {all_stats['entropy']['mean']:.4f}")
            print(f"- Average entropy std: {all_stats['entropy']['std']:.4f}")
            
            if all_stats['entropy']['quartiles']:
                print(f"- Entropy by generation quartile:")
                for i, q in enumerate(all_stats['entropy']['quartiles']):
                    print(f"  - Q{i+1}: {q:.4f}")
            
            if all_stats['transitions']:
                print(f"- Entropy around thought transitions:")
                for phrase, stats in all_stats['transitions'].items():
                    change = stats['entropy_change']
                    change_dir = "increases" if change > 0 else "decreases"
                    print(f"  - {phrase} (n={stats['count']}): Entropy {change_dir} by {abs(change):.4f}")
                    print(f"    - Before: {stats['before_mean']:.4f}, After: {stats['after_mean']:.4f}")
        
        if correct_stats['count'] > 0 and incorrect_stats['count'] > 0:
            print("\nEntropy Comparison (Correct vs Incorrect Attempts):")
            
            if (correct_stats['entropy'] and correct_stats['entropy']['mean'] and 
                incorrect_stats['entropy'] and incorrect_stats['entropy']['mean']):
                
                correct_entropy = correct_stats['entropy']['mean']
                incorrect_entropy = incorrect_stats['entropy']['mean']
                diff = correct_entropy - incorrect_entropy
                
                print(f"- Correct attempts avg entropy: {correct_entropy:.4f}")
                print(f"- Incorrect attempts avg entropy: {incorrect_entropy:.4f}")
                print(f"- Difference (correct - incorrect): {diff:.4f}")
                
                # Compare entropy progression
                if (correct_stats['entropy']['quartiles'] and incorrect_stats['entropy']['quartiles']):
                    print(f"- Entropy progression through generation:")
                    
                    for i in range(min(len(correct_stats['entropy']['quartiles']), 
                                       len(incorrect_stats['entropy']['quartiles']))):
                        c_q = correct_stats['entropy']['quartiles'][i]
                        i_q = incorrect_stats['entropy']['quartiles'][i]
                        q_diff = c_q - i_q
                        
                        print(f"  - Q{i+1}: Correct: {c_q:.4f}, Incorrect: {i_q:.4f}, Diff: {q_diff:.4f}")
                
                # Compare transitions
                common_transitions = set(correct_stats['transitions'].keys()) & set(incorrect_stats['transitions'].keys())
                
                if common_transitions:
                    print(f"- Entropy changes around thought transitions:")
                    
                    for phrase in common_transitions:
                        c_stats = correct_stats['transitions'][phrase]
                        i_stats = incorrect_stats['transitions'][phrase]
                        
                        c_change = c_stats['entropy_change']
                        i_change = i_stats['entropy_change']
                        
                        print(f"  - {phrase}:")
                        print(f"    - Correct: {c_stats['before_mean']:.4f} → {c_stats['after_mean']:.4f} (Δ {c_change:.4f})")
                        print(f"    - Incorrect: {i_stats['before_mean']:.4f} → {i_stats['after_mean']:.4f} (Δ {i_change:.4f})")
                        print(f"    - Difference in entropy change: {c_change - i_change:.4f}")
    
    print("\n=== Incorrect Problems ===")
    for r in results:
        if not r['is_correct']:
            print(f"Problem {r['index']}:")
            print(f"Expected: {r['correct_answer']}")
            print("Predicted answers across attempts:", [
                attempt['predicted_answer'] for attempt in r['attempts']
            ])
            print("---")

def save_raw_response(filename: str, problem_id: int, response_data: Dict):
    """Save raw response data (including logprobs) to a separate file."""
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    
    # Create a timestamped ID for this response
    timestamp = int(time.time())
    response_id = f"{problem_id}_{timestamp}"
    
    # Create or update the raw responses file
    try:
        with open(filename, 'r') as f:
            raw_responses = json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        raw_responses = {}
    
    # Add this response to the collection
    raw_responses[response_id] = response_data
    
    # Save the updated collection
    with open(filename, 'w') as f:
        json.dump(raw_responses, f)
    
    return response_id

def main(model: str, n_attempts: int, year: int = 2024, analyze_thoughts: bool = False, analyze_logits: bool = False, test_time_compute: bool = False, approach_name: str = None, extra_body: dict = None):
    """Main evaluation function that handles gaps in processed indexes."""
    os.makedirs("results", exist_ok=True)

    # Create suffix based on analysis flags
    suffix_parts = []
    if year != 2024:
        suffix_parts.append(f"aime{year}")
    if analyze_thoughts:
        suffix_parts.append("thought_analysis")
    if analyze_logits:
        suffix_parts.append("logit_analysis")
    if approach_name:
        suffix_parts.append(approach_name)

    suffix = "_" + "_".join(suffix_parts) if suffix_parts else ""
    results_file = f"results/evaluation_results_{model.replace('/', '_')}_pass_at_{n_attempts}{suffix}.json"

    dataset = load_dataset_by_year(year)
    existing_results = load_existing_results(results_file)
    
    # Create a set of already processed indexes for efficient lookup
    processed_indexes = {result['index'] for result in existing_results}
    
    for _, item in enumerate(tqdm(dataset, desc="Evaluating problems")):
        id = int(item['id'])
        # Skip if this index has already been processed
        if id in processed_indexes:
            continue
            
        problem_text = item['problem']
        correct_answer = int(item['answer'])

        print(f"\n🔬 Processing Problem {id}: {problem_text[:100]}...")
        print(f"   Expected answer: {correct_answer}")
        if extra_body and 'optillm_approach' in extra_body:
            print(f"   Using approach: {extra_body['optillm_approach']}")

        # Make n attempts for each problem
        attempts = make_n_attempts(problem_text, model, n_attempts, analyze_thoughts, analyze_logits, extra_body)
        is_correct, first_correct = evaluate_pass_at_n(attempts, correct_answer)

        # Report result
        predicted_answers = [attempt.get('predicted_answer') for attempt in attempts]
        print(f"   Predicted: {predicted_answers}")
        if is_correct:
            print(f"   ✅ CORRECT!")
        else:
            print(f"   ❌ Incorrect")

        result = {
            "index": id,
            "problem": problem_text,
            "attempts": attempts,
            "correct_answer": correct_answer,
            "is_correct": is_correct,
            "first_correct_attempt": first_correct
        }
        save_result(results_file, result)
    
    final_results = load_existing_results(results_file)
    analyze_results(final_results, n_attempts, analyze_thoughts, analyze_logits)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate LLM performance on AIME problems")
    parser.add_argument("--model", type=str, required=True, help="OpenAI model to use (e.g., gpt-4, gpt-3.5-turbo)")
    parser.add_argument("--n", type=int, default=1, help="Number of attempts per problem (for pass@n evaluation)")
    parser.add_argument("--year", type=int, default=2024, choices=[2024, 2025], help="AIME year to evaluate (2024 or 2025)")
    parser.add_argument("--approach", type=str, help="OptILLM approach to use (e.g., mars, moa, bon)")
    parser.add_argument("--analyze-thoughts", action="store_true", help="Analyze thinking patterns in responses")
    parser.add_argument("--analyze-logits", action="store_true", help="Analyze token probability distributions")
    parser.add_argument("--test-time-compute", action="store_true", help="Evaluate test-time compute scaling approaches")
    args = parser.parse_args()
    
    if args.test_time_compute:
        # Define test-time compute approaches with same config as eval_optillmbench.py
        TEST_TIME_COMPUTE_APPROACHES = [
            # Baseline
            ("none", "Baseline without any optimization", {}),
            
            # Sequential test-time compute using thinkdeeper with controlled thinking budgets
            ("thinkdeeper_2k", "ThinkDeeper with 2K thinking tokens", {
                "decoding": "thinkdeeper",
                "min_thinking_tokens": 2048,
                "max_thinking_tokens": 2560,  # min + 512 for flexibility
                "max_tokens": 3072  # Total budget: max_thinking_tokens + 512
            }),
            ("thinkdeeper_4k", "ThinkDeeper with 4K thinking tokens", {
                "decoding": "thinkdeeper", 
                "min_thinking_tokens": 4096,
                "max_thinking_tokens": 4608,  # min + 512 for flexibility
                "max_tokens": 5120  # Total budget: max_thinking_tokens + 512
            }),
            ("thinkdeeper_8k", "ThinkDeeper with 8K thinking tokens", {
                "decoding": "thinkdeeper",
                "min_thinking_tokens": 8192,
                "max_thinking_tokens": 8704,  # min + 512 for flexibility
                "max_tokens": 9216  # Total budget: max_thinking_tokens + 512
            }),
            
            # Parallel test-time compute using majority voting with different k values
            ("majority_voting_3", "Majority Voting with k=3", {"k": 3}),
            ("majority_voting_6", "Majority Voting with k=6", {"k": 6}),
            ("majority_voting_9", "Majority Voting with k=9", {"k": 9}),
        ]
        
        # Run evaluation for each approach
        for approach_slug, approach_name, extra_body in TEST_TIME_COMPUTE_APPROACHES:
            print(f"\n{'=' * 80}")
            print(f"Evaluating: {approach_name}")
            print(f"Model: {args.model}")
            print(f"Approach: {approach_slug}")
            print(f"Extra body: {extra_body}")
            print(f"{'=' * 80}\n")
            
            main(args.model, args.n, args.year, args.analyze_thoughts, args.analyze_logits,
                 test_time_compute=True, approach_name=approach_slug, extra_body=extra_body)
    else:
        # Handle approach parameter - only set extra_body if approach is not "none"
        extra_body = {"optillm_approach": args.approach} if args.approach and args.approach != "none" else None
        approach_name = args.approach if args.approach and args.approach != "none" else None

        main(args.model, args.n, args.year, args.analyze_thoughts, args.analyze_logits,
             approach_name=approach_name, extra_body=extra_body)