visualize_data.py

import re
import argparse
from pathlib import Path
from collections import defaultdict
from typing import Dict, Any
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.ticker import MaxNLocator


def load_data_from_directory(directory: str) -> Dict[str, Dict[int, Dict[int, Dict[str, Dict[str, Any]]]]]:
    """
    Load and aggregate data from experiment result files in a directory.

    Args:
        directory: Path to directory containing result .txt files

    Returns:
        Nested dictionary with structure:
        data[dataset][N][M][model] = {
            "Accuracy": float,
            "Total Correct": int,
            "Total Problems": int,
            "Average Tokens": float,
            "Average Input Tokens": float,
            "Average Output Tokens": float,
            "Average Latency": float,
            "Average LLM Calls": float,
            "Accuracy by Level": dict (optional, for MATH dataset)
        }
    """
    # Initialize nested defaultdict structure
    data = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))

    # Get all .txt files in directory
    txt_files = Path(directory).glob("*.txt")

    for file_path in txt_files:
        # Parse filename to extract N, M, and dataset
        filename = file_path.stem

        # Try multi_sample_revisions pattern first
        match = re.match(r'multi_sample_revisions_N(\d+)_M(\d+)_(\w+)', filename)

        if match:
            N = int(match.group(1))
            M = int(match.group(2))
            dataset = match.group(3).lower()  # Convert MATH/AIME to lowercase
        else:
            # Try one_shot pattern
            match = re.match(r'one_shot_(\w+)', filename)
            if match:
                N = 1
                M = 0
                dataset = match.group(1).lower()
            else:
                print(f"Skipping file with unrecognized format: {filename}")
                continue

        # Read and parse file content
        with open(file_path, 'r') as f:
            content = f.read()

        # Split into individual run entries
        entries = content.strip().split('\n\n\n')

        for entry in entries:
            if not entry.strip():
                continue

            # Parse the entry
            parsed_data = parse_entry(entry)

            if parsed_data:
                model_name = parsed_data['model']
                # Remove the model name from dict before storing
                metrics = {k: v for k, v in parsed_data.items() if k != 'model'}
                data[dataset][N][M][model_name] = metrics

    # Convert defaultdicts to regular dicts for cleaner output
    return {
        dataset: {
            N: {
                M: dict(models)
                for M, models in Ms.items()
            }
            for N, Ms in Ns.items()
        }
        for dataset, Ns in data.items()
    }


def parse_entry(entry: str) -> Dict[str, Any]:
    """
    Parse a single model run entry from a file.

    Args:
        entry: String containing one model's results

    Returns:
        Dictionary with model name and all metrics
    """
    lines = entry.strip().split('\n')

    if not lines:
        return None

    result = {}

    # Parse run name to extract model
    # Try multi_sample_revisions pattern first
    run_name_match = re.match(r'Run Name: (.+?)_multi_sample_revisions', lines[0])

    if run_name_match:
        result['model'] = run_name_match.group(1)
    else:
        # Try one_shot pattern
        run_name_match = re.match(r'Run Name: (.+?)_one_shot', lines[0])
        if run_name_match:
            result['model'] = run_name_match.group(1)
        else:
            return None

    # Parse metrics from subsequent lines
    for line in lines[1:]:
        line = line.strip()

        if line.startswith('==='):
            continue

        # Match key-value pairs like "Accuracy: 0.76"
        kv_match = re.match(r'(.+?):\s*(.+)', line)
        if kv_match:
            key = kv_match.group(1).strip()
            value_str = kv_match.group(2).strip()

            # Special handling for "Accuracy by Level"
            if key == "Accuracy by Level":
                result[key] = {}
            elif '\t' in line:
                # This is a level entry under "Accuracy by Level"
                level_match = re.match(r'\s*Level (\d+):\s*(.+)', line)
                if level_match and "Accuracy by Level" in result:
                    level = int(level_match.group(1))
                    accuracy = float(level_match.group(2))
                    result["Accuracy by Level"][level] = accuracy
            else:
                # Try to convert to appropriate type
                try:
                    # Check if it's an integer
                    if '.' not in value_str:
                        result[key] = int(value_str)
                    else:
                        result[key] = float(value_str)
                except ValueError:
                    # Keep as string if conversion fails
                    result[key] = value_str

    return result


def plot_accuracy(data: Dict, fixed_param: str, fixed_value: int, model: str = None, dataset: str = None):
    """
    Plot accuracy vs varying parameter (N or M).

    Args:
        data: The loaded data structure from load_data_from_directory()
        fixed_param: Either 'M' or 'N' (the parameter to hold fixed)
        fixed_value: The value of the fixed parameter
        model: The model name to plot (mutually exclusive with dataset)
        dataset: The dataset to plot all models for (mutually exclusive with model)
    """
    varying_param = 'N' if fixed_param == 'M' else 'M'

    if model is not None:
        # Model mode: plot both datasets for one model
        plot_data = {'math': {}, 'aime': {}}

        for ds in data.keys():
            for N in data[ds].keys():
                for M in data[ds][N].keys():
                    # Check if this matches our fixed parameter
                    if fixed_param == 'M' and M == fixed_value:
                        varying_value = N
                    elif fixed_param == 'N' and N == fixed_value:
                        varying_value = M
                    else:
                        continue

                    # Check if the model exists in this configuration
                    if model in data[ds][N][M]:
                        accuracy = data[ds][N][M][model].get('Accuracy')
                        if accuracy is not None:
                            plot_data[ds][varying_value] = accuracy

        # Check if we have any data to plot
        if not any(plot_data.values()):
            print(f"Error: No data found for model '{model}' with {fixed_param}={fixed_value}")
            print(f"\nAvailable models:")
            models_found = set()
            for ds in data.keys():
                for N in data[ds].keys():
                    for M in data[ds][N].keys():
                        if fixed_param == 'M' and M == fixed_value:
                            models_found.update(data[ds][N][M].keys())
                        elif fixed_param == 'N' and N == fixed_value:
                            models_found.update(data[ds][N][M].keys())
            for m in sorted(models_found):
                print(f"  - {m}")
            return

        # Create the plot
        plt.figure(figsize=(10, 6))

        # Plot each dataset as a separate line (only if 2+ points)
        for ds, points in plot_data.items():
            if points and len(points) >= 2:
                x_values = sorted(points.keys())
                y_values = [points[x] for x in x_values]
                plt.plot(x_values, y_values, marker='o', label=ds.upper(), linewidth=2, markersize=8)

        # Configure the plot
        plt.xlabel(varying_param, fontsize=12)
        plt.ylabel('Accuracy', fontsize=12)
        plt.title(f'Accuracy vs {varying_param} for {model}\n({fixed_param}={fixed_value})', fontsize=14)
        plt.grid(True, alpha=0.3)
        plt.legend(fontsize=10)
        plt.ylim(0, 1.0)

    else:
        # Dataset mode: plot all models for one dataset
        plot_data = defaultdict(dict)  # model -> {varying_value: accuracy}

        if dataset not in data:
            print(f"Error: Dataset '{dataset}' not found in data")
            print(f"Available datasets: {', '.join(sorted(data.keys()))}")
            return

        for N in data[dataset].keys():
            for M in data[dataset][N].keys():
                # Check if this matches our fixed parameter
                if fixed_param == 'M' and M == fixed_value:
                    varying_value = N
                elif fixed_param == 'N' and N == fixed_value:
                    varying_value = M
                else:
                    continue

                # Collect all models at this configuration
                for model_name, metrics in data[dataset][N][M].items():
                    accuracy = metrics.get('Accuracy')
                    if accuracy is not None:
                        plot_data[model_name][varying_value] = accuracy

        # Check if we have any data to plot
        if not plot_data:
            print(f"Error: No data found for dataset '{dataset}' with {fixed_param}={fixed_value}")
            return

        # Create the plot
        plt.figure(figsize=(12, 6))

        # Plot each model as a separate line (only if 2+ points)
        for model_name, points in sorted(plot_data.items()):
            if points and len(points) >= 2:
                x_values = sorted(points.keys())
                y_values = [points[x] for x in x_values]
                # Use shorter model names for legend
                short_name = model_name.split('/')[-1] if '/' in model_name else model_name
                plt.plot(x_values, y_values, marker='o', label=short_name, linewidth=2, markersize=8)

        # Configure the plot
        plt.xlabel(varying_param, fontsize=12)
        plt.ylabel('Accuracy', fontsize=12)
        plt.title(f'Accuracy vs {varying_param} for {dataset.upper()} dataset\n({fixed_param}={fixed_value})', fontsize=14)
        plt.grid(True, alpha=0.3)
        plt.legend(fontsize=9, loc='best')
        plt.ylim(0, 1.0)

    # Format x-axis to show integer values
    plt.gca().xaxis.set_major_locator(plt.MaxNLocator(integer=True))

    plt.tight_layout()
    plt.show()


def plot_heatmap_grid(data: Dict, model: str, dataset: str, metric: str = 'accuracy'):
    """
    Create a 2D heatmap grid showing metrics across N and M dimensions.

    Args:
        data: The loaded data structure from load_data_from_directory()
        model: The model name to plot
        dataset: The dataset to visualize ('math' or 'aime')
        metric: Metric to visualize ('accuracy', 'tokens', 'latency', 'llm_calls')
    """
    # Check if dataset exists
    if dataset not in data:
        print(f"Error: Dataset '{dataset}' not found in data")
        print(f"Available datasets: {', '.join(sorted(data.keys()))}")
        return

    # Collect data points and find dimensions
    data_dict = {}
    all_N = set()
    all_M = set()

    for N in data[dataset].keys():
        for M in data[dataset][N].keys():
            if model in data[dataset][N][M]:
                metrics = data[dataset][N][M][model]
                accuracy = metrics.get('Accuracy')
                tokens = metrics.get('Average Tokens')
                latency = metrics.get('Average Latency')
                llm_calls = metrics.get('Average LLM Calls')

                if accuracy is not None:
                    data_dict[(N, M)] = (accuracy, tokens, latency, llm_calls)
                    all_N.add(N)
                    all_M.add(M)

    # Check if we have data to plot
    if not data_dict:
        print(f"Error: No data found for model '{model}' on dataset '{dataset}'")
        print(f"\nAvailable models for {dataset} dataset:")
        models_found = set()
        for N in data[dataset].keys():
            for M in data[dataset][N].keys():
                models_found.update(data[dataset][N][M].keys())
        for m in sorted(models_found):
            print(f"  - {m}")
        return

    # Create sorted lists of unique N and M values that have data
    sorted_N = sorted(all_N)
    sorted_M = sorted(all_M)

    # Create mappings from actual values to grid indices
    N_to_index = {n: i for i, n in enumerate(sorted_N)}
    M_to_index = {m: i for i, m in enumerate(sorted_M)}

    # Grid dimensions based on actual data points
    grid_height = len(sorted_N)
    grid_width = len(sorted_M)

    # Create grid initialized with NaN for missing data
    grid = np.full((grid_height, grid_width), np.nan)

    # Store both the grid values and the raw metrics for labeling
    raw_metrics_grid = {}  # Store raw metrics for text labels

    # Populate grid with calculated metric values
    for (n, m), metrics_tuple in data_dict.items():
        accuracy, tokens, latency, llm_calls = metrics_tuple

        # Store raw metrics
        raw_metrics_grid[(N_to_index[n], M_to_index[m])] = (accuracy, tokens, latency, llm_calls)

        # Calculate the appropriate metric value
        if metric == 'accuracy':
            value = accuracy
        elif metric == 'tokens':
            # Tokens per 1% accuracy (inverted and scaled by 100)
            value = (tokens / accuracy) / 100 if accuracy and accuracy > 0 else np.nan
        elif metric == 'latency':
            value = accuracy / latency if latency and latency > 0 else np.nan
        elif metric == 'llm_calls':
            value = accuracy / llm_calls if llm_calls and llm_calls > 0 else np.nan
        else:
            value = np.nan

        grid[N_to_index[n], M_to_index[m]] = value

    # Determine colormap scaling based on metric
    if metric == 'accuracy':
        vmin, vmax = 0.0, 1.0
    else:
        # Use data-driven min/max for efficiency metrics
        vmin = np.nanmin(grid)
        vmax = np.nanmax(grid)

    # Determine top cells to mark with ranks (works for all metrics)
    cell_ranks = {}  # Maps (i, j) -> rank (1, 2, or 3)

    # Get all cells with their raw accuracy scores
    valid_accuracies = []
    for (i, j), metrics_tuple in raw_metrics_grid.items():
        accuracy = metrics_tuple[0]  # First element is accuracy
        if accuracy is not None:
            valid_accuracies.append((accuracy, i, j))

    if valid_accuracies:
        # Sort by accuracy descending
        valid_accuracies.sort(key=lambda x: x[0], reverse=True)

        # Get unique accuracy scores in descending order
        unique_accuracies = sorted(set(acc for acc, _, _ in valid_accuracies), reverse=True)

        # Progressively add tiers until we have at least 3 cells
        total_marked = 0
        for tier_idx in range(min(3, len(unique_accuracies))):
            tier_accuracy = unique_accuracies[tier_idx]
            rank = tier_idx + 1  # 1, 2, or 3

            # Add all cells with this accuracy
            tier_cells = [(i, j) for acc, i, j in valid_accuracies if acc == tier_accuracy]
            for i, j in tier_cells:
                cell_ranks[(i, j)] = rank

            total_marked += len(tier_cells)

            # Stop if we have 3 or more cells
            if total_marked >= 3:
                break

    # Create figure and axes
    fig, ax = plt.subplots(figsize=(max(10, grid_width * 1.5), max(8, grid_height * 1.2)))

    # Choose colormap (reverse for tokens metric where lower is better)
    colormap = 'viridis_r' if metric == 'tokens' else 'viridis'

    # Create heatmap using imshow
    im = ax.imshow(grid, cmap=colormap, aspect='auto', origin='lower',
                   vmin=vmin, vmax=vmax, interpolation='nearest')

    # Add cell boundaries
    for i in range(grid_height + 1):
        ax.axhline(i - 0.5, color='black', linewidth=1.5)
    for j in range(grid_width + 1):
        ax.axvline(j - 0.5, color='black', linewidth=1.5)

    # Fill missing data cells with green and add text annotations
    for i in range(grid_height):
        for j in range(grid_width):
            if np.isnan(grid[i, j]):
                # Add green rectangle for missing data
                rect = mpatches.Rectangle((j-0.5, i-0.5), 1, 1,
                                         linewidth=0, edgecolor='none',
                                         facecolor='#90EE90', alpha=0.6)
                ax.add_patch(rect)
                # Add "N/A" text
                ax.text(j, i, 'N/A', ha="center", va="center",
                       color="darkgray", fontsize=10, fontweight='bold')
            else:
                # Add metric value text with appropriate formatting
                value = grid[i, j]

                if metric == 'accuracy':
                    text_value = f'{value:.3f}'
                    fontsize = 10
                else:
                    # For efficiency metrics, include accuracy and base metric value
                    if (i, j) in raw_metrics_grid:
                        accuracy, tokens, latency, llm_calls = raw_metrics_grid[(i, j)]

                        if metric == 'tokens':
                            base_value = f'{tokens:.0f}' if tokens else 'N/A'
                            metric_name = 'Tokens'
                            # Format tokens per 1% accuracy (larger values)
                            efficiency = f'{value:.0f}' if value and value >= 1 else f'{value:.1f}'
                            eff_label = 'Tok/1%'
                        elif metric == 'latency':
                            base_value = f'{latency:.1f}' if latency else 'N/A'
                            metric_name = 'Latency'
                            # Format as before for latency
                            if value < 0.001:
                                efficiency = f'{value:.2e}'
                            else:
                                efficiency = f'{value:.4f}'
                            eff_label = 'Eff'
                        elif metric == 'llm_calls':
                            base_value = f'{llm_calls:.1f}' if llm_calls else 'N/A'
                            metric_name = 'Calls'
                            # Format as before for llm_calls
                            if value < 0.001:
                                efficiency = f'{value:.2e}'
                            else:
                                efficiency = f'{value:.4f}'
                            eff_label = 'Eff'

                        # Multi-line text: Accuracy / Base Metric / Efficiency
                        text_value = f'Acc: {accuracy:.3f}\n{metric_name}: {base_value}\n{eff_label}: {efficiency}'
                        fontsize = 8
                    else:
                        # Fallback if raw metrics not found
                        if value < 0.001:
                            text_value = f'{value:.2e}'
                        else:
                            text_value = f'{value:.4f}'
                        fontsize = 10

                ax.text(j, i, text_value, ha="center", va="center",
                       color="white", fontsize=fontsize, fontweight='bold')

                # Add rank number for top scores in efficiency modes
                if (i, j) in cell_ranks:
                    rank = cell_ranks[(i, j)]
                    # Add rank number in top-right corner
                    ax.text(j + 0.4, i + 0.4, str(rank), ha="center", va="center",
                           color="gold", fontsize=14, fontweight='bold',
                           bbox=dict(boxstyle='circle,pad=0.1', facecolor='black',
                                   edgecolor='gold', linewidth=2))

    # Configure axes labels
    ax.set_xticks(np.arange(grid_width))
    ax.set_yticks(np.arange(grid_height))
    ax.set_xticklabels([str(m) for m in sorted_M])
    ax.set_yticklabels([str(n) for n in sorted_N])

    # Metric display labels
    metric_labels = {
        'accuracy': 'Accuracy',
        'tokens': 'Tokens per 1% Accuracy',
        'latency': 'Time Efficiency (Accuracy/Latency)',
        'llm_calls': 'Call Efficiency (Accuracy/LLM Calls)'
    }

    ax.set_xlabel('M', fontsize=14, fontweight='bold')
    ax.set_ylabel('N', fontsize=14, fontweight='bold')
    ax.set_title(f'{metric_labels[metric]} Heatmap for {model}\non {dataset.upper()} dataset',
                fontsize=16, fontweight='bold', pad=20)

    # Add colorbar
    cbar = plt.colorbar(im, ax=ax, pad=0.02, shrink=0.8)
    cbar.set_label(metric_labels[metric], fontsize=12, fontweight='bold')

    plt.tight_layout()
    plt.show()


# Keep old name as alias for backward compatibility
plot_3d_accuracy = plot_heatmap_grid


def plot_bar_graph(data: Dict, dataset: str, N: int, M: int, metric: str = 'accuracy', compare: bool = False):
    """
    Create a bar graph comparing all models at a specific N,M configuration.

    Args:
        data: The loaded data structure from load_data_from_directory()
        dataset: Dataset to visualize ('math' or 'aime')
        N: N value for the configuration
        M: M value for the configuration
        metric: Metric to display ('accuracy', 'tokens', 'latency', 'llm_calls')
        compare: If True, show max accuracy across all N,M configs as comparison bars
    """
    # Check if dataset exists
    if dataset not in data:
        print(f"Error: Dataset '{dataset}' not found in data")
        print(f"Available datasets: {', '.join(sorted(data.keys()))}")
        return

    # Check if N,M configuration exists
    if N not in data[dataset]:
        print(f"Error: N={N} not found in {dataset} dataset")
        print(f"Available N values: {', '.join(str(n) for n in sorted(data[dataset].keys()))}")
        return

    if M not in data[dataset][N]:
        print(f"Error: M={M} not found for N={N} in {dataset} dataset")
        print(f"Available M values for N={N}: {', '.join(str(m) for m in sorted(data[dataset][N].keys()))}")
        return

    # Extract all models and their metric values at this N,M configuration
    models_data = data[dataset][N][M]

    if not models_data:
        print(f"Error: No models found at N={N}, M={M} for {dataset} dataset")
        return

    # Collect model names and values
    model_names = []
    values = []

    for model_name, metrics in models_data.items():
        accuracy = metrics.get('Accuracy')
        if accuracy is None:
            continue

        # Calculate the appropriate metric value
        if metric == 'accuracy':
            value = accuracy
        elif metric == 'tokens':
            tokens = metrics.get('Average Tokens')
            value = (tokens / accuracy) / 100 if accuracy and tokens and accuracy > 0 else None
        elif metric == 'latency':
            latency = metrics.get('Average Latency')
            value = accuracy / latency if latency and latency > 0 else None
        elif metric == 'llm_calls':
            llm_calls = metrics.get('Average LLM Calls')
            value = accuracy / llm_calls if llm_calls and llm_calls > 0 else None
        else:
            value = None

        if value is not None:
            model_names.append(model_name)
            values.append(value)

    if not model_names:
        print(f"Error: No valid data found for metric '{metric}' at N={N}, M={M}")
        return

    # Sort by value descending
    sorted_indices = np.argsort(values)[::-1]
    model_names = [model_names[i] for i in sorted_indices]
    values = [values[i] for i in sorted_indices]

    # If compare mode, find max accuracy for each model across all N,M configs
    # Also get current accuracies to check if they differ
    max_accuracies = {}
    current_accuracies = {}
    if compare:
        for model_name in model_names:
            # Get current accuracy for this model at current N,M
            current_acc = models_data[model_name].get('Accuracy', 0.0)
            current_accuracies[model_name] = current_acc

            max_acc = 0.0
            # Search all N,M configurations for this model in this dataset
            for N_iter in data[dataset].keys():
                for M_iter in data[dataset][N_iter].keys():
                    if model_name in data[dataset][N_iter][M_iter]:
                        acc = data[dataset][N_iter][M_iter][model_name].get('Accuracy')
                        if acc and acc > max_acc:
                            max_acc = acc

            # Only store if we found a max accuracy > 0 and it differs from current
            if max_acc > 0 and abs(max_acc - current_acc) > 0.0001:
                max_accuracies[model_name] = max_acc

    # Shorten model names for display (use last part after /)
    display_names = [name.split('/')[-1] if '/' in name else name for name in model_names]

    # Create figure
    fig, ax = plt.subplots(figsize=(max(10, len(model_names) * 0.8), 8))

    # Set up positions and widths
    x_positions = np.arange(len(model_names))
    bar_width = 0.6

    # If compare mode, plot comparison bars first (behind main bars)
    # Only plot bars for models where max differs from current
    if compare and max_accuracies:
        # Create comparison values array with 0 for models without differences
        comparison_values = []
        for name in model_names:
            if name in max_accuracies:
                comparison_values.append(max_accuracies[name])
            else:
                comparison_values.append(0)

        # Only plot if we have at least one comparison value
        if any(comparison_values):
            ax.bar(x_positions, comparison_values, width=bar_width * 1.2,
                   color='lightgray', alpha=0.5, edgecolor='black', linewidth=1,
                   label='Max Accuracy (all configs)', zorder=1)

    # Create main bar chart with color gradient
    colors = plt.cm.viridis(np.linspace(0.3, 0.9, len(values)))
    bars = ax.bar(x_positions, values, width=bar_width, color=colors,
                  edgecolor='black', linewidth=1.5, zorder=2)

    # Add value labels on top of main bars
    for i, (bar, value) in enumerate(zip(bars, values)):
        height = bar.get_height()
        if metric == 'accuracy':
            label = f'{value:.3f}'
        elif metric == 'tokens':
            label = f'{value:.0f}'
        else:
            label = f'{value:.4f}' if value >= 0.001 else f'{value:.2e}'

        ax.text(bar.get_x() + bar.get_width()/2., height,
               label, ha='center', va='bottom', fontsize=10, fontweight='bold')

    # Add value labels on top of comparison bars if compare mode is active
    # Only show labels for models where max differs from current
    if compare and max_accuracies:
        for i, model_name in enumerate(model_names):
            if model_name in max_accuracies:
                max_acc = max_accuracies[model_name]
                # Format as accuracy (always 3 decimal places since it's accuracy)
                label = f'{max_acc:.3f}'
                # Position label on top of comparison bar
                ax.text(x_positions[i], max_acc,
                       label, ha='center', va='bottom',
                       fontsize=9, fontweight='bold', color='dimgray',
                       bbox=dict(boxstyle='round,pad=0.2', facecolor='white',
                                edgecolor='gray', alpha=0.7))

    # Configure axes
    ax.set_xticks(range(len(model_names)))
    ax.set_xticklabels(display_names, rotation=45, ha='right', fontsize=10)

    # Metric labels
    metric_labels = {
        'accuracy': 'Accuracy',
        'tokens': 'Tokens per 1% Accuracy',
        'latency': 'Time Efficiency (Accuracy/Latency)',
        'llm_calls': 'Call Efficiency (Accuracy/LLM Calls)'
    }

    ax.set_ylabel(metric_labels[metric], fontsize=12, fontweight='bold')
    ax.set_xlabel('Model', fontsize=12, fontweight='bold')
    ax.set_title(f'Model Comparison at N={N}, M={M} on {dataset.upper()} dataset\n({metric_labels[metric]})',
                fontsize=14, fontweight='bold', pad=20)

    # Add grid for readability
    ax.grid(axis='y', alpha=0.3, linestyle='--')
    ax.set_axisbelow(True)

    # Add legend if compare mode is active and there are actual comparison bars
    if compare and max_accuracies and any(max_accuracies.values()):
        ax.legend(loc='upper right', fontsize=10, framealpha=0.9)

    # Set y-axis limits with some padding
    if metric == 'accuracy':
        ax.set_ylim(0, 1.0)
    else:
        ax.set_ylim(0, max(values) * 1.15)

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='Visualize model accuracy across different N and M values'
    )

    # M and N arguments (can be used independently or together)
    parser.add_argument('-M', type=int, help='M value: alone = vary N, with -N = bar graph mode')
    parser.add_argument('-N', type=int, help='N value: alone = vary M, with -M = bar graph mode')

    # Model and dataset arguments (constraints depend on mode)
    parser.add_argument('--model', help='Model name to plot')
    parser.add_argument('--dataset', choices=['math', 'aime'], help='Dataset to plot')

    parser.add_argument(
        '--data-dir',
        default='data/baseline_centralized_judge',
        help='Directory containing result files (default: baseline_centralized_judge)'
    )

    parser.add_argument(
        '--metric',
        choices=['accuracy', 'tokens', 'latency', 'llm_calls'],
        default='accuracy',
        help='Metric to visualize: accuracy (default), tokens (accuracy/tokens), latency (accuracy/latency), llm_calls (accuracy/llm_calls)'
    )

    parser.add_argument(
        '--compare',
        action='store_true',
        help='In bar graph mode, show max accuracy across all N,M configs as comparison bars'
    )

    args = parser.parse_args()

    # Load data
    print(f"Loading data from {args.data_dir}...")
    data = load_data_from_directory(args.data_dir)

    if not data:
        print(f"Error: No data found in {args.data_dir}")
        exit(1)

    print(f"Data loaded successfully from {args.data_dir}")

    # Determine visualization mode based on which arguments are provided
    if args.M is not None and args.N is not None:
        # Bar graph mode: both M and N specified, requires dataset
        if not args.dataset:
            print("Error: Bar graph mode requires --dataset")
            print("Usage: python visualize_data.py --dataset <math|aime> -M <value> -N <value>")
            exit(1)
        if args.model:
            print("Error: Bar graph mode uses --dataset, not --model")
            exit(1)

        print(f"Creating bar graph for all models at N={args.N}, M={args.M} on {args.dataset} dataset with metric: {args.metric}...")
        plot_bar_graph(data, args.dataset, args.N, args.M, metric=args.metric, compare=args.compare)

    elif args.M is None and args.N is None:
        # Heatmap mode: neither M nor N specified, requires both model and dataset
        if not args.model or not args.dataset:
            print("Error: Heatmap mode requires both --model and --dataset arguments")
            print("Usage: python visualize_data.py --model <model_name> --dataset <math|aime>")
            exit(1)

        print(f"Creating heatmap for {args.model} on {args.dataset} dataset with metric: {args.metric}...")
        plot_3d_accuracy(data, args.model, args.dataset, metric=args.metric)

    else:
        # 2D line plot mode: either M or N specified (but not both)
        if args.model and args.dataset:
            print("Error: In 2D line plot mode, specify either --model or --dataset, not both")
            exit(1)

        if not args.model and not args.dataset:
            print("Error: In 2D line plot mode, you must specify either --model or --dataset")
            exit(1)

        # Determine which parameter is fixed
        if args.M is not None:
            fixed_param = 'M'
            fixed_value = args.M
        else:
            fixed_param = 'N'
            fixed_value = args.N

        # Create the 2D line plot (metric flag not applicable in this mode)
        plot_accuracy(data, fixed_param, fixed_value, model=args.model, dataset=args.dataset)