#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
compare_tokens_multi.py

统计多个模型在指定类别文件下（如 coding.json, knowledge.json, math.json 等）
的 total_tokens 平均值。
"""

import json
from pathlib import Path

# ========= 配置区：按需修改 ========= #

# 关键词：必须是 ["coding", "knowledge", "math", "reasoning", "roleplay", "writing"]
KEYWORD = "writing"

# 19个模型文件夹路径
MODEL_DIRS = [
    
    
    
    
    
    
    
    
    
    

    
    
    
    
    
    
    
    
    
]

# =================================== #


def load_json(path: Path):
    """读取 JSON（数组形式）。"""
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def calc_avg_total_tokens(model_file: Path):
    """计算某个模型文件的 total_tokens 平均值"""
    if not model_file.exists():
        print(f"[Skip] 文件不存在: {model_file}")
        return None

    data = load_json(model_file)
    total_sum = 0
    count = 0

    for item in data:
        cost_tokens = item.get("cost_tokens", {})
        total_tokens = cost_tokens.get("total_tokens")
        if total_tokens is not None:
            total_sum += total_tokens
            count += 1

    if count == 0:
        return None

    avg = total_sum / count
    return avg


def main():
    print(f"\n=== 多模型 {KEYWORD}.json 的 total_tokens 平均值 ===")
    for model_dir in MODEL_DIRS:
        model_file = model_dir / f"{KEYWORD}.json"
        avg = calc_avg_total_tokens(model_file)
        if avg is None:
            continue
        print(f"{model_dir.name} 的 {KEYWORD} 文件 total_tokens 平均值为: {avg:.2f}")


if __name__ == "__main__":
    main()
