#!/usr/bin/env python3
"""Scan EvalScope output directories and aggregate per-model benchmark scores."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# Display order for comparison charts
CATEGORY_ORDER = [
    ('math', '数学', [
        'aime24', 'aime25', 'aime26', 'hmmt26',
        'imo_answerbench', 'competition_math', 'gsm8k',
    ]),
    ('code', '代码', ['humaneval', 'live_code_bench', 'bigcodebench']),
    ('science', '科学', ['gpqa_diamond', 'super_gpqa', 'hle']),
    ('knowledge', '知识', [
        'mmlu', 'mmlu_pro', 'cmmlu', 'bbh', 'arc',
        'drop', 'hellaswag', 'winogrande', 'simple_qa', 'trivia_qa',
    ]),
    ('long_context', '长文本', ['longbench_v2', 'openai_mrcr']),
    ('tool_agent', '智能体', ['bfcl_v3', 'general_fc', 'tau2_bench']),
]

BENCHMARK_ALIAS = {
    'hle_low': 'hle',
}

SKIP_DIR_NAMES = {
    'active_time', 'perf_stats_backup', 'predictions_archive',
    'logs', 'configs', 'reports', 'predictions', 'reviews',
}


def _is_seed_dir(path: Path) -> bool:
    return path.is_dir() and path.name.startswith('seed_')


def _is_benchmark_dir(path: Path) -> bool:
    if not path.is_dir() or path.name in SKIP_DIR_NAMES:
        return False
    try:
        return any(_is_seed_dir(child) for child in path.iterdir())
    except OSError:
        return False


def _is_model_folder(path: Path) -> bool:
    if not path.is_dir() or path.name.startswith('.'):
        return False
    try:
        return any(_is_benchmark_dir(child) for child in path.iterdir())
    except OSError:
        return False


def extract_score(report_data: dict) -> Optional[float]:
    score = report_data.get('score')
    if score is not None:
        try:
            return float(score)
        except (TypeError, ValueError):
            pass
    for metric in report_data.get('metrics') or []:
        if metric.get('name') in ('mean_acc', 'acc', 'accuracy', 'pass@1', 'Score'):
            for key in ('score', 'macro_score'):
                if metric.get(key) is not None:
                    try:
                        return float(metric[key])
                    except (TypeError, ValueError):
                        continue
    # fallback: first metric with a score
    for metric in report_data.get('metrics') or []:
        for key in ('score', 'macro_score'):
            if metric.get(key) is not None:
                try:
                    return float(metric[key])
                except (TypeError, ValueError):
                    continue
    return None


def _find_report_files(seed_dir: Path, benchmark: str) -> List[Path]:
    reports_dir = seed_dir / 'reports'
    if not reports_dir.is_dir():
        return []
    preferred = [
        reports_dir / f'{benchmark}.json',
        reports_dir / f'{BENCHMARK_ALIAS.get(benchmark, benchmark)}.json',
    ]
    found = [p for p in preferred if p.is_file()]
    if found:
        return found
    # nested: reports/<model>/<benchmark>.json (legacy)
    nested = sorted(reports_dir.glob(f'*/*.json')) + sorted(reports_dir.glob('*.json'))
    return [p for p in nested if p.is_file()]


def _read_report(path: Path) -> Optional[dict]:
    try:
        return json.loads(path.read_text(encoding='utf-8'))
    except Exception:
        return None


def collect_benchmark_scores(model_dir: Path, benchmark: str) -> Optional[dict]:
    bench_dir = model_dir / benchmark
    if not bench_dir.is_dir():
        # alias: hle stored as hle_low
        for alias_src, alias_dst in BENCHMARK_ALIAS.items():
            if benchmark == alias_dst:
                alt = model_dir / alias_src
                if alt.is_dir():
                    bench_dir = alt
                    break
        else:
            return None

    scores: List[float] = []
    runs: List[dict] = []
    model_name = None
    num_samples = None

    for seed_dir in sorted(p for p in bench_dir.iterdir() if _is_seed_dir(p)):
        for report_path in _find_report_files(seed_dir, benchmark):
            data = _read_report(report_path)
            if not data:
                continue
            score = extract_score(data)
            if score is None:
                continue
            scores.append(score)
            if model_name is None:
                model_name = data.get('model_name')
            if num_samples is None:
                num_samples = data.get('num')
                if num_samples is None:
                    metrics = data.get('metrics') or []
                    if metrics:
                        num_samples = metrics[0].get('num')
            try:
                rel = str(report_path.relative_to(model_dir))
            except ValueError:
                rel = str(report_path)
            runs.append({
                'seed_dir': seed_dir.name,
                'score': score,
                'report': rel,
            })
            break  # one report per seed dir

    if not scores:
        return None

    avg = sum(scores) / len(scores)
    return {
        'benchmark': BENCHMARK_ALIAS.get(benchmark, benchmark),
        'score': round(avg, 6),
        'scores': scores,
        'n_runs': len(scores),
        'num_samples': num_samples,
        'model_name': model_name,
        'runs': runs,
    }


def list_benchmarks_in_model(model_dir: Path) -> List[str]:
    names = []
    for child in sorted(model_dir.iterdir()):
        if not _is_benchmark_dir(child):
            continue
        names.append(BENCHMARK_ALIAS.get(child.name, child.name))
    # unique preserve order
    seen = set()
    out = []
    for n in names:
        if n not in seen:
            seen.add(n)
            out.append(n)
    return out


def scan_output_dir(output_dir: Path) -> dict:
    output_dir = Path(output_dir)
    if not output_dir.is_dir():
        return {
            'output_dir': str(output_dir),
            'models': [],
            'benchmarks': [],
            'categories': [],
            'matrix': {},
        }

    models = []
    all_benchmarks = set()
    matrix: Dict[str, Dict[str, dict]] = {}

    for child in sorted(output_dir.iterdir()):
        if not _is_model_folder(child):
            continue
        folder = child.name
        bench_names = list_benchmarks_in_model(child)
        model_scores: Dict[str, dict] = {}
        served_model = None
        for bench in bench_names:
            info = collect_benchmark_scores(child, bench)
            if not info:
                continue
            model_scores[bench] = info
            all_benchmarks.add(bench)
            if served_model is None and info.get('model_name'):
                served_model = info['model_name']

        if not model_scores:
            continue

        models.append({
            'folder': folder,
            'model_name': served_model or folder,
            'benchmarks': sorted(model_scores.keys()),
            'n_benchmarks': len(model_scores),
            'path': str(child),
        })
        matrix[folder] = model_scores

    # Ordered benchmark list by category, then leftovers
    ordered = []
    seen = set()
    categories = []
    for cat_id, cat_name, items in CATEGORY_ORDER:
        present = [b for b in items if b in all_benchmarks]
        if present:
            categories.append({'id': cat_id, 'name': cat_name, 'items': present})
            for b in present:
                if b not in seen:
                    seen.add(b)
                    ordered.append(b)
    others = sorted(all_benchmarks - seen)
    if others:
        categories.append({'id': 'other', 'name': '其他', 'items': others})
        ordered.extend(others)

    # Compact matrix for API: folder -> benchmark -> score summary
    compact = {}
    for folder, benches in matrix.items():
        compact[folder] = {
            b: {
                'score': info['score'],
                'n_runs': info['n_runs'],
                'num_samples': info.get('num_samples'),
                'model_name': info.get('model_name'),
            }
            for b, info in benches.items()
        }

    return {
        'output_dir': str(output_dir),
        'models': models,
        'benchmarks': ordered,
        'categories': categories,
        'matrix': compact,
    }


def compare_models(
    output_dir: Path,
    folders: Optional[List[str]] = None,
    benchmarks: Optional[List[str]] = None,
) -> dict:
    overview = scan_output_dir(output_dir)
    available = {m['folder'] for m in overview['models']}
    if folders:
        selected = [f for f in folders if f in available]
    else:
        selected = [m['folder'] for m in overview['models']]

    if benchmarks:
        bench_list = [b for b in benchmarks if b in overview['benchmarks']]
    else:
        bench_list = list(overview['benchmarks'])

    series = []
    for folder in selected:
        scores = []
        for b in bench_list:
            cell = overview['matrix'].get(folder, {}).get(b)
            scores.append(cell['score'] if cell else None)
        model_meta = next((m for m in overview['models'] if m['folder'] == folder), None)
        series.append({
            'folder': folder,
            'label': model_meta['model_name'] if model_meta else folder,
            'display': folder,
            'scores': scores,
        })

    # Per-benchmark ranking
    ranking = []
    for i, b in enumerate(bench_list):
        rows = []
        for s in series:
            if s['scores'][i] is not None:
                rows.append({'folder': s['folder'], 'label': s['label'], 'score': s['scores'][i]})
        rows.sort(key=lambda x: x['score'], reverse=True)
        ranking.append({'benchmark': b, 'rows': rows})

    # Capability-domain aggregation: mean score over selected benches in each category
    bench_set = set(bench_list)
    active_categories = []
    for cat in overview['categories']:
        items = [b for b in cat['items'] if b in bench_set]
        if items:
            active_categories.append({
                'id': cat['id'],
                'name': cat['name'],
                'items': items,
            })

    category_labels = [c['name'] for c in active_categories]
    category_series = []
    for folder in selected:
        scores = []
        details = []
        for cat in active_categories:
            vals = []
            for b in cat['items']:
                cell = overview['matrix'].get(folder, {}).get(b)
                if cell and cell.get('score') is not None:
                    vals.append(float(cell['score']))
            if vals:
                avg = sum(vals) / len(vals)
                scores.append(round(avg, 6))
            else:
                avg = None
                scores.append(None)
            details.append({
                'category': cat['name'],
                'score': avg,
                'n_benchmarks': len(vals),
                'benchmarks': cat['items'],
            })
        model_meta = next((m for m in overview['models'] if m['folder'] == folder), None)
        category_series.append({
            'folder': folder,
            'label': model_meta['model_name'] if model_meta else folder,
            'display': folder,
            'scores': scores,
            'details': details,
        })

    category_ranking = []
    for i, cat in enumerate(active_categories):
        rows = []
        for s in category_series:
            if s['scores'][i] is not None:
                rows.append({
                    'folder': s['folder'],
                    'label': s['label'],
                    'score': s['scores'][i],
                    'n_benchmarks': s['details'][i]['n_benchmarks'],
                })
        rows.sort(key=lambda x: x['score'], reverse=True)
        category_ranking.append({
            'category': cat['name'],
            'id': cat['id'],
            'items': cat['items'],
            'rows': rows,
        })

    # Overall aggregates
    overall = []
    for s in series:
        vals = [v for v in s['scores'] if v is not None]
        overall.append({
            'folder': s['folder'],
            'label': s['label'],
            'display': s['display'],
            'mean': round(sum(vals) / len(vals), 6) if vals else None,
            'n': len(vals),
        })
    overall.sort(key=lambda x: (x['mean'] if x['mean'] is not None else -1), reverse=True)

    # Win count: how many benchmarks each model ranks #1
    win_counts = {s['folder']: 0 for s in series}
    for i, b in enumerate(bench_list):
        best_score = -1
        best_folders = []
        for s in series:
            v = s['scores'][i]
            if v is None:
                continue
            if v > best_score:
                best_score = v
                best_folders = [s['folder']]
            elif v == best_score:
                best_folders.append(s['folder'])
        for f in best_folders:
            win_counts[f] += 1 / len(best_folders)

    # Head-to-head win matrix
    folders = [s['folder'] for s in series]
    win_matrix = {f: {g: 0 for g in folders} for f in folders}
    for i, b in enumerate(bench_list):
        scores_here = [(s['folder'], s['scores'][i]) for s in series if s['scores'][i] is not None]
        for f, sv in scores_here:
            for g, gv in scores_here:
                if sv > gv:
                    win_matrix[f][g] += 1

    return {
        'output_dir': overview['output_dir'],
        'benchmarks': bench_list,
        'categories': overview['categories'],
        'models': [m for m in overview['models'] if m['folder'] in selected],
        'series': series,
        'ranking': ranking,
        'category_labels': category_labels,
        'category_series': category_series,
        'category_ranking': category_ranking,
        'overall': overall,
        'win_counts': win_counts,
        'win_matrix': win_matrix,
        'matrix': {
            f: {b: overview['matrix'].get(f, {}).get(b) for b in bench_list}
            for f in selected
        },
    }