import json
import os
import re
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def extract_checkpoint_number(dirname):
    """从目录名中提取checkpoint数字"""
    match = re.search(r'checkpoint-(\d+)', dirname)
    if match:
        return int(match.group(1))
    return None

def extract_gsm8k_accuracy(json_file_path):
    """从JSON文件中提取GSM8K严格匹配准确率"""
    try:
        with open(json_file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # 提取GSM8K严格匹配准确率
        if 'results' in data and 'gsm8k' in data['results']:
            gsm8k_results = data['results']['gsm8k']
            if 'exact_match,flexible-extract' in gsm8k_results:
                return gsm8k_results['exact_match,flexible-extract']
    except Exception as e:
        print(f"读取文件 {json_file_path} 时出错: {e}")
    
    return None

def main():
    # 基础目录路径
    base_dir = "eval_results/finetune_qwen_1.5B_nemotron_shizhe_no_reasoning_pad_bd_size_128_8nodes_seq_length_2k_gsm8k"
    
    # 存储结果的列表
    checkpoint_numbers = []
    accuracies = []
    
    # 遍历所有checkpoint目录
    for item in os.listdir(base_dir):
        item_path = os.path.join(base_dir, item)
        if os.path.isdir(item_path) and item.startswith('checkpoint-'):
            checkpoint_num = extract_checkpoint_number(item)
            if checkpoint_num is not None:
                # 查找JSON结果文件
                json_files = [f for f in os.listdir(item_path) if f.endswith('.json')]
                if json_files:
                    json_file_path = os.path.join(item_path, json_files[0])
                    accuracy = extract_gsm8k_accuracy(json_file_path)
                    if accuracy is not None:
                        checkpoint_numbers.append(checkpoint_num)
                        accuracies.append(accuracy)
                        print(f"Checkpoint {checkpoint_num}: {accuracy:.6f}")
    
    # 按checkpoint数字排序
    sorted_data = sorted(zip(checkpoint_numbers, accuracies))
    checkpoint_numbers, accuracies = zip(*sorted_data)
    
    # 转换为numpy数组
    checkpoint_numbers = np.array(checkpoint_numbers)
    accuracies = np.array(accuracies)
    
    # 绘制图表
    plt.figure(figsize=(12, 8))
    
    # 主图表
    # plt.subplot(2, 1, 1)
    plt.plot(checkpoint_numbers, accuracies, 'b-o', linewidth=2, markersize=6, label='GSM8K Acc')
    plt.xlabel('Checkpoint steps')
    plt.ylabel('Accuracy')
    plt.title('GSM8K Acc')
    plt.grid(True, alpha=0.3)
    plt.legend()
    
    # 添加数值标签
    for i, (x, y) in enumerate(zip(checkpoint_numbers, accuracies)):
        plt.annotate(f'{y:.4f}', (x, y), textcoords="offset points", 
                    xytext=(0,10), ha='center', fontsize=8)
    
    # 统计信息
    # plt.subplot(2, 1, 2)
    # plt.bar(range(len(checkpoint_numbers)), accuracies, alpha=0.7, color='skyblue')
    # plt.xlabel('Checkpoint index')
    # plt.ylabel('Accuracy')
    # plt.title('GSM8K Acc Distribution')
    # plt.xticks(range(len(checkpoint_numbers)), [f'{c}' for c in checkpoint_numbers], rotation=45)
    # plt.grid(True, alpha=0.3)
    
    # # 添加统计信息
    mean_acc = np.mean(accuracies)
    max_acc = np.max(accuracies)
    min_acc = np.min(accuracies)
    std_acc = np.std(accuracies)
    
    # plt.text(0.02, 0.98, f'平均准确率: {mean_acc:.6f}\n最高准确率: {max_acc:.6f}\n最低准确率: {min_acc:.6f}\n标准差: {std_acc:.6f}', 
    #          transform=plt.gca().transAxes, verticalalignment='top',
    #          bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    plt.tight_layout()
    plt.savefig('gsm8k_accuracy_analysis.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 打印统计信息
    print(f"\n=== 统计信息 ===")
    print(f"总checkpoint数量: {len(checkpoint_numbers)}")
    print(f"平均准确率: {mean_acc:.6f}")
    print(f"最高准确率: {max_acc:.6f} (Checkpoint {checkpoint_numbers[np.argmax(accuracies)]})")
    print(f"最低准确率: {min_acc:.6f} (Checkpoint {checkpoint_numbers[np.argmin(accuracies)]})")
    print(f"标准差: {std_acc:.6f}")
    
    # 保存数据到CSV文件
    # import pandas as pd
    # df = pd.DataFrame({
    #     'checkpoint': checkpoint_numbers,
    #     'gsm8k_strict_match_accuracy': accuracies
    # })
    # df.to_csv('gsm8k_accuracy_data.csv', index=False)
    # print(f"\n数据已保存到 gsm8k_accuracy_data.csv")

if __name__ == "__main__":
    main()
