#!/usr/bin/env python
"""
Aggregate evaluation results from multiple models/checkpoints.
Produces a summary table for easy comparison.
"""

import argparse
import json
import os
from pathlib import Path
from typing import Dict, List, Any
import glob


def load_results(results_file: str) -> Dict[str, Any]:
    """Load results from a JSON file."""
    if os.path.exists(results_file):
        with open(results_file) as f:
            return json.load(f)
    return {}


def get_metric(result: Dict[str, Any], task: str) -> str:
    """Get the main metric for a task."""
    if not result or "error" in result:
        return "N/A"
    
    if task == "pope":
        # POPE uses F1 as main metric
        f1 = result.get("f1", result.get("accuracy", 0))
        return f"{f1:.2f}"
    else:
        # Other tasks use accuracy
        acc = result.get("accuracy", 0)
        return f"{acc:.2f}"


def aggregate_results(
    output_dirs: List[str],
    tasks: List[str],
    model_names: List[str] = None,
) -> Dict[str, Dict[str, str]]:
    """
    Aggregate results from multiple output directories.
    
    Args:
        output_dirs: List of output directories containing results.json
        tasks: List of task names to include
        model_names: Optional custom names for each model
    
    Returns:
        Dictionary mapping model name to task results
    """
    results = {}
    
    for i, output_dir in enumerate(output_dirs):
        # Determine model name
        if model_names and i < len(model_names):
            model_name = model_names[i]
        else:
            # Use directory name as model name
            model_name = Path(output_dir).name
        
        # Load results
        results_file = os.path.join(output_dir, "results.json")
        data = load_results(results_file)
        
        results[model_name] = {}
        for task in tasks:
            task_result = data.get(task, {})
            results[model_name][task] = get_metric(task_result, task)
    
    return results


def print_table(results: Dict[str, Dict[str, str]], tasks: List[str]):
    """Print results as a formatted table."""
    # Header
    header = ["Model"] + tasks
    col_widths = [max(20, max(len(m) for m in results.keys()))] + [12] * len(tasks)
    
    # Print header
    header_line = " | ".join(
        h.center(w) for h, w in zip(header, col_widths)
    )
    print("=" * len(header_line))
    print(header_line)
    print("=" * len(header_line))
    
    # Print rows
    for model_name, task_results in results.items():
        row = [model_name] + [task_results.get(t, "N/A") for t in tasks]
        row_line = " | ".join(
            str(v).center(w) for v, w in zip(row, col_widths)
        )
        print(row_line)
    
    print("=" * len(header_line))


def print_markdown_table(results: Dict[str, Dict[str, str]], tasks: List[str]):
    """Print results as a Markdown table."""
    # Header
    header = ["Model"] + tasks
    print("| " + " | ".join(header) + " |")
    print("| " + " | ".join(["---"] * len(header)) + " |")
    
    # Rows
    for model_name, task_results in results.items():
        row = [model_name] + [task_results.get(t, "N/A") for t in tasks]
        print("| " + " | ".join(row) + " |")


def find_result_dirs(base_dir: str, pattern: str = "*") -> List[str]:
    """Find all result directories matching a pattern."""
    dirs = []
    search_path = os.path.join(base_dir, pattern)
    for path in glob.glob(search_path):
        if os.path.isdir(path):
            results_file = os.path.join(path, "results.json")
            if os.path.exists(results_file):
                dirs.append(path)
    return sorted(dirs)


def main():
    parser = argparse.ArgumentParser(description="Aggregate VLM evaluation results")
    parser.add_argument(
        "--output_dirs", 
        type=str, 
        nargs="+",
        help="Output directories containing results.json"
    )
    parser.add_argument(
        "--base_dir",
        type=str,
        default="eval_outputs",
        help="Base directory to search for results"
    )
    parser.add_argument(
        "--pattern",
        type=str,
        default="*",
        help="Glob pattern to match result directories"
    )
    parser.add_argument(
        "--tasks",
        type=str,
        default="gqa,pope,textvqa,mmmu,mmbench,mathvista",
        help="Comma-separated tasks to include"
    )
    parser.add_argument(
        "--names",
        type=str,
        nargs="*",
        help="Custom names for each model (optional)"
    )
    parser.add_argument(
        "--format",
        type=str,
        choices=["table", "markdown", "json"],
        default="table",
        help="Output format"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Output file (optional, defaults to stdout)"
    )
    
    args = parser.parse_args()
    
    tasks = [t.strip() for t in args.tasks.split(",")]
    
    # Find output directories
    if args.output_dirs:
        output_dirs = args.output_dirs
    else:
        output_dirs = find_result_dirs(args.base_dir, args.pattern)
        if not output_dirs:
            print(f"No result directories found in {args.base_dir}/{args.pattern}")
            return
    
    print(f"Found {len(output_dirs)} result directories")
    
    # Aggregate results
    results = aggregate_results(output_dirs, tasks, args.names)
    
    # Output
    if args.format == "table":
        print_table(results, tasks)
    elif args.format == "markdown":
        print_markdown_table(results, tasks)
    elif args.format == "json":
        output = json.dumps(results, indent=2)
        if args.output:
            with open(args.output, "w") as f:
                f.write(output)
            print(f"Results saved to {args.output}")
        else:
            print(output)
    
    # Save to file if specified
    if args.output and args.format != "json":
        with open(args.output, "w") as f:
            if args.format == "markdown":
                header = ["Model"] + tasks
                f.write("| " + " | ".join(header) + " |\n")
                f.write("| " + " | ".join(["---"] * len(header)) + " |\n")
                for model_name, task_results in results.items():
                    row = [model_name] + [task_results.get(t, "N/A") for t in tasks]
                    f.write("| " + " | ".join(row) + " |\n")
            else:
                f.write(json.dumps(results, indent=2))
        print(f"Results saved to {args.output}")


if __name__ == "__main__":
    main()

