import trimesh
import numpy as np
import shutil
import os

def load_obj(file_path):
    """Carrega um arquivo OBJ usando trimesh e retorna a primeira malha encontrada."""
    scene = trimesh.load(file_path, file_type='obj')
    if isinstance(scene, trimesh.Scene):
        # Se for uma cena, extrair a primeira malha
        mesh = trimesh.util.concatenate(scene.dump())
    else:
        mesh = scene
    return mesh

def simplify_mesh(mesh, reduction_factor=0.5):
    """Reduz o número de triângulos da malha."""
    target_faces = int(len(mesh.faces) * reduction_factor)
    simplified_mesh = mesh.simplify_quadric_decimation(target_faces)
    return simplified_mesh

def save_obj_with_materials(mesh, input_obj_path, output_obj_path):
    """Salva a malha em formato OBJ mantendo as referências aos materiais."""
    # Exportar a malha simplificada para um arquivo OBJ temporário
    temp_obj_path = output_obj_path + '.temp.obj'
    mesh.export(temp_obj_path)

    # Ler o conteúdo do arquivo OBJ original
    with open(input_obj_path, 'r') as file:
        original_lines = file.readlines()

    # Ler o conteúdo do arquivo OBJ simplificado temporário
    with open(temp_obj_path, 'r') as file:
        simplified_lines = file.readlines()

    # Encontrar a linha que referencia o arquivo MTL no original
    mtl_line = None
    for line in original_lines:
        if line.startswith('mtllib'):
            mtl_line = line
            break

    # Adicionar a linha do MTL no início do arquivo simplificado
    if mtl_line:
        simplified_lines.insert(0, mtl_line)

    # Preservar as linhas usemtl e associar com as faces
    usemtl_lines = []
    for i, line in enumerate(original_lines):
        if line.startswith('usemtl'):
            usemtl_lines.append((i, line))

    # Inserir as linhas usemtl antes das faces correspondentes
    face_indices = [i for i, line in enumerate(simplified_lines) if line.startswith('f ')]
    for face_index in face_indices:
        for usemtl_index, usemtl_line in usemtl_lines:
            if usemtl_index < face_index:
                simplified_lines.insert(face_index, usemtl_line)
                break

    # Salvar o arquivo OBJ final com a referência ao MTL
    with open(output_obj_path, 'w') as file:
        file.writelines(simplified_lines)

    # Remover o arquivo temporário
    os.remove(temp_obj_path)

def copy_mtl_file(input_obj_path, output_obj_path):
    """Copia o arquivo MTL associado ao OBJ."""
    with open(input_obj_path, 'r') as file:
        lines = file.readlines()
        for line in lines:
            if line.startswith('mtllib'):
                mtl_file = line.split()[1]
                input_mtl_path = os.path.join(os.path.dirname(input_obj_path), mtl_file)
                output_mtl_path = os.path.join(os.path.dirname(output_obj_path), mtl_file)
                shutil.copy(input_mtl_path, output_mtl_path)
                break

def simplify_obj_mesh(input_obj_path, output_obj_path, reduction_factor=0.5):
    """Simplifica a malha de um arquivo OBJ e mantém a referência ao MTL."""
    # Carregar a malha do arquivo OBJ
    mesh = load_obj(input_obj_path)
    
    # Simplificar a malha
    simplified_mesh = simplify_mesh(mesh, reduction_factor)
    
    # Salvar a malha simplificada em formato OBJ mantendo as referências aos materiais
    save_obj_with_materials(simplified_mesh, input_obj_path, output_obj_path)
    
    # Copiar o arquivo MTL associado
    copy_mtl_file(input_obj_path, output_obj_path)

# Caminhos dos arquivos
input_obj_file = 'Mont.obj'
output_obj_file = 'Mont_simplificado.obj'

# Fator de redução (0.5 para reduzir pela metade)
reduction_factor = 0.05

# Simplificar a malha
simplify_obj_mesh(input_obj_file, output_obj_file, reduction_factor)