import os
import numpy as np
import pandas as pd
from math import cos, sin, radians
# 从原有文件导入所需函数
from  rotateLdr import (checkEncoding, readTextFile, 
                          apply_rotation, apply_rotation_matrix,
                          getModulePart, process_ldr)

def create_rotation_matrix(y_rotation_code):
    """
    根据旋转编码创建Y轴旋转矩阵
    0: 不旋转, 1: 绕Y轴90度, 2: 绕Y轴180度, 3: 绕Y轴270度
    """
    rotation_angles = {0: 0, 1: -90, 2: -180, 3: -270}
    y_angle = rotation_angles.get(y_rotation_code, 0)
    y_rad = radians(y_angle)
    
    # 绕Y轴旋转矩阵
    rot_mat = [
        [cos(y_rad), 0, sin(y_rad)],
        [0, 1, 0],
        [-sin(y_rad), 0, cos(y_rad)]
    ]
    return rot_mat, y_angle

def extract_prefix(filename):
    """从文件名中提取前缀（去除扩展名和可能的_车底）"""
    # 先去除文件扩展名
    prefix = os.path.splitext(str(filename))[0]
    # 去除可能存在的_车底后缀
    if prefix.endswith("_车底"):
        prefix = prefix[:-3]  # "_车底"是3个字符，截取前面部分
    return prefix

def batch_process_ldr(input_dir, output_dir, excel_path):
    """
    批量处理LDR文件，通过前缀匹配Excel中的记录
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 读取Excel文件
    try:
        df = pd.read_excel(excel_path, engine='openpyxl')
        # 确保旋转编码列是整数类型
        rotation_code_col = df.columns[1]  # 第二列是旋转编码
        df[rotation_code_col] = df[rotation_code_col].astype(int, errors='ignore')
        
        # 提取文件名前缀并添加到DataFrame
        df['prefix'] = df['Name'].apply(extract_prefix)
        print(f"成功读取Excel文件，共 {len(df)} 条记录")
    except Exception as e:
        print(f"读取Excel文件出错: {e}")
        return
    
    # 获取所有LDR文件
    ldr_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.ldr')]
    print(f"在 {input_dir} 中找到 {len(ldr_files)} 个LDR文件")
    
    # 处理每个LDR文件
    processed_count = 0
    for ldr_file in ldr_files:
        # 提取LDR文件的前缀和扩展名
        file_base, file_ext = os.path.splitext(ldr_file)
        # 构建带_rot后缀的输出文件名 (例如 car_29.ldr → car_29_rot.ldr)
        output_filename = f"{file_base}_rot{file_ext}"
        # 提取LDR文件的前缀
        ldr_prefix = extract_prefix(ldr_file)
        
        # 查找匹配的前缀
        matched_rows = df[df['prefix'] == ldr_prefix]
        if matched_rows.empty:
            print(f"警告: 在Excel中未找到与 {ldr_file} (前缀: {ldr_prefix}) 匹配的记录，跳过")
            continue
        
        # 获取旋转编码
        try:
            rotation_code = int(matched_rows.iloc[0][rotation_code_col])
            if rotation_code not in [0, 1, 2, 3]:
                print(f"警告: {ldr_file} 的旋转编码 {rotation_code} 无效，使用默认值0")
                rotation_code = 0
        except (ValueError, IndexError, KeyError) as e:
            print(f"警告: {ldr_file} 的旋转编码获取失败: {e}，使用默认值0")
            rotation_code = 0
        
        # 创建旋转矩阵
        rot_mat, y_angle = create_rotation_matrix(rotation_code)
        print(f"处理文件: {ldr_file}, 旋转编码: {rotation_code} (绕Y轴旋转 {y_angle} 度)")
        
        # 处理文件
        input_path = os.path.join(input_dir, ldr_file)
        output_path = os.path.join(output_dir, output_filename)
        
        if process_ldr(input_path, output_path, rot_mat):
            processed_count += 1
    
    print(f"批量处理完成，成功处理 {processed_count}/{len(ldr_files)} 个文件")


if __name__ == "__main__":
    INPUT_DIR = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/ldr_bottom_gather_test/"  # 包含多个LDR文件的目录
    OUTPUT_DIR = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/ldr_rot_test/"  # 输出目录
    EXCEL_PATH = "/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/Pose_test.xlsx"  # 包含旋转角度的Excel文件
    batch_process_ldr(INPUT_DIR, OUTPUT_DIR, EXCEL_PATH)