#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import re
from pathlib import Path
from transformers import AutoTokenizer
import numpy as np
from collections import defaultdict
import time
import multiprocessing as mp
from functools import partial
import os

def analyze_conversations_parallel(input_dir, tokenizer_name="Qwen/Qwen2.5-0.5B-Instruct", 
                                 file_pattern="*.json", num_processes=None):
    """
    并行分析conversation文件中的reasoning数据和token长度分布
    
    Args:
        input_dir: 输入目录路径
        tokenizer_name: tokenizer模型名称
        file_pattern: 文件匹配模式，默认"*.json"
        num_processes: 并行进程数，默认为CPU核心数
    """
    input_dir_path = Path(input_dir)
    
    if not input_dir_path.exists():
        print(f"错误：输入目录 {input_dir_path} 不存在")
        return
    
    # 查找所有匹配的json文件
    json_files = list(input_dir_path.glob(file_pattern))
    
    if not json_files:
        print(f"错误：在目录 {input_dir_path} 中没有找到匹配 '{file_pattern}' 的文件")
        return
    
    # 按文件名排序
    json_files.sort()
    
    # 设置进程数
    if num_processes is None:
        num_processes = min(mp.cpu_count(), len(json_files))
    
    print(f"使用 {num_processes} 个进程进行并行处理")
    print(f"找到 {len(json_files)} 个JSON文件需要处理")
    
    # 准备文件列表
    file_list = []
    for filepath in json_files:
        filename = filepath.name
        file_list.append((filename, str(filepath)))
        print(f"  - {filename}")
    
    print(f"\n将处理 {len(file_list)} 个文件")
    
    # 创建进程池
    with mp.Pool(processes=num_processes) as pool:
        # 使用partial函数固定tokenizer_name参数
        process_file_func = partial(process_single_file, tokenizer_name=tokenizer_name)
        
        # 并行处理文件
        print("开始并行处理文件...")
        start_time = time.time()
        
        results = pool.map(process_file_func, file_list)
        
        elapsed_time = time.time() - start_time
        print(f"并行处理完成，总耗时: {elapsed_time:.2f} 秒")
    
    # 合并结果
    print("正在合并处理结果...")
    merged_results = merge_processing_results(results)
    
    # 生成报告
    generate_report_from_merged_results(merged_results)

def process_single_file(file_info, tokenizer_name):
    """
    处理单个文件的函数，用于并行处理
    
    Args:
        file_info: (filename, filepath) 元组
        tokenizer_name: tokenizer模型名称
    
    Returns:
        dict: 处理结果
    """
    filename, filepath = file_info
    
    print(f"进程 {os.getpid()}: 开始处理文件 {filename}")
    
    # 加载tokenizer（每个进程独立加载）
    try:
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    except Exception as e:
        print(f"进程 {os.getpid()}: Tokenizer加载失败: {e}")
        return None
    
    # 统计变量
    file_instances = 0
    file_reasoning = 0
    file_non_reasoning = 0
    file_token_lengths = []
    file_reasoning_lengths = []
    file_non_reasoning_lengths = []
    file_turns = []
    file_reasoning_turns = []
    file_non_reasoning_turns = []
    
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
            
            if 'instances' not in data:
                print(f"进程 {os.getpid()}: 文件 {filename} 不包含 'instances' 字段")
                return None
            
            instances = data['instances']
            print(f"进程 {os.getpid()}: 文件包含 {len(instances):,} 个对话实例")
            
            for i, instance in enumerate(instances):
                if i % 10000 == 0:
                    print(f"进程 {os.getpid()}: 已处理 {i:,} 个实例...")
                
                # 检查是否为reasoning数据
                is_reasoning = check_reasoning_data(instance)
                
                # 计算token长度
                total_tokens = calculate_token_length(instance, tokenizer)
                
                # 计算对话轮次
                turns = calculate_conversation_turns(instance)
                
                # 更新统计
                file_instances += 1
                file_token_lengths.append(total_tokens)
                file_turns.append(turns)
                
                if is_reasoning:
                    file_reasoning += 1
                    file_reasoning_lengths.append(total_tokens)
                    file_reasoning_turns.append(turns)
                else:
                    file_non_reasoning += 1
                    file_non_reasoning_lengths.append(total_tokens)
                    file_non_reasoning_turns.append(turns)
            
            # 准备返回结果
            result = {
                'filename': filename,
                'total_instances': file_instances,
                'reasoning_instances': file_reasoning,
                'non_reasoning_instances': file_non_reasoning,
                'reasoning_ratio': float(file_reasoning / file_instances if file_instances > 0 else 0),
                'avg_token_length': float(np.mean(file_token_lengths) if file_token_lengths else 0),
                'min_token_length': int(np.min(file_token_lengths) if file_token_lengths else 0),
                'max_token_length': int(np.max(file_token_lengths) if file_token_lengths else 0),
                'std_token_length': float(np.std(file_token_lengths) if file_token_lengths else 0),
                'avg_reasoning_length': float(np.mean(file_reasoning_lengths) if file_reasoning_lengths else 0),
                'avg_non_reasoning_length': float(np.mean(file_non_reasoning_lengths) if file_non_reasoning_lengths else 0),
                'avg_turns': float(np.mean(file_turns) if file_turns else 0),
                'min_turns': int(np.min(file_turns) if file_turns else 0),
                'max_turns': int(np.max(file_turns) if file_turns else 0),
                'avg_reasoning_turns': float(np.mean(file_reasoning_turns) if file_reasoning_turns else 0),
                'avg_non_reasoning_turns': float(np.mean(file_non_reasoning_turns) if file_non_reasoning_turns else 0),
                'token_lengths': file_token_lengths,
                'reasoning_token_lengths': file_reasoning_lengths,
                'non_reasoning_token_lengths': file_non_reasoning_lengths,
                'turns': file_turns,
                'reasoning_turns': file_reasoning_turns,
                'non_reasoning_turns': file_non_reasoning_turns
            }
            
            print(f"进程 {os.getpid()}: 文件 {filename} 处理完成")
            return result
            
    except Exception as e:
        print(f"进程 {os.getpid()}: 处理文件 {filename} 时发生错误: {e}")
        return None

def merge_processing_results(results):
    """
    合并并行处理的结果
    
    Args:
        results: 并行处理的结果列表
    
    Returns:
        dict: 合并后的结果
    """
    # 过滤掉None结果
    valid_results = [r for r in results if r is not None]
    
    if not valid_results:
        print("没有有效的处理结果")
        return None
    
    # 合并统计
    total_instances = sum(r['total_instances'] for r in valid_results)
    reasoning_instances = sum(r['reasoning_instances'] for r in valid_results)
    non_reasoning_instances = sum(r['non_reasoning_instances'] for r in valid_results)
    
    # 合并token长度列表
    all_token_lengths = []
    all_reasoning_token_lengths = []
    all_non_reasoning_token_lengths = []
    
    # 合并轮次列表
    all_turns = []
    all_reasoning_turns = []
    all_non_reasoning_turns = []
    
    # 文件统计
    file_stats = {}
    
    for result in valid_results:
        # 合并token长度
        all_token_lengths.extend(result['token_lengths'])
        all_reasoning_token_lengths.extend(result['reasoning_token_lengths'])
        all_non_reasoning_token_lengths.extend(result['non_reasoning_token_lengths'])
        
        # 合并轮次
        all_turns.extend(result['turns'])
        all_reasoning_turns.extend(result['reasoning_turns'])
        all_non_reasoning_turns.extend(result['non_reasoning_turns'])
        
        # 保存文件统计
        filename = result['filename']
        file_stats[filename] = {
            'total_instances': result['total_instances'],
            'reasoning_instances': result['reasoning_instances'],
            'non_reasoning_instances': result['non_reasoning_instances'],
            'reasoning_ratio': result['reasoning_ratio'],
            'avg_token_length': result['avg_token_length'],
            'min_token_length': result['min_token_length'],
            'max_token_length': result['max_token_length'],
            'std_token_length': result['std_token_length'],
            'avg_reasoning_length': result['avg_reasoning_length'],
            'avg_non_reasoning_length': result['avg_non_reasoning_length'],
            'avg_turns': result['avg_turns'],
            'min_turns': result['min_turns'],
            'max_turns': result['max_turns'],
            'avg_reasoning_turns': result['avg_reasoning_turns'],
            'avg_non_reasoning_turns': result['avg_non_reasoning_turns']
        }
    
    return {
        'total_instances': total_instances,
        'reasoning_instances': reasoning_instances,
        'non_reasoning_instances': non_reasoning_instances,
        'token_lengths': all_token_lengths,
        'reasoning_token_lengths': all_reasoning_token_lengths,
        'non_reasoning_token_lengths': all_non_reasoning_token_lengths,
        'turns': all_turns,
        'reasoning_turns': all_reasoning_turns,
        'non_reasoning_turns': all_non_reasoning_turns,
        'file_stats': file_stats
    }

def generate_report_from_merged_results(merged_results):
    """
    从合并结果生成报告
    """
    if merged_results is None:
        print("没有有效的结果可以生成报告")
        return
    
    generate_report(
        merged_results['file_stats'],
        merged_results['total_instances'],
        merged_results['reasoning_instances'],
        merged_results['non_reasoning_instances'],
        merged_results['token_lengths'],
        merged_results['reasoning_token_lengths'],
        merged_results['non_reasoning_token_lengths'],
        merged_results['turns'],
        merged_results['reasoning_turns'],
        merged_results['non_reasoning_turns']
    )

def analyze_conversations_chunked(input_dir, tokenizer_name="Qwen/Qwen2.5-0.5B-Instruct", 
                                file_pattern="*.json", chunk_size=1000, num_processes=None):
    """
    按实例块进行并行分析，适用于大文件
    
    Args:
        input_dir: 输入目录路径
        tokenizer_name: tokenizer模型名称
        file_pattern: 文件匹配模式，默认"*.json"
        chunk_size: 每个块的大小
        num_processes: 并行进程数
    """
    input_dir_path = Path(input_dir)
    
    if not input_dir_path.exists():
        print(f"错误：输入目录 {input_dir_path} 不存在")
        return
    
    # 查找所有匹配的json文件
    json_files = list(input_dir_path.glob(file_pattern))
    
    if not json_files:
        print(f"错误：在目录 {input_dir_path} 中没有找到匹配 '{file_pattern}' 的文件")
        return
    
    # 按文件名排序
    json_files.sort()
    
    # 设置进程数
    if num_processes is None:
        num_processes = mp.cpu_count()
    
    print(f"使用 {num_processes} 个进程进行分块并行处理")
    print(f"每个块大小: {chunk_size} 个实例")
    print(f"找到 {len(json_files)} 个JSON文件")
    
    # 准备所有实例块
    all_chunks = []
    
    for filepath in json_files:
        filename = filepath.name
        print(f"正在分割文件: {filename}")
        
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
                
                if 'instances' not in data:
                    print(f"警告：文件 {filename} 不包含 'instances' 字段")
                    continue
                
                instances = data['instances']
                print(f"文件包含 {len(instances):,} 个实例")
                
                # 分割实例
                for i in range(0, len(instances), chunk_size):
                    chunk = instances[i:i + chunk_size]
                    chunk_info = {
                        'filename': filename,
                        'chunk_id': i // chunk_size,
                        'instances': chunk,
                        'start_idx': i,
                        'end_idx': min(i + chunk_size, len(instances))
                    }
                    all_chunks.append(chunk_info)
                
        except Exception as e:
            print(f"处理文件 {filename} 时发生错误: {e}")
            continue
    
    if not all_chunks:
        print("没有找到有效的实例块")
        return
    
    print(f"总共创建了 {len(all_chunks)} 个实例块")
    
    # 创建进程池
    with mp.Pool(processes=num_processes) as pool:
        # 使用partial函数固定tokenizer_name参数
        process_chunk_func = partial(process_instance_chunk, tokenizer_name=tokenizer_name)
        
        # 并行处理实例块
        print("开始并行处理实例块...")
        start_time = time.time()
        
        results = pool.map(process_chunk_func, all_chunks)
        
        elapsed_time = time.time() - start_time
        print(f"并行处理完成，总耗时: {elapsed_time:.2f} 秒")
    
    # 合并结果
    print("正在合并处理结果...")
    merged_results = merge_chunk_results(results)
    
    # 生成报告
    generate_report_from_merged_results(merged_results)

def process_instance_chunk(chunk_info, tokenizer_name):
    """
    处理实例块的函数，用于并行处理
    
    Args:
        chunk_info: 包含实例块信息的字典
        tokenizer_name: tokenizer模型名称
    
    Returns:
        dict: 处理结果
    """
    filename = chunk_info['filename']
    chunk_id = chunk_info['chunk_id']
    instances = chunk_info['instances']
    start_idx = chunk_info['start_idx']
    end_idx = chunk_info['end_idx']
    
    print(f"进程 {os.getpid()}: 开始处理 {filename} 的块 {chunk_id} (实例 {start_idx}-{end_idx})")
    
    # 加载tokenizer（每个进程独立加载）
    try:
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    except Exception as e:
        print(f"进程 {os.getpid()}: Tokenizer加载失败: {e}")
        return None
    
    # 统计变量
    chunk_instances = 0
    chunk_reasoning = 0
    chunk_non_reasoning = 0
    chunk_token_lengths = []
    chunk_reasoning_lengths = []
    chunk_non_reasoning_lengths = []
    chunk_turns = []
    chunk_reasoning_turns = []
    chunk_non_reasoning_turns = []
    
    for i, instance in enumerate(instances):
        # 检查是否为reasoning数据
        is_reasoning = check_reasoning_data(instance)
        
        # 计算token长度
        total_tokens = calculate_token_length(instance, tokenizer)
        
        # 计算对话轮次
        turns = calculate_conversation_turns(instance)
        
        # 更新统计
        chunk_instances += 1
        chunk_token_lengths.append(total_tokens)
        chunk_turns.append(turns)
        
        if is_reasoning:
            chunk_reasoning += 1
            chunk_reasoning_lengths.append(total_tokens)
            chunk_reasoning_turns.append(turns)
        else:
            chunk_non_reasoning += 1
            chunk_non_reasoning_lengths.append(total_tokens)
            chunk_non_reasoning_turns.append(turns)
    
    # 准备返回结果
    result = {
        'filename': filename,
        'chunk_id': chunk_id,
        'total_instances': chunk_instances,
        'reasoning_instances': chunk_reasoning,
        'non_reasoning_instances': chunk_non_reasoning,
        'token_lengths': chunk_token_lengths,
        'reasoning_token_lengths': chunk_reasoning_lengths,
        'non_reasoning_token_lengths': chunk_non_reasoning_lengths,
        'turns': chunk_turns,
        'reasoning_turns': chunk_reasoning_turns,
        'non_reasoning_turns': chunk_non_reasoning_turns
    }
    
    print(f"进程 {os.getpid()}: 块 {chunk_id} 处理完成")
    return result

def merge_chunk_results(results):
    """
    合并分块处理的结果
    
    Args:
        results: 分块处理的结果列表
    
    Returns:
        dict: 合并后的结果
    """
    # 过滤掉None结果
    valid_results = [r for r in results if r is not None]
    
    if not valid_results:
        print("没有有效的处理结果")
        return None
    
    # 按文件名分组
    file_groups = defaultdict(list)
    for result in valid_results:
        file_groups[result['filename']].append(result)
    
    # 合并统计
    total_instances = sum(r['total_instances'] for r in valid_results)
    reasoning_instances = sum(r['reasoning_instances'] for r in valid_results)
    non_reasoning_instances = sum(r['non_reasoning_instances'] for r in valid_results)
    
    # 合并token长度列表
    all_token_lengths = []
    all_reasoning_token_lengths = []
    all_non_reasoning_token_lengths = []
    
    # 合并轮次列表
    all_turns = []
    all_reasoning_turns = []
    all_non_reasoning_turns = []
    
    # 文件统计
    file_stats = {}
    
    for filename, chunks in file_groups.items():
        # 合并该文件的所有块
        file_token_lengths = []
        file_reasoning_lengths = []
        file_non_reasoning_lengths = []
        file_turns = []
        file_reasoning_turns = []
        file_non_reasoning_turns = []
        
        file_total_instances = 0
        file_reasoning_instances = 0
        file_non_reasoning_instances = 0
        
        for chunk in chunks:
            file_total_instances += chunk['total_instances']
            file_reasoning_instances += chunk['reasoning_instances']
            file_non_reasoning_instances += chunk['non_reasoning_instances']
            
            file_token_lengths.extend(chunk['token_lengths'])
            file_reasoning_lengths.extend(chunk['reasoning_token_lengths'])
            file_non_reasoning_lengths.extend(chunk['non_reasoning_token_lengths'])
            file_turns.extend(chunk['turns'])
            file_reasoning_turns.extend(chunk['reasoning_turns'])
            file_non_reasoning_turns.extend(chunk['non_reasoning_turns'])
        
        # 合并到全局列表
        all_token_lengths.extend(file_token_lengths)
        all_reasoning_token_lengths.extend(file_reasoning_lengths)
        all_non_reasoning_token_lengths.extend(file_non_reasoning_lengths)
        all_turns.extend(file_turns)
        all_reasoning_turns.extend(file_reasoning_turns)
        all_non_reasoning_turns.extend(file_non_reasoning_turns)
        
        # 保存文件统计
        file_stats[filename] = {
            'total_instances': file_total_instances,
            'reasoning_instances': file_reasoning_instances,
            'non_reasoning_instances': file_non_reasoning_instances,
            'reasoning_ratio': float(file_reasoning_instances / file_total_instances if file_total_instances > 0 else 0),
            'avg_token_length': float(np.mean(file_token_lengths) if file_token_lengths else 0),
            'min_token_length': int(np.min(file_token_lengths) if file_token_lengths else 0),
            'max_token_length': int(np.max(file_token_lengths) if file_token_lengths else 0),
            'std_token_length': float(np.std(file_token_lengths) if file_token_lengths else 0),
            'avg_reasoning_length': float(np.mean(file_reasoning_lengths) if file_reasoning_lengths else 0),
            'avg_non_reasoning_length': float(np.mean(file_non_reasoning_lengths) if file_non_reasoning_lengths else 0),
            'avg_turns': float(np.mean(file_turns) if file_turns else 0),
            'min_turns': int(np.min(file_turns) if file_turns else 0),
            'max_turns': int(np.max(file_turns) if file_turns else 0),
            'avg_reasoning_turns': float(np.mean(file_reasoning_turns) if file_reasoning_turns else 0),
            'avg_non_reasoning_turns': float(np.mean(file_non_reasoning_turns) if file_non_reasoning_turns else 0)
        }
    
    return {
        'total_instances': total_instances,
        'reasoning_instances': reasoning_instances,
        'non_reasoning_instances': non_reasoning_instances,
        'token_lengths': all_token_lengths,
        'reasoning_token_lengths': all_reasoning_token_lengths,
        'non_reasoning_token_lengths': all_non_reasoning_token_lengths,
        'turns': all_turns,
        'reasoning_turns': all_reasoning_turns,
        'non_reasoning_turns': all_non_reasoning_turns,
        'file_stats': file_stats
    }

def check_reasoning_data(instance):
    """
    检查实例是否为reasoning数据（包含<think>和</think>标签）
    """
    if 'messages' not in instance:
        return False
    
    # 检查所有消息中是否包含<think>和</think>
    for message in instance['messages']:
        if 'content' in message:
            content = message['content']
            if '<think>' in content and '</think>' in content:
                return True
    
    return False

def calculate_token_length(instance, tokenizer):
    """
    计算实例的总token长度
    """
    if 'messages' not in instance:
        return 0
    
    total_tokens = 0
    
    for message in instance['messages']:
        if 'content' in message:
            content = message['content']
            # 使用tokenizer计算token数量
            tokens = tokenizer.encode(content, add_special_tokens=False)
            total_tokens += len(tokens)
    
    return total_tokens

def calculate_conversation_turns(instance):
    """
    计算对话轮次数量
    """
    if 'messages' not in instance:
        return 0
    
    # 计算对话轮次：每两个消息为一轮（user + assistant）
    messages = instance['messages']
    if len(messages) < 2:
        return 0
    
    # 计算完整的对话轮次
    turns = len(messages) // 2
    
    return turns

def generate_report(file_stats, total_instances, reasoning_instances, non_reasoning_instances,
                   token_lengths, reasoning_token_lengths, non_reasoning_token_lengths,
                   conversation_turns, reasoning_turns, non_reasoning_turns):
    """
    生成分析报告
    """
    print("\n" + "="*80)
    print("数据分析报告")
    print("="*80)
    
    # 总体统计
    print(f"\n总体统计:")
    print(f"  总对话实例数: {total_instances:,}")
    print(f"  Reasoning实例数: {reasoning_instances:,}")
    print(f"  非Reasoning实例数: {non_reasoning_instances:,}")
    print(f"  Reasoning比例: {reasoning_instances/total_instances*100:.2f}%")
    
    # Token长度统计
    if token_lengths:
        print(f"\nToken长度统计 (总体):")
        print(f"  平均长度: {np.mean(token_lengths):.1f}")
        print(f"  中位数长度: {np.median(token_lengths):.1f}")
        print(f"  最小长度: {np.min(token_lengths):,}")
        print(f"  最大长度: {np.max(token_lengths):,}")
        print(f"  标准差: {np.std(token_lengths):.1f}")
        
        # 分位数统计
        percentiles = [25, 50, 75, 90, 95, 99]
        for p in percentiles:
            value = np.percentile(token_lengths, p)
            print(f"  {p}%分位数: {value:.1f}")
    
    # Reasoning数据统计
    if reasoning_token_lengths:
        print(f"\nReasoning数据Token长度统计:")
        print(f"  平均长度: {np.mean(reasoning_token_lengths):.1f}")
        print(f"  中位数长度: {np.median(reasoning_token_lengths):.1f}")
        print(f"  最小长度: {np.min(reasoning_token_lengths):,}")
        print(f"  最大长度: {np.max(reasoning_token_lengths):,}")
        print(f"  标准差: {np.std(reasoning_token_lengths):.1f}")
    
    # 非Reasoning数据统计
    if non_reasoning_token_lengths:
        print(f"\n非Reasoning数据Token长度统计:")
        print(f"  平均长度: {np.mean(non_reasoning_token_lengths):.1f}")
        print(f"  中位数长度: {np.median(non_reasoning_token_lengths):.1f}")
        print(f"  最小长度: {np.min(non_reasoning_token_lengths):,}")
        print(f"  最大长度: {np.max(non_reasoning_token_lengths):,}")
        print(f"  标准差: {np.std(non_reasoning_token_lengths):.1f}")
    
    # 对话轮次统计
    if conversation_turns:
        print(f"\n对话轮次统计 (总体):")
        print(f"  平均轮次: {np.mean(conversation_turns):.1f}")
        print(f"  中位数轮次: {np.median(conversation_turns):.1f}")
        print(f"  最小轮次: {np.min(conversation_turns):,}")
        print(f"  最大轮次: {np.max(conversation_turns):,}")
        print(f"  标准差: {np.std(conversation_turns):.1f}")
        
        # 轮次分布统计
        turn_counts = defaultdict(int)
        for turns in conversation_turns:
            turn_counts[turns] += 1
        
        print(f"\n轮次分布:")
        for turns in sorted(turn_counts.keys()):
            count = turn_counts[turns]
            percentage = count / len(conversation_turns) * 100
            print(f"  {turns}轮: {count:,} 个 ({percentage:.1f}%)")
    
    # Reasoning数据轮次统计
    if reasoning_turns:
        print(f"\nReasoning数据轮次统计:")
        print(f"  平均轮次: {np.mean(reasoning_turns):.1f}")
        print(f"  中位数轮次: {np.median(reasoning_turns):.1f}")
        print(f"  最小轮次: {np.min(reasoning_turns):,}")
        print(f"  最大轮次: {np.max(reasoning_turns):,}")
        print(f"  标准差: {np.std(reasoning_turns):.1f}")
    
    # 非Reasoning数据轮次统计
    if non_reasoning_turns:
        print(f"\n非Reasoning数据轮次统计:")
        print(f"  平均轮次: {np.mean(non_reasoning_turns):.1f}")
        print(f"  中位数轮次: {np.median(non_reasoning_turns):.1f}")
        print(f"  最小轮次: {np.min(non_reasoning_turns):,}")
        print(f"  最大轮次: {np.max(non_reasoning_turns):,}")
        print(f"  标准差: {np.std(non_reasoning_turns):.1f}")
    
    # 按文件统计
    print(f"\n按文件统计:")
    print(f"{'文件名':<30} {'总实例':<10} {'Reasoning':<10} {'比例(%)':<8} {'平均长度':<10} {'平均轮次':<10}")
    print("-" * 80)
    
    for filename, stats in sorted(file_stats.items()):
        print(f"{filename:<30} {stats['total_instances']:<10,} {stats['reasoning_instances']:<10,} "
              f"{stats['reasoning_ratio']*100:<8.2f} {stats['avg_token_length']:<10.1f} {stats['avg_turns']:<10.1f}")
    
    # 保存详细报告到文件
    save_detailed_report(file_stats, total_instances, reasoning_instances, non_reasoning_instances,
                        token_lengths, reasoning_token_lengths, non_reasoning_token_lengths,
                        conversation_turns, reasoning_turns, non_reasoning_turns)

def save_detailed_report(file_stats, total_instances, reasoning_instances, non_reasoning_instances,
                        token_lengths, reasoning_token_lengths, non_reasoning_token_lengths,
                        conversation_turns, reasoning_turns, non_reasoning_turns):
    """
    保存详细报告到文件
    """
    report_data = {
        'summary': {
            'total_instances': total_instances,
            'reasoning_instances': reasoning_instances,
            'non_reasoning_instances': non_reasoning_instances,
            'reasoning_ratio': float(reasoning_instances / total_instances if total_instances > 0 else 0)
        },
        'token_length_stats': {
            'overall': {
                'mean': float(np.mean(token_lengths)) if token_lengths else 0,
                'median': float(np.median(token_lengths)) if token_lengths else 0,
                'min': int(np.min(token_lengths)) if token_lengths else 0,
                'max': int(np.max(token_lengths)) if token_lengths else 0,
                'std': float(np.std(token_lengths)) if token_lengths else 0,
                'percentiles': {
                    '25': float(np.percentile(token_lengths, 25)) if token_lengths else 0,
                    '50': float(np.percentile(token_lengths, 50)) if token_lengths else 0,
                    '75': float(np.percentile(token_lengths, 75)) if token_lengths else 0,
                    '90': float(np.percentile(token_lengths, 90)) if token_lengths else 0,
                    '95': float(np.percentile(token_lengths, 95)) if token_lengths else 0,
                    '99': float(np.percentile(token_lengths, 99)) if token_lengths else 0
                }
            },
            'reasoning': {
                'mean': float(np.mean(reasoning_token_lengths)) if reasoning_token_lengths else 0,
                'median': float(np.median(reasoning_token_lengths)) if reasoning_token_lengths else 0,
                'min': int(np.min(reasoning_token_lengths)) if reasoning_token_lengths else 0,
                'max': int(np.max(reasoning_token_lengths)) if reasoning_token_lengths else 0,
                'std': float(np.std(reasoning_token_lengths)) if reasoning_token_lengths else 0
            },
            'non_reasoning': {
                'mean': float(np.mean(non_reasoning_token_lengths)) if non_reasoning_token_lengths else 0,
                'median': float(np.median(non_reasoning_token_lengths)) if non_reasoning_token_lengths else 0,
                'min': int(np.min(non_reasoning_token_lengths)) if non_reasoning_token_lengths else 0,
                'max': int(np.max(non_reasoning_token_lengths)) if non_reasoning_token_lengths else 0,
                'std': float(np.std(non_reasoning_token_lengths)) if non_reasoning_token_lengths else 0
            }
        },
        'turn_stats': {
            'overall': {
                'mean': float(np.mean(conversation_turns)) if conversation_turns else 0,
                'median': float(np.median(conversation_turns)) if conversation_turns else 0,
                'min': int(np.min(conversation_turns)) if conversation_turns else 0,
                'max': int(np.max(conversation_turns)) if conversation_turns else 0,
                'std': float(np.std(conversation_turns)) if conversation_turns else 0
            },
            'reasoning': {
                'mean': float(np.mean(reasoning_turns)) if reasoning_turns else 0,
                'median': float(np.median(reasoning_turns)) if reasoning_turns else 0,
                'min': int(np.min(reasoning_turns)) if reasoning_turns else 0,
                'max': int(np.max(reasoning_turns)) if reasoning_turns else 0,
                'std': float(np.std(reasoning_turns)) if reasoning_turns else 0
            },
            'non_reasoning': {
                'mean': float(np.mean(non_reasoning_turns)) if non_reasoning_turns else 0,
                'median': float(np.median(non_reasoning_turns)) if non_reasoning_turns else 0,
                'min': int(np.min(non_reasoning_turns)) if non_reasoning_turns else 0,
                'max': int(np.max(non_reasoning_turns)) if non_reasoning_turns else 0,
                'std': float(np.std(non_reasoning_turns)) if non_reasoning_turns else 0
            }
        },
        'file_stats': file_stats
    }
    
    output_file = "conversation_analysis_report.json"
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(report_data, f, ensure_ascii=False, indent=2)
    
    print(f"\n详细报告已保存到: {output_file}")

def preview_file_content(filepath, num_samples=3):
    """
    预览文件内容，显示几个样本
    """
    print(f"\n预览文件内容: {filepath}")
    print("-" * 50)
    
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
            
            if 'instances' not in data:
                print("文件不包含 'instances' 字段")
                return
            
            instances = data['instances']
            print(f"文件包含 {len(instances):,} 个实例")
            
            for i in range(min(num_samples, len(instances))):
                instance = instances[i]
                is_reasoning = check_reasoning_data(instance)
                
                print(f"\n样本 {i+1} (Reasoning: {is_reasoning}):")
                if 'messages' in instance:
                    for j, message in enumerate(instance['messages']):
                        role = message.get('role', 'unknown')
                        content = message.get('content', '')
                        # 截断长内容
                        if len(content) > 200:
                            content = content[:200] + "..."
                        print(f"  {role}: {content}")
                
    except Exception as e:
        print(f"预览文件时发生错误: {e}")

def analyze_conversations(input_dir, tokenizer_name="Qwen/Qwen2.5-0.5B-Instruct", file_pattern="*.json"):
    """
    分析conversation文件中的reasoning数据和token长度分布（串行版本）
    
    Args:
        input_dir: 输入目录路径
        tokenizer_name: tokenizer模型名称
        file_pattern: 文件匹配模式，默认"*.json"
    """
    input_dir_path = Path(input_dir)
    
    if not input_dir_path.exists():
        print(f"错误：输入目录 {input_dir_path} 不存在")
        return
    
    # 查找所有匹配的json文件
    json_files = list(input_dir_path.glob(file_pattern))
    
    if not json_files:
        print(f"错误：在目录 {input_dir_path} 中没有找到匹配 '{file_pattern}' 的文件")
        return
    
    # 按文件名排序
    json_files.sort()
    print(f"找到 {len(json_files)} 个JSON文件需要处理")
    
    # 加载tokenizer
    print(f"正在加载tokenizer: {tokenizer_name}")
    try:
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        print("Tokenizer加载成功")
    except Exception as e:
        print(f"Tokenizer加载失败: {e}")
        return
    
    # 统计变量
    total_instances = 0
    reasoning_instances = 0
    non_reasoning_instances = 0
    token_lengths = []
    reasoning_token_lengths = []
    non_reasoning_token_lengths = []
    conversation_turns = []  # 对话轮次统计
    reasoning_turns = []     # reasoning数据的对话轮次
    non_reasoning_turns = [] # 非reasoning数据的对话轮次
    
    # 按文件统计
    file_stats = {}
    
    # 处理每个文件
    for filepath in json_files:
        filename = filepath.name
        
        print(f"\n正在分析文件: {filename}")
        start_time = time.time()
        
        file_instances = 0
        file_reasoning = 0
        file_non_reasoning = 0
        file_token_lengths = []
        file_reasoning_lengths = []
        file_non_reasoning_lengths = []
        file_turns = []
        file_reasoning_turns = []
        file_non_reasoning_turns = []
        
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
                
                if 'instances' not in data:
                    print(f"警告：文件 {filename} 不包含 'instances' 字段")
                    continue
                
                instances = data['instances']
                print(f"文件包含 {len(instances):,} 个对话实例")
                
                for i, instance in enumerate(instances):
                    if i % 10000 == 0:
                        print(f"  已处理 {i:,} 个实例...")
                    
                    # 检查是否为reasoning数据
                    is_reasoning = check_reasoning_data(instance)
                    
                    # 计算token长度
                    total_tokens = calculate_token_length(instance, tokenizer)
                    
                    # 计算对话轮次
                    turns = calculate_conversation_turns(instance)
                    
                    # 更新统计
                    file_instances += 1
                    total_instances += 1
                    file_token_lengths.append(total_tokens)
                    token_lengths.append(total_tokens)
                    file_turns.append(turns)
                    conversation_turns.append(turns)
                    
                    if is_reasoning:
                        file_reasoning += 1
                        reasoning_instances += 1
                        file_reasoning_lengths.append(total_tokens)
                        reasoning_token_lengths.append(total_tokens)
                        file_reasoning_turns.append(turns)
                        reasoning_turns.append(turns)
                    else:
                        file_non_reasoning += 1
                        non_reasoning_instances += 1
                        file_non_reasoning_lengths.append(total_tokens)
                        non_reasoning_token_lengths.append(total_tokens)
                        file_non_reasoning_turns.append(turns)
                        non_reasoning_turns.append(turns)
                
                # 保存文件统计
                file_stats[filename] = {
                    'total_instances': file_instances,
                    'reasoning_instances': file_reasoning,
                    'non_reasoning_instances': file_non_reasoning,
                    'reasoning_ratio': float(file_reasoning / file_instances if file_instances > 0 else 0),
                    'avg_token_length': float(np.mean(file_token_lengths) if file_token_lengths else 0),
                    'min_token_length': int(np.min(file_token_lengths) if file_token_lengths else 0),
                    'max_token_length': int(np.max(file_token_lengths) if file_token_lengths else 0),
                    'std_token_length': float(np.std(file_token_lengths) if file_token_lengths else 0),
                    'avg_reasoning_length': float(np.mean(file_reasoning_lengths) if file_reasoning_lengths else 0),
                    'avg_non_reasoning_length': float(np.mean(file_non_reasoning_lengths) if file_non_reasoning_lengths else 0),
                    'avg_turns': float(np.mean(file_turns) if file_turns else 0),
                    'min_turns': int(np.min(file_turns) if file_turns else 0),
                    'max_turns': int(np.max(file_turns) if file_turns else 0),
                    'avg_reasoning_turns': float(np.mean(file_reasoning_turns) if file_reasoning_turns else 0),
                    'avg_non_reasoning_turns': float(np.mean(file_non_reasoning_turns) if file_non_reasoning_turns else 0)
                }
                
                elapsed_time = time.time() - start_time
                print(f"文件处理完成，耗时: {elapsed_time:.2f} 秒")
                
        except Exception as e:
            print(f"处理文件 {filename} 时发生错误: {e}")
            continue
    
    # 生成总体统计报告
    generate_report(file_stats, total_instances, reasoning_instances, non_reasoning_instances,
                   token_lengths, reasoning_token_lengths, non_reasoning_token_lengths,
                   conversation_turns, reasoning_turns, non_reasoning_turns)

def get_first_json_file(input_dir, file_pattern="*.json"):
    """
    获取目录中的第一个JSON文件用于预览
    """
    input_dir_path = Path(input_dir)
    json_files = list(input_dir_path.glob(file_pattern))
    
    if json_files:
        json_files.sort()
        return json_files[0]
    return None

def print_usage_examples():
    """
    打印脚本使用示例
    """
    print("\n" + "="*80)
    print("脚本使用说明")
    print("="*80)
    print("这个脚本可以分析任意目录下的所有JSON对话文件，支持:")
    print("  - 灵活的文件匹配模式")
    print("  - 三种处理模式: 串行、文件级并行、实例块并行")
    print("  - 自动检测reasoning数据（包含<think>标签）")
    print("  - 详细的token长度和对话轮次统计")
    print()
    print("使用示例:")
    print("  1. 分析所有JSON文件:")
    print("     python analyze_conversations.py")
    print()
    print("  2. 指定特定目录:")
    print("     修改脚本中的input_dir变量")
    print()
    print("  3. 使用特定文件模式:")
    print("     修改脚本中的file_pattern变量，如:")
    print("     - '*.json' (所有JSON文件)")  
    print("     - 'nemotron_*.json' (以nemotron开头的JSON文件)")
    print("     - 'converted_*.json' (以converted开头的JSON文件)")
    print()
    print("输出:")
    print("  - 终端显示详细统计报告")
    print("  - 生成conversation_analysis_report.json详细报告文件")
    print("="*80)

if __name__ == "__main__":
    input_dir = "data/nemotron_v2_converted"  # 默认目录，可以修改
    tokenizer_name = "Qwen/Qwen2.5-0.5B-Instruct"
    file_pattern = "*.json"  # 文件匹配模式
    
    print("=== Conversation数据分析工具 ===")
    print(f"输入目录: {input_dir}")
    print(f"Tokenizer: {tokenizer_name}")
    print(f"文件匹配模式: {file_pattern}")
    
    # 显示使用说明
    show_help = input("\n是否显示使用说明? (y/n，默认n): ").strip().lower()
    if show_help == 'y':
        print_usage_examples()
    
    # 检查目录是否存在
    if not Path(input_dir).exists():
        print(f"错误：输入目录 {input_dir} 不存在")
        # 尝试使用用户提供的目录
        custom_dir = input("请输入有效的目录路径: ").strip()
        if custom_dir and Path(custom_dir).exists():
            input_dir = custom_dir
        else:
            print("无效目录，程序退出")
            exit(1)
    
    # 预览第一个JSON文件的内容
    first_file = get_first_json_file(input_dir, file_pattern)
    if first_file:
        print(f"\n预览第一个文件: {first_file.name}")
        preview_file_content(first_file)
    else:
        print(f"目录 {input_dir} 中没有找到JSON文件")
        exit(1)
    
    # 选择处理模式
    print("\n请选择处理模式:")
    print("1. 串行处理 (适合小数据集)")
    print("2. 文件级并行处理 (适合中等数据集)")
    print("3. 实例块并行处理 (适合大数据集)")
    
    try:
        mode = input("请输入选择 (1/2/3，默认为2): ").strip()
        if not mode:
            mode = "2"
        
        if mode == "1":
            print("\n使用串行处理模式...")
            analyze_conversations(input_dir, tokenizer_name, file_pattern)
        elif mode == "2":
            print("\n使用文件级并行处理模式...")
            # 可以设置进程数
            num_processes = input("请输入并行进程数 (默认为CPU核心数): ").strip()
            if num_processes:
                num_processes = int(num_processes)
            else:
                num_processes = None
            
            analyze_conversations_parallel(input_dir, tokenizer_name, file_pattern, num_processes)
        elif mode == "3":
            print("\n使用实例块并行处理模式...")
            # 可以设置块大小和进程数
            chunk_size = input("请输入每个块的大小 (默认为1000): ").strip()
            if chunk_size:
                chunk_size = int(chunk_size)
            else:
                chunk_size = 1000
            
            num_processes = input("请输入并行进程数 (默认为CPU核心数): ").strip()
            if num_processes:
                num_processes = int(num_processes)
            else:
                num_processes = None
            
            analyze_conversations_chunked(input_dir, tokenizer_name, file_pattern, chunk_size, num_processes)
        else:
            print("无效选择，使用默认的文件级并行处理模式...")
            analyze_conversations_parallel(input_dir, tokenizer_name, file_pattern)
    
    except KeyboardInterrupt:
        print("\n用户中断，程序退出")
    except Exception as e:
        print(f"发生错误: {e}")
        print("使用默认的文件级并行处理模式...")
        analyze_conversations_parallel(input_dir, tokenizer_name, file_pattern)
    
    print("\n分析完成！") 