content: # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
  # except for the third-party components listed below.
  # Hunyuan 3D does not impose any additional limitations beyond what is outlined
  # in the repsective licenses of these third-party components.
  # Users must comply with all terms and conditions of original licenses of these third-party
  # components and must ensure that the usage of the third party components adheres to
  # all relevant laws and regulations.
  
  # For avoidance of doubts, Hunyuan 3D means the large language models and
  # their software and algorithms, including trained model weights, parameters (including
  # optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
  # fine-tuning enabling code and other elements of the foregoing made publicly available
  # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
  
  import argparse
  import igl
  import numpy as np
  import os
  from scipy.stats import truncnorm
  import trimesh
  
  def random_sample_pointcloud(mesh, num = 30000):
      points, face_idx = mesh.sample(num, return_index=True)
      normals = mesh.face_normals[face_idx]
      rng = np.random.default_rng()
      index = rng.choice(num, num, replace=False)
      return points[index], normals[index]
  
  def sharp_sample_pointcloud(mesh, num=16384):
      V = mesh.vertices
      N = mesh.face_normals
      VN = mesh.vertex_normals
      F = mesh.faces
      VN2 = np.ones(V.shape[0])
      for i in range(3):
          dot = np.stack((VN2[F[:,i]], np.sum(VN[F[:,i]] * N, axis=-1)), axis=-1)
          VN2[F[:,i]] = np.min(dot, axis=-1)
  
      sharp_mask = VN2<0.985
      # collect edge
      edge_a = np.concatenate((F[:,0],F[:,1],F[:,2]))
      edge_b = np.concatenate((F[:,1],F[:,2],F[:,0]))
      sharp_edge = ((sharp_mask[edge_a] * sharp_mask[edge_b]))
      edge_a = edge_a[sharp_edge>0]
      edge_b = edge_b[sharp_edge>0]
  
      sharp_verts_a = V[edge_a]
      sharp_verts_b = V[edge_b]
      sharp_verts_an = VN[edge_a]
      sharp_verts_bn = VN[edge_b]
  
      weights = np.linalg.norm(sharp_verts_b - sharp_verts_a, axis=-1)
      weights /= np.sum(weights)
  
      random_number = np.random.rand(num)
      w = np.random.rand(num,1)
      index = np.searchsorted(weights.cumsum(), random_number)
      samples = w * sharp_verts_a[index] + (1 - w) * sharp_verts_b[index]
      normals = w * sharp_verts_an[index] + (1 - w) * sharp_verts_bn[index]
      return samples, normals
  
  def sample_sdf(mesh, random_surface, sharp_surface):
      n_volume_points = sharp_surface.shape[0] * 2
      vol_points = (np.random.rand(n_volume_points, 3) - 0.5) * 2 * 1.05
  
      a, b = -0.25, 0.25
      mu = 0
  
      # get near points (add offset on surface points)
      offset1 = truncnorm.rvs((a - mu) / 0.005, (b - mu) / 0.005, loc=mu, scale=0.005, size=(len(random_surface), 3))
      offset2 = truncnorm.rvs((a - mu) / 0.05, (b - mu) / 0.05, loc=mu, scale=0.05,  size=(len(random_surface), 3))
      random_near_points = np.concatenate([
          random_surface + offset1,
          random_surface + offset2
      ], axis=0)
  
      unit_num = len(sharp_surface) // 6
      sharp_near_points = np.concatenate([
          sharp_surface[:unit_num] + np.random.normal(scale=0.001, size=(unit_num, 3)),
          sharp_surface[unit_num:unit_num*2] + np.random.normal(scale=0.003, size=(unit_num,3)),
          sharp_surface[unit_num*2:unit_num*3] + np.random.normal(scale=0.06, size=(unit_num,3)),
          sharp_surface[unit_num*3:unit_num*4] + np.random.normal(scale=0.01, size=(unit_num,3)),
          sharp_surface[unit_num*4:unit_num*5] + np.random.normal(scale=0.02, size=(unit_num,3)),
          sharp_surface[unit_num*5:] + np.random.normal(scale=0.04, size=(len(sharp_surface)-5*unit_num,3))
      ], axis=0)
  
      np.random.shuffle(random_near_points)
      np.random.shuffle(sharp_near_points)
  
      sign_type = igl.SIGNED_DISTANCE_TYPE_FAST_WINDING_NUMBER
      try:
          vol_sdf, I, C = igl.signed_distance(
              vol_points.astype(np.float32), 
              mesh.vertices, mesh.faces, 
              return_normals=False,
              sign_type=sign_type)
      except:
          vol_sdf, I, C = igl.signed_distance(
              vol_points.astype(np.float32), 
              mesh.vertices, mesh.faces, 
              return_normals=False)
      try:
          random_near_sdf, I, C = igl.signed_distance(
              random_near_points.astype(np.float32), 
              mesh.vertices, mesh.faces, 
              return_normals=False,
              sign_type=sign_type)
      except:
          random_near_sdf, I, C = igl.signed_distance(
              random_near_points.astype(np.float32), 
              mesh.vertices, mesh.faces, 
              return_normals=False)
      try:
          sharp_near_sdf, I, C = igl.signed_distance(
              sharp_near_points.astype(np.float32), 
              mesh.vertices, mesh.faces, 
              return_normals=False,
              sign_type=sign_type)
      except:
          sharp_near_sdf, I, C = igl.signed_distance(
              sharp_near_points.astype(np.float32), 
              mesh.vertices, mesh.faces, 
              return_normals=False)
          
      vol_label = -vol_sdf
      random_near_label = -random_near_sdf
      sharp_near_label = -sharp_near_sdf
  
      data = {
          "vol_points": vol_points.astype(np.float16),
          "vol_label": vol_label.astype(np.float16),
          "random_near_points": random_near_points.astype(np.float16),
          "random_near_label": random_near_label.astype(np.float16),
          "sharp_near_points": sharp_near_points.astype(np.float16),
          "sharp_near_label": sharp_near_label.astype(np.float16)
      }
      return data
  
  def SampleMesh(V, F):
      mesh = trimesh.Trimesh(vertices=V, faces=F)
  
      area = mesh.area
      sample_num = 499712//4
  
      random_surface, random_normal = random_sample_pointcloud(mesh, num=sample_num)
      random_sharp_surface, sharp_normal = sharp_sample_pointcloud(mesh, num=sample_num)
  
      #save_surface
      surface = np.concatenate((random_surface, random_normal), axis = 1).astype(np.float16)
      sharp_surface = np.concatenate((random_sharp_surface, sharp_normal), axis=1).astype(np.float16)
  
      surface_data = {
          "random_surface": surface,
          "sharp_surface": sharp_surface,
      }
  
      sdf_data = sample_sdf(mesh, random_surface, random_sharp_surface)
      return surface_data, sdf_data
  
  def normalize_to_unit_box(V):
      """
      Normalize the vertices V to fit inside a unit bounding box [0,1]^3.
      V: (n,3) numpy array of vertex positions.
      Returns: normalized V
      """
      V_min = V.min(axis=0)
      V_max = V.max(axis=0)
      scale = (V_max - V_min).max() * 1.01
      V_normalized = (V - V_min) / scale
      return V_normalized
  
  # Given: V (n x 3 array of vertices), F (m x 3 array of faces)
  # Parameters epsilon/grid_res
  def Watertight(V, F, epsilon = 2.0/256, grid_res = 256):
      # Compute bounding box
      min_corner = V.min(axis=0)
      max_corner = V.max(axis=0)
      padding = 0.05 * (max_corner - min_corner)
      min_corner -= padding
      max_corner += padding
  
      # Create a uniform grid
      x = np.linspace(min_corner[0], max_corner[0], grid_res)
      y = np.linspace(min_corner[1], max_corner[1], grid_res)
      z = np.linspace(min_corner[2], max_corner[2], grid_res)
      X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
      grid_points = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T
  
      # Compute SDF at grid points using igl.signed_distance with pseudo normals
      sdf, _, _ = igl.signed_distance(
          grid_points, V, F, sign_type=igl.SIGNED_DISTANCE_TYPE_PSEUDONORMAL
      )
   
      # igl.marching_cubes returns (vertices, faces)
      mc_verts, mc_faces = igl.marching_cubes(epsilon - np.abs(sdf), grid_points, grid_res, grid_res, grid_res, 0.0)
  
      # mc_verts: (k x 3) array of vertices of the epsilon contour
      # mc_faces: (l x 3) array of faces of the epsilon contour
      return mc_verts, mc_faces
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser(description='Process an OBJ file and output surface and SDF data.')
      parser.add_argument('--input_obj', type=str, help='Path to the input OBJ file')
      parser.add_argument('--output_prefix', type=str, default=None, 
          help='Base name for output files (default: input OBJ filename without extension)')
      args = parser.parse_args()
  
      input_obj = args.input_obj
      name = args.output_prefix
  
      V, F = igl.read_triangle_mesh(input_obj)
      V = normalize_to_unit_box(V)
  
      mc_verts, mc_faces = Watertight(V, F)
      surface_data, sdf_data = SampleMesh(mc_verts, mc_faces)
  
      parent_folder = os.path.dirname(args.output_prefix)
      os.makedirs(parent_folder, exist_ok=True)
      export_surface = f'{name}_surface.npz'
      np.savez(export_surface, **surface_data)
      export_sdf = f'{name}_sdf.npz'
      np.savez(export_sdf, **sdf_data)
      igl.write_obj(f'{name}_watertight.obj', mc_verts, mc_faces)
  