import json
import os

def process_mappings(original_mapping_path, inverse_mapping_path, frequency_data_path, output_dir=None):
    """
    处理正向和反向映射关系，根据使用频率重新排序并生成新的ID
    
    参数:
        original_mapping_path: 原始映射关系JSON文件路径
        inverse_mapping_path: 反向映射关系JSON文件路径
        frequency_data_path: 包含使用频率数据的JSON文件路径
        output_dir: 输出新映射文件的目录，默认为原文件所在目录
    """
    # 读取原正向映射关系
    with open(original_mapping_path, 'r', encoding='utf-8') as f:
        original_mapping = json.load(f)
    
    # 读取原反向映射关系
    with open(inverse_mapping_path, 'r', encoding='utf-8') as f:
        inverse_mapping = json.load(f)
    
    # 读取频率数据
    with open(frequency_data_path, 'r', encoding='utf-8') as f:
        frequency_data = json.load(f)
    
    # 创建part_name到usage_count的映射
    part_frequency = {item['part_name']: item['usage_count'] for item in frequency_data}
    
    # 提取所有part_name，确保不遗漏任何数据
    all_parts = set(original_mapping.values())
    all_parts.update(inverse_mapping.keys())
    all_parts.update(part_frequency.keys())
    
    # 按照使用频率排序，频率高的在前
    # 对于没有频率数据的part，使用0作为默认值
    sorted_parts = sorted(all_parts, key=lambda x: part_frequency.get(x, 0), reverse=True)
    
    # 生成新的正向映射关系，ID从0开始
    new_forward_mapping = {str(i): part for i, part in enumerate(sorted_parts)}
    
    # 生成新的反向映射关系 (part_name -> ID)
    #new_inverse_mapping = {part: str(i) for i, part in enumerate(sorted_parts)}
    new_inverse_mapping = {part: int(i) for i, part in enumerate(sorted_parts)}
    
    # 确定输出目录
    if output_dir is None:
        output_dir = os.path.dirname(original_mapping_path)
    
    # 构建输出文件路径
    forward_output = os.path.join(output_dir, "new_forward_mapping.json")
    inverse_output = os.path.join(output_dir, "new_inverse_mapping.json")
    
    # 保存新的正向映射
    with open(forward_output, 'w', encoding='utf-8') as f:
        json.dump(new_forward_mapping, f, ensure_ascii=False, indent=4)
    
    # 保存新的反向映射
    with open(inverse_output, 'w', encoding='utf-8') as f:
        json.dump(new_inverse_mapping, f, ensure_ascii=False, indent=4)
    
    print(f"新的正向映射已生成: {forward_output}")
    print(f"新的反向映射已生成: {inverse_output}")
    print(f"共处理 {len(new_forward_mapping)} 个part")
    
    return new_forward_mapping, new_inverse_mapping

if __name__ == "__main__":
    # 示例用法，根据实际情况修改文件路径
    original_mapping_file = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_self/label_mapping.json"  # 原正向映射文件
    inverse_mapping_file = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_self/label_inverse_mapping.json"   # 原反向映射文件
    frequency_data_file = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_self/label_frequency.json"     # 频率数据文件
    
    # 处理映射关系
    process_mappings(original_mapping_file, inverse_mapping_file, frequency_data_file)