import geomagic.app.v3
from geomagic.app.v3.imports import *
import os

def doctor(input_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    stl_files = [f for f in os.listdir(input_dir) if f.endswith('.stl')]
    stl_files.sort()

    for i in range(len(stl_files)):
        stl_file = stl_files[i]
        
        input_path = os.path.join(input_dir, stl_file)
        
        base_name = os.path.splitext(stl_file)[0]
        output_path = os.path.join(output_dir, base_name + '.stl')
        
        geo.open(0, 1, input_path)
        
        geo.mesh_doctor("smallcompsize", 0.0034721, "smalltunnelsize", 0.0017361, "holesize", 0.0017361, "spikesens", 50, "spikelevel", 0.5, "defeatureoption", 2, "fillholeoption", 2, "autoexpand", 2, "operations", "IntersectionCheck+", "SmallComponentCheck+", "SpikeCheck+", "HighCreaseCheck+", "Update", "Auto-Repair", "Update", "Auto-Repair", "Update", "Auto-Repair", "Update", "Auto-Repair")
        
        geo.saveas(output_path, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)

def remesh(input_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    stl_files = [f for f in os.listdir(input_dir) if f.endswith('.stl')]
    stl_files.sort()

    for i in range(len(stl_files)):
        stl_file = stl_files[i]
        input_path = os.path.join(input_dir, stl_file)
        
        base_name = os.path.splitext(stl_file)[0]
        output_path = os.path.join(output_dir, base_name + '.stl')
        
        geo.open(0, 1, input_path)

        geo.remesh(0.0003339, 0, 1, 45, 0.000806, 1, 0, 0, 0.0003339, 1)
        geo.refine_polygons(0, 0, 1)
        geo.optimize_edges()
        geo.enhance_mesh(5, 0, 50)
        geo.make_open_manifold()

        geo.saveas(output_path, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)

def stl2stp(input_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    stl_files = [f for f in os.listdir(input_dir) if f.endswith('.stl')]
    stl_files.sort()

    for i in range(len(stl_files)):
        stl_file = stl_files[i]
        print(f"[{i+1}/{len(stl_files)}] {stl_file}")
        input_path = os.path.join(input_dir, stl_file)
        
        base_name = os.path.splitext(stl_file)[0]
        output_path = os.path.join(output_dir, base_name + '.stp')
        
        geo.open(0, 1, input_path)

        geo.start_exact_surfacing(0, 0, 0, 0)
        geo.construct_patches(1, 1, 1000, 0)
        geo.construct_grids(20, 1)
        geo.fit_surfaces(0, 18, 8.05842e-06, 0.25, 0.5, 2, 0.5, 0, 0, 1, 0, 0)

        geo.saveas(output_path, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)

    print("Finished")

def test_doctor(input_dir):
    stl_files = [f for f in os.listdir(input_dir) if f.endswith('.stl')]
    stl_files.sort()
    geo.open(0, 1, os.path.join(input_dir, stl_files[0]))
    geo.mesh_doctor("smallcompsize", 0.0034721, "smalltunnelsize", 0.0017361, "holesize", 0.0017361, "spikesens", 50, "spikelevel", 0.5, "defeatureoption", 2, "fillholeoption", 2, "autoexpand", 2, "operations", "IntersectionCheck+", "SmallComponentCheck+", "SpikeCheck+", "HighCreaseCheck+", "Update", "Auto-Repair", "Update", "Auto-Repair", "Update", "Auto-Repair", "Update", "Auto-Repair")

def check_stp(input_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    stl_files = [f for f in os.listdir(input_dir) if f.endswith('.stp')]
    stl_files.sort()
    
    for i in range(len(stl_files)):
        stl_file = stl_files[i]
        print(f"[{i+1}/{len(stl_files)}] {stl_file}")
        input_path = os.path.join(input_dir, stl_file)
        
        geo.open(0, 1, input_path)

        geo.fit_model_to_view()
        geo.fit_model_to_view()
        geo.fit_model_to_view()
        geo.fit_model_to_view()
		
        geo.screen_capture_views(os.path.join(output_dir, stl_file.replace('.stp', '.png')))

root_dir = r'root_dir\\'
doctor(root_dir+'stl_2_cut', root_dir+'stl_3_doctor')
remesh(root_dir+'stl_3_doctor', root_dir+'stl_4_remesh')
stl2stp(root_dir+'stl_4_remesh', root_dir+'stp_1')

root_dir = r'root_dir\\'
test_doctor(root_dir+'stl_2_cut')
check_stp(root_dir+'stp_1', root_dir+'stp_2_screen')