import os
import json

# 处理单个ldr文件数据（统一转为小写，不区分大小写）
def process_ldr_data(lines, label_mapping, label_inverse_mapping, label_counter):
    filenames = []
    all_labels = []

    for line in lines:
        if line.startswith('1'):  # 只处理零件数据行
            parts = line.split()
            if len(parts) < 15:  # 增加简单的边界检查
                print(f"Skipping invalid line: {line.strip()}")
                continue

            # 核心修改：将文件名统一转为小写（包括后缀DAT/dat）
            filename = parts[14].lower()  # 例如 Brick1.DAT → brick1.dat
            print(f"Processing file: {filename}")
            
            if filename not in label_mapping:
                label_mapping[filename] = label_counter
                label_inverse_mapping[label_counter] = filename
                label_counter += 1
            all_labels.append(label_mapping[filename])

    return label_mapping, label_inverse_mapping, label_counter

# 处理一个文件夹中的所有ldr文件
def process_all_ldr_in_folder(folder_path):
    overall_label_mapping = {}
    overall_label_inverse_mapping = {}
    label_counter = 0

    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.ldr'):
                file_path = os.path.join(root, file)
                with open(file_path, 'r') as f:
                    lines = f.readlines()
                    overall_label_mapping, overall_label_inverse_mapping, label_counter = process_ldr_data(
                        lines, overall_label_mapping, overall_label_inverse_mapping, label_counter)

    return overall_label_mapping, overall_label_inverse_mapping

# 保存mapping到文件
def save_mappings(label_mapping, label_inverse_mapping, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    
    with open(os.path.join(output_dir, 'label_mapping.json'), 'w') as f:
        json.dump(label_mapping, f, indent=4)
    
    with open(os.path.join(output_dir, 'label_inverse_mapping.json'), 'w') as f:
        json.dump(label_inverse_mapping, f, indent=4)


if __name__ == "__main__":
    input_folder = '/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/ldr_rot_expand_trans'
    output_folder = '/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300'

    label_mapping, label_inverse_mapping = process_all_ldr_in_folder(input_folder)
    save_mappings(label_mapping, label_inverse_mapping, output_folder)
    print(f"Label mappings have been saved to {output_folder}")