import os
import numpy as np
from process_single_ldr import process_ldr_data, process_ldr_flatten, process_ldr_flatten_bottom 
import re

# 处理所有ldr文件
def process_all_ldr_in_folder(folder_path, output_folder):
    all_data = []
    all_label_inverse_mappings = []  # 存储所有文件的label_inverse_mapping
    filenames = []  # 存储所有文件的文件名
    
    line_count = 0

    ldr_files = [f for f in os.listdir(input_folder) if 
                 os.path.isfile(os.path.join(input_folder, f)) and 
                 re.match(r'modified_car_\d+_车底_rot\.ldr', f)]

    ldr_files.sort(key=lambda x: int(re.search(r'(\d+)', x).group(1)))

    #import ipdb; ipdb.set_trace()
    # 遍历文件夹中的所有文件夹
    #for file in os.listdir(folder_path):
    for idx, file in enumerate(ldr_files, 1):
        # if file == 'modified_car_81_rot.ldr' or file == 'modified_car_1320_rot.ldr'  or file == 'modified_car_1688_rot.ldr' :
        #     continue 
        if file.endswith('.ldr') and file.startswith('modified'):  # 只处理ldr文件
            file_path = os.path.join(folder_path, file)

            with open(file_path, 'r') as f:
                lines = f.readlines()
                #if len(lines)>110:
                #print(file, len(lines))
                #     #break
                line_count += len(lines)
                
                #data, _ = process_ldr_data(lines)  
                data, _ = process_ldr_flatten_bottom(lines)  
                if data.shape[0]>100:
#                     import ipdb; ipdb.set_trace()
                    print(file, idx)

                print(file, data.shape[0])
                
                #sort
                sort_cols = data[:, [-4, -5, -3]]
                sort_idx = np.lexsort((sort_cols[:, 2], sort_cols[:, 1], sort_cols[:, 0]))
                data = data[sort_idx]

                all_data.append(data)
                filenames.append(file)

            #print(f"Processed {file}")

    print(f"Total lines processed: {line_count}") #1263
    #import ipdb; ipdb.set_trace()  
    output_file = os.path.join(output_folder, "bottom_train.npz")
    np.savez_compressed(output_file, 
                        data=np.array(all_data, dtype=object),  # 将所有数据以列表的形式存储
                        filenames=filenames)  # 也可以保存文件名，方便后续查找

    print(f"All LDR data have been processed and saved to {output_file}")

# 主程序
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'  # 输出文件夹路径
    
    os.makedirs(output_folder, exist_ok=True)  # 确保输出文件夹存在
    
    # 处理所有ldr文件并保存为npz文件
    process_all_ldr_in_folder(input_folder, output_folder)