import json
from typing import Dict, Any, Optional, List, Union, Tuple

import re

def extract_number_before_year(file_path: str) -> str | None:
    """
    从文件路径中提取年份（四位数字）前的数字
    
    Args:
        file_path: 待处理的文件路径字符串
    
    Returns:
        提取到的数字字符串（如"0"）；若未找到匹配项，返回None
    
    Example:
        >>> path = "/xxx/self_attn_Llama-2-7b-hf_0_2025-11-27.json"
        >>> extract_number_before_year(path)
        '0'
    """
    # 正则模式解释：
    # _(\d+)_ 匹配下划线+一串数字+下划线（括号捕获数字）
    # \d{4}   匹配四位数字（年份，如2025、2024）
    pattern = r'_(\d+)_\d{4}'
    
    # 执行正则匹配
    match = re.search(pattern, file_path)
    
    # 返回结果：有匹配则返回捕获的数字，否则返回None
    return match.group(1) if match else None



def read_json_acc_scores(
    file_path: str, 
    target_keys: Optional[Union[str, List[str]]] = None
) -> Union[float, Tuple[Dict[str, float], float]]:
    """
    读取JSON文件中的results数据，提取并计算acc,none分数
    
    Args:
        file_path: JSON文件的路径
        target_keys: 要提取的key（单个字符串或字符串列表），为None时处理全部results
    
    Returns:
        - 单个key时：返回对应的acc,none分数（float）
        - 多个key/全部时：返回(各key的acc分数字典, 平均分)
    
    Raises:
        FileNotFoundError: 文件不存在时抛出
        json.JSONDecodeError: JSON格式错误时抛出
        KeyError: 指定的key不存在或缺少acc,none字段时抛出
    """

    benchmark_metrics = {
        "piqa": "acc,none",
        "hellaswag": "acc_norm,none",
        "arc_challenge": "acc_norm,none",
        "boolq": "acc,none",
        "winogrande": "acc,none",
    }
    try:
        # 打开并读取JSON文件
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 提取results部分
        results = data.get('results', {})
        if not results:
            raise KeyError("JSON文件中未找到'results'字段")
        
        # 确定要处理的keys
        if target_keys is None:
            process_keys = list(results.keys())  # 全部keys
        elif isinstance(target_keys, str):
            process_keys = [target_keys]  # 单个key转为列表处理
        elif isinstance(target_keys, list):
            process_keys = target_keys  # 列表keys
        else:
            raise TypeError("target_keys参数必须是字符串、列表或None")
        
        # 验证keys是否存在并提取acc,none分数
        acc_scores = {}
        conc_ans = 0 
        for key in benchmark_metrics:
            if key not in results:
                continue

            acc_scores[key] = results[key][benchmark_metrics[key]]
            conc_ans += results[key][benchmark_metrics[key]]

        # 多个key或全部，返回分数字典和平均分
        # average_score = sum(acc_scores.values()) / len(acc_scores)
        return acc_scores, conc_ans
    
    except FileNotFoundError:
        raise FileNotFoundError(f"文件不存在: {file_path}")
    except json.JSONDecodeError:
        raise json.JSONDecodeError("JSON文件格式错误", file_path, 0)

import os
def get_all_paths(folder):
    """用os模块获取路径下所有文件/目录的完整路径"""
    all_paths = []
    # 先获取当前目录下的所有内容
    for name in os.listdir(folder):
        full_path = os.path.join(folder, name)
        all_paths.append(full_path)
        # 如果是目录，递归进去继续获取
    return all_paths

dir_path = \
    "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results2/Qwen2.5-14B-quantization-layer"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results3/Llama-3.1-8B-quantization-layer-mlp"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results2/Llama-3.1-8B-quantization-layer"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results2/Qwen2.5-7B-quantization-layer"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results3/Llama-3.1-8B-quantization-layer-mlp"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results2/Llama-3.1-8B-quantization-layer"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results2/Qwen2.5-7B-quantization-layer"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results3/Llama-3.1-8B-quantization-layer-mlp"
    # "/mnt/bn/life-mllm/users/cxr/quantization/lm-evaluation-harness/results2/Qwen2.5-7B-quantization-layer"

mode = "self_attn"


paths = get_all_paths(dir_path)

filter_paths = [path for path in paths if mode in path]
print(filter_paths)
total_score = [0 for _ in range(len(filter_paths)-1)]
for filter_path in filter_paths:
    final_paths = get_all_paths(filter_path)
    idx = filter_path.split("/")[-1]
    ans = []
    for final_path in final_paths:
        acc_scores, conc_ans = read_json_acc_scores(final_path)
        # print(acc_scores)
        if conc_ans !=0:
            ans.append(conc_ans)  
    print(idx, sum(ans)/len(ans))
    if int(idx.split("_")[-1]) != -1:
        total_score[int(idx.split("_")[-1])] = sum(ans)/len(ans)
print("total_score:", total_score)

index_value_pairs = list(enumerate(total_score))
# 步骤2：按数值从大到小排序（key=lambda x: x[1] 取元组的第二个值作为排序依据，reverse=True 降序）
sorted_pairs = sorted(index_value_pairs, key=lambda x: x[1], reverse=True)

# 步骤3：提取排序后的索引
sorted_indices = [pair[0] for pair in sorted_pairs]
print(sorted_indices)
print(' '.join(str(num) for num in sorted_indices[:14]))
# acc_scores, average_score = read_json_acc_scores(path, "piqa")
# print(average_score)