Project: 020_TetSphere_Splatting_Representing_High-Quality_Geom | Function: forward (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:** Differentiable mesh rasterization implementing TetSphere splatting with MVP transformation, anti-aliasing, and gradient-based optimization for high-quality 3D rendering

**KEY TERMS:** {
  "order": "barrier function order ∈ {2,4}",
  "multiplier": "2^(4*|sin(...)| adaptive coefficient",
  "smooth_eng_func": "tetrahedral smoothness energy computation",
  "barrier_coeff": "barrier energy weighting coefficient",
  "tet_sp": "tetrahedral sphere geometry structure"
}

**CONTEXT CODE:**         self.cfg = parse_structured(self.Config, cfg)

        if self.cfg.context_type == "cuda":
            self.glctx = dr.RasterizeCudaContext()
        elif self.cfg.context_type == "gl":
            self.glctx = dr.RasterizeGLContext()
        else:
            raise ValueError("Invalid context type")

        self.device = get_device()

        self.geometry = geometry

        if materials is not None:
            self.materials = materials
        else:
            self.materials = None

        # dummy tensors
        self.ones_surface_v = torch.ones(
            [self.geometry.surface_vid.shape[0], 1]).to(self.device)
        self.zeros_surface_v = torch.zeros(
            [self.geometry.surface_vid.shape[0], 1]).to(self.device)

        self.tri_hash = None

    def transform_pos(self, mtx, pos, is_vec=False):
        t_mtx = torch.from_numpy(mtx).to(self.device) if isinstance(
            mtx, np.ndarray) else mtx.to(self.device)

        if is_vec:
            # posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
            posw = torch.cat([pos, self.zeros_surface_v],
                             axis=1).to(self.device)
            # return torch.matmul(posw, t_mtx.t())[None, ...]
            res = torch.matmul(posw, t_mtx.transpose(1, 2))  # [None, ...]
        else:
            # posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
            posw = torch.cat([pos, self.ones_surface_v],
                             axis=1).to(self.device)

            # return torch.matmul(posw, t_mtx.t())[None, ...]
            res = torch.matmul(posw, t_mtx.transpose(1, 2))  # [None, ...]

            # TODO check
            if self.cfg.is_orhto:
                res[..., 2] /= 6

        return res

    def forward(self, mvp: torch.Tensor,
                only_alpha: bool,
                iter_num: int, resolution: int,
                permute_surface_scheduler=None,
                fit_normal: bool = False, fit_depth: bool = False,
                background: Optional[torch.Tensor] = None,
                campos: Optional[torch.Tensor] = None):
        """Differentiable mesh rasterization implementing TetSphere splatting with MVP transformation, anti-aliasing, and gradient-based optimization for high-quality 3D rendering"""
        raise NotImplementedError("Function implementation has been masked for testing")
    def export(self, path: str, folder: str, texture_res: int = 1024):
        assert self.materials is not None

        os.makedirs(os.path.join(path, folder), exist_ok=True)

        v_pos = self.geometry.tet_v.clone().detach()[self.geometry.surface_vid]
        t_pos_idx = self.geometry.surface_fid
        v_tex = self.geometry.uv
        t_tex_idx = self.geometry.uv_idx

        # uv_clip = v_tex * 2 - 1
        # uv_clip4 = torch.cat([uv_clip,
        #                       torch.zeros_like(uv_clip[..., 0:1]),
        #                       torch.ones_like(uv_clip[..., 0:1])],
        #                      dim=-1)

        # rast, _ = dr.rasterize(self.glctx, uv_clip4[None, ...], t_pos_idx, resolution=[
        #                        texture_res, texture_res])
        # hole_mask = ~(rast[0, :, :, 3] > 0)

        # def uv_padding(image):
        #     uv_padding_size = 2
        #     inpaint_image = (
        #         cv2.inpaint(
        #             (image.detach().cpu().numpy() * 255).astype(np.uint8),
        #             (hole_mask.detach().cpu().numpy() * 255).astype(np.uint8),
        #             uv_padding_size,
        #             cv2.INPAINT_TELEA,
        #         )
        #     )
        #     return inpaint_image

        # gb_pos, _ = dr.interpolate(v_pos[None, ...], rast, t_pos_idx)
        # gb_pos = gb_pos[0]

        # mat_out = self.materials.forward(gb_pos)
        # save_mat = uv_padding(mat_out["color"].clone())
        # import pdb; pdb.set_trace()
        # # cv2.imwrite(os.path.join(path, folder, "exported_surface.png"), save_mat)

        # _save_obj(path=os.path.join(path, folder),
        #           filename="exported_surface.obj",
        #           v_pos=v_pos.cpu().numpy(),
        #           t_pos_idx=t_pos_idx.cpu().numpy(),
        #           v_tex=v_tex.cpu().numpy(),
        #           t_tex_idx=t_tex_idx.cpu().numpy(),
        #           matname="exported_surface",

        #           mtllib="exported_surface.mtl")
        # _save_mtl(path=os.path.join(path, folder),
        #           filename="exported_surface.mtl",
        #           map_Kd=save_mat,
        #           matname="exported_surface",
        #           map_format="png")
        
        with torch.no_grad():
            mat_out = self.materials.forward(v_pos)
        ## export surface obj with vertex color
        surface_mesh = trimesh.Trimesh(vertices=v_pos.cpu().numpy(), faces=t_pos_idx.cpu().numpy())
        surface_mesh.visual.vertex_colors = mat_out["color"].cpu().numpy()
        surface_mesh.export(os.path.join(path, folder, "exported_surface.obj"))

        # use PyMeshLab
        ms = pymeshlab.MeshSet()
        ms.load_new_mesh(os.path.join(path, folder, "exported_surface.obj"))

        # Check if the mesh has vertex colors
        if not ms.current_mesh().has_vertex_color():
            print("The mesh does not have vertex colors.")
            exit(0)

        # Create UV coordinates using a trivial per-triangle parameterization
        ms.compute_texcoord_parametrization_triangle_trivial_per_wedge(border=1)
        # Bake vertex colors to a texture
        ms.transfer_attributes_to_texture_per_vertex(textw=texture_res, texth=texture_res)
        # Save the mesh with the new UV coordinates
        ms.save_current_mesh(os.path.join(path, folder,"exported_surface.obj"), save_textures=True)


Complete the function implementation replacing "raise NotImplementedError".