Project: 020_TetSphere_Splatting_Representing_High-Quality_Geom | Function: solve_milp (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Core Mixed Integer Linear Programming optimization implementing set cover problem to find minimal tetrahedral sphere placement for high-quality geometry coverage

**KEY TERMS:** {
  "D": "coverage matrix where D[i,j] = 1 if sphere j covers point i",
  "c": "objective coefficients (ones for minimal sphere count)",
  "A, b": "constraint matrices for coverage requirement Dx >= 1",
  "radius_scaled": "sphere radius after scaling and offset",
  "batch_construct_D": "D matrix construction with coverage analysis"
}

**CONTEXT CODE:**     print(f"inner_set:{inner_set.shape}")
    inner_set_mesh = trimesh.Trimesh(vertices=inner_set.cpu().numpy(), faces=[])
    inner_set_mesh.export(os.path.join(save_path, "inner_set_filtered.obj"))
    #
    dist = batch_cdist(inner_set, point_set, batch_size=500)
    radius = dist.topk(10, largest=False).values.mean(dim=1, keepdim=True)
    radius = radius * radius_scale + offset

    dist = dist.permute(1, 0)  # [N, Nin]
    radius_g = radius.permute(1, 0)  # [1, Nin]

    D = batch_construct_D(radius_g, dist, batch_size=500)
    D = D.cpu().numpy()

    c = np.ones(len(inner_set))
    options = {"disp": True, "time_limit": 300}
    A, b = D, np.ones(len(point_set))
    integrality = np.ones(len(inner_set))
    lb, ub = np.zeros(len(inner_set)), np.ones(len(inner_set))
    variable_bounds = Bounds(lb, ub)
    constraints = LinearConstraint(A, lb=b)
    res_milp = milp(
        c,
        integrality=integrality,
        bounds=variable_bounds,
        constraints=constraints,
        options=options)
    
    res_milp.x = [int(x_i) for x_i in res_milp.x]
    binary_x = res_milp.x
    
    # res_lp = linprog(
    #     c,
    #     A_ub=D * -1,
    #     b_ub=b * -1,
    #     bounds=[(0, 1) for _ in range(len(inner_set))],
    #     options=options
    # )
    # threshold = 0.3  # Adjust threshold based on your problem's requirements
    # binary_x = (res_lp.x > threshold).astype(int)
    # print(res_lp)
    print(np.sum(binary_x))
    value_pos = np.nonzero(binary_x)[0]

    final_selected_pts = inner_set.cpu().numpy()[value_pos]
    final_radius = radius.cpu().numpy()[value_pos]

    return final_selected_pts, final_radius


def solve_milp(inner_set, point_set, radius, radius_scale, offset, options):
    raise NotImplementedError("Function implementation has been masked for testing")

def generate_spheres(input_mesh_filename, sample_mesh_filename, radius_scale, offset, remesh_edge_length, save_path):
    input_mesh = trimesh.load_mesh(input_mesh_filename)
    
    # remesh
    trimeshgeo = pypgo.create_trimeshgeo(
        input_mesh.vertices.flatten(), input_mesh.faces.flatten())

    trimesh_remeshed = pypgo.mesh_isotropic_remeshing(
        trimeshgeo, remesh_edge_length, 5, 180)
    tri_v = pypgo.trimeshgeo_get_vertices(trimesh_remeshed)
    tri_t = pypgo.trimeshgeo_get_triangles(trimesh_remeshed)
    input_mesh_remeshed = trimesh.Trimesh(vertices=tri_v.reshape(-1, 3), faces=tri_t.reshape(-1, 3))
    input_mesh_remeshed.export(os.path.join(save_path, "input_mesh_remeshed.obj"))
    # sample_mesh = trimesh.load_mesh(sample_mesh_filename)
    
    skeleton = get_min_sdf_skel(input_mesh_remeshed)
    candidate_points = skeleton
    print(f"candidate_points:{candidate_points.shape}")
    
    point_set = torch.tensor(input_mesh.vertices).cuda().double()
    inner_set = torch.tensor(candidate_points).cuda().double()
    
    inner_set_mesh = trimesh.Trimesh(vertices=inner_set.cpu().numpy(), faces=[])
    inner_set_mesh.export(os.path.join(save_path, "inner_set_filtered.obj"))
    
    
    #
    options = {"disp": True, "time_limit": 30000, "mip_rel_gap": 0.20}
    dist = batch_cdist(inner_set, point_set, batch_size=500)
    radius = dist.topk(10, largest=False).values.mean(dim=1, keepdim=True)
    
    ## debug
    template_sphere = trimesh.creation.uv_sphere(radius=1.0)
    debug_vtx = []
    debug_faces = []
    vtx_offset = 0
    for ct, r in zip(inner_set.cpu().numpy().tolist(), (radius.cpu().numpy() * radius_scale + offset).tolist()):
        scaled_vertices = template_sphere.vertices.copy() * r + ct
        debug_vtx.append(scaled_vertices)
        
        offset_faces = template_sphere.faces.copy() + vtx_offset
        debug_faces.append(offset_faces)
        
        vtx_offset += len(template_sphere.vertices)

    # Combine all vertices and faces into a single mesh
    debug_mesh = trimesh.Trimesh(
        vertices=np.vstack(debug_vtx),
        faces=np.vstack(debug_faces))
    debug_mesh.export(os.path.join(save_path, "debug_mesh.obj"))
        
    
    res_milp, D, point_set = solve_milp(inner_set, point_set, radius, radius_scale, offset, options)
    
    res_milp.x = [int(x_i) for x_i in res_milp.x]
    binary_x = res_milp.x
    value_pos = np.nonzero(binary_x)[0]
    print("phase 1: ", np.sum(binary_x))
    
    ## check if there are points that are not covered
    covered_flag = D @ binary_x
    uncovered_flag = covered_flag < 0.5
    uncovered_points = point_set[uncovered_flag]
    print(f"uncovered_points:{np.sum(uncovered_flag)}")
    
    if np.sum(uncovered_flag) > 0:
        # solve again
        options = {"disp": True, "time_limit": 30000, "mip_rel_gap": 0.0}
        res_milp_2, _,  _ = solve_milp(inner_set, uncovered_points, radius, radius_scale, offset, options)
        res_milp_2.x = [int(x_i) for x_i in res_milp_2.x]
        binary_x_2 = res_milp_2.x
        print("phase 2: ", np.sum(binary_x_2))
        value_pos_2 = np.nonzero(binary_x_2)[0]
    
        value_pos = np.concatenate([value_pos, value_pos_2])
    
    print("all: ", value_pos.shape)

    final_selected_pts = inner_set.cpu().numpy()[value_pos]
    radius_scaled = radius * radius_scale + offset
    radius_scaled += offset * 0.3
    final_radius = radius_scaled.cpu().numpy()[value_pos]

    return final_selected_pts, final_radius


def main(tgt_path, mesh_name, save_path, radius_scale=1.1, offset=0.06, pc_res=50, surf_res=50, remesh_edge_length=0.08):
    img, mvp = load_data(tgt_path)
    os.makedirs(save_path, exist_ok=True)

    t1 = time.time()
    surf_mesh_path = os.path.join(save_path, "{}_surf.obj".format(mesh_name))
    pt_mesh_path = os.path.join(save_path, "{}_pt.obj".format(mesh_name))

    ################

    # generate_rough_sdf(img, mvp, pc_res,  pt_mesh_path, True)


Complete the function implementation replacing "raise NotImplementedError".