import numpy as np
import sys
import os

def compute_vertex_normals(vertices, faces):
    """
    计算顶点法线
    Args:
        vertices: 顶点列表，每个顶点是[x, y, z]
        faces: 面列表，每个面是顶点索引列表（三角形面为3个索引）
    
    Returns:
        normals: 顶点法线列表，每个法线是[nx, ny, nz]
    """
    # 初始化法线数组
    normals = np.zeros_like(vertices, dtype=np.float64)
    
    # 计算每个面的法线并累加到对应的顶点
    for face in faces:
        if len(face) < 3:
            continue
            
        # 获取三角形的三个顶点
        v0 = vertices[face[0]]
        v1 = vertices[face[1]]
        v2 = vertices[face[2]]
        
        # 计算两条边向量
        edge1 = v1 - v0
        edge2 = v2 - v0
        
        # 计算面法线（叉积）
        face_normal = np.cross(edge1, edge2)
        
        # 面法线长度（用于加权平均）
        face_area = np.linalg.norm(face_normal)
        
        # 如果面法线长度不为零，则归一化并加权
        if face_area > 1e-10:
            face_normal_normalized = face_normal / face_area
            
            # 将面法线累加到每个顶点
            for vertex_index in face:
                normals[vertex_index] += face_normal_normalized
    
    # 归一化所有顶点法线
    for i in range(len(normals)):
        length = np.linalg.norm(normals[i])
        if length > 1e-10:
            normals[i] /= length
        else:
            # 如果法线长度为零（可能是孤立顶点），设置为默认向上法线
            normals[i] = np.array([0.0, 1.0, 0.0])
    
    return normals

def parse_obj_file(filename):
    """
    解析OBJ文件
    Returns:
        vertices: 顶点列表
        faces: 面索引列表
        has_normals: 是否已经包含法线
    """
    vertices = []
    faces = []
    existing_normals = []
    has_normals = False
    
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            for line in file:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                
                parts = line.split()
                if not parts:
                    continue
                
                if parts[0] == 'v':  # 顶点
                    if len(parts) >= 4:
                        vertex = [float(parts[1]), float(parts[2]), float(parts[3])]
                        vertices.append(vertex)
                
                elif parts[0] == 'vn':  # 法线
                    if len(parts) >= 4:
                        normal = [float(parts[1]), float(parts[2]), float(parts[3])]
                        existing_normals.append(normal)
                        has_normals = True
                
                elif parts[0] == 'f':  # 面
                    face_vertices = []
                    for part in parts[1:]:
                        # 处理 f v1/vt1/vn1 v2/vt2/vn2 格式
                        vertex_data = part.split('/')
                        if vertex_data[0]:  # 顶点索引
                            # OBJ索引从1开始，转换为从0开始
                            vertex_index = int(vertex_data[0]) - 1
                            face_vertices.append(vertex_index)
                    
                    if len(face_vertices) >= 3:
                        # 如果是四边形或多边形，分解为三角形
                        for i in range(1, len(face_vertices) - 1):
                            faces.append([face_vertices[0], face_vertices[i], face_vertices[i + 1]])
    
    except UnicodeDecodeError:
        print(f"错误: 无法读取文件 {filename}，请确保文件是UTF-8编码。")
        sys.exit(1)
    
    return np.array(vertices), faces, has_normals

def write_obj_file(filename, vertices, faces, normals, original_has_normals=False):
    """
    写入OBJ文件
    """
    with open(filename, 'w', encoding='utf-8') as file:
        file.write("# OBJ文件 - 自动计算法线\n")
        file.write(f"# 顶点数: {len(vertices)}\n")
        file.write(f"# 面数: {len(faces)}\n\n")
        
        # 写入顶点
        # for i, vertex in enumerate(vertices):
        #     file.write(f"v {vertex[0]:.6f} {vertex[1]:.6f} {vertex[2]:.6f}\n")
        for i, vertex in enumerate(vertices):
            file.write(f"v {vertex[0]*100:.6f} {vertex[1]*100:.6f} {vertex[2]*100:.6f}\n")
        
        file.write("\n")
        
        # 写入法线
        for i, normal in enumerate(normals):
            file.write(f"vn {normal[0]:.6f} {normal[1]:.6f} {normal[2]:.6f}\n")
        
        file.write("\n")
        
        # 写入面（使用顶点索引/法线索引格式）
        for face in faces:
            face_line = "f"
            for vertex_index in face:
                # 如果原始文件有法线或者我们计算了法线，使用 v//vn 格式
                face_line += f" {vertex_index + 1}//{vertex_index + 1}"
            file.write(face_line + "\n")

def process_obj_file(input_file, output_file=None):
    """
    处理OBJ文件的主要函数
    """
    if output_file is None:
        base_name = os.path.splitext(input_file)[0]
        output_file = f"{base_name}_with_normals.obj"
    
    print(f"读取文件: {input_file}")
    vertices, faces, has_normals = parse_obj_file(input_file)
    
    print(f"找到 {len(vertices)} 个顶点, {len(faces)} 个面")
    
    if has_normals:
        print("警告: 输入文件已包含法线数据，将被重新计算")
    
    if len(vertices) == 0:
        print("错误: 没有找到顶点数据")
        return False
    
    if len(faces) == 0:
        print("错误: 没有找到面数据")
        return False
    
    print("计算顶点法线中...")
    normals = compute_vertex_normals(vertices, faces)
    
    print("写入输出文件...")
    write_obj_file(output_file, vertices, faces, normals, has_normals)
    
    print(f"完成! 输出文件: {output_file}")
    return True


if __name__ == "__main__":
    input_file = "face.obj"  # 默认输入文件
    output_file = "face_with_normals.obj"  # 默认输出文件
    process_obj_file(input_file, output_file)