from typing import *
import time
import numpy as np
import torch
import utils3d
import nvdiffrast.torch as dr
from tqdm import tqdm
import trimesh
import trimesh.visual
import xatlas
import pyvista as pv
from pymeshfix import _meshfix
import igraph
import cv2
from PIL import Image
from .random_utils import sphere_hammersley_sequence
from .render_utils import render_multiview, render_multiview_gpu
from ..renderers import GaussianRenderer
from ..representations import Strivec, Gaussian, MeshExtractResult

from api_spz.core.exceptions import CancelledException

BAKE_SYNC_FREE = True
TEXTURE_OPT_FUSED_ADAM = True


def postprocess_mesh(
    vertices: Union[np.ndarray, torch.Tensor],
    faces: Union[np.ndarray, torch.Tensor],
    simplify: bool = True,
    simplify_ratio: float = 0.9,
    fill_holes: bool = True,
    fill_holes_max_hole_size: float = 0.04,
    fill_holes_max_hole_nbe: int = 32,
    fill_holes_resolution: int = 1024,
    fill_holes_num_views: int = 1000,
    debug: bool = False,
    verbose: bool = False,
    decimate_backend: str = "pyvista",
    timings: Optional[dict] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Postprocess a mesh by simplifying, removing invisible faces, and removing isolated pieces.
    Maintains data on GPU where possible and uses fp16 precision.

    Args:
        vertices: Vertices of the mesh. Shape (V, 3). Can be numpy array or torch tensor.
        faces: Faces of the mesh. Shape (F, 3). Can be numpy array or torch tensor.
        simplify: Whether to simplify the mesh, using quadric edge collapse.
        simplify_ratio: Ratio of faces to keep after simplification.
        fill_holes: Whether to fill holes in the mesh.
        fill_holes_max_hole_size: Maximum area of a hole to fill.
        fill_holes_max_hole_nbe: Maximum number of boundary edges of a hole to fill.
        fill_holes_resolution: Resolution of the rasterization.
        fill_holes_num_views: Number of views to rasterize the mesh.
        verbose: Whether to print progress.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Processed vertices and faces as torch tensors on GPU
    """
    if verbose:
        tqdm.write(f'Before postprocess: {vertices.shape[0]} vertices, {faces.shape[0]} faces')

    # Convert inputs to torch tensors if needed and ensure float32
    if isinstance(vertices, np.ndarray):
        vertices = torch.from_numpy(vertices).float()
    if isinstance(faces, np.ndarray):
        faces = torch.from_numpy(faces)
    
    # Ensure tensors are on GPU and in float32
    vertices = vertices.cuda().float()
    faces = faces.cuda()
    
    started = time.perf_counter()
    if timings is not None:
        timings["faces_in"] = int(faces.shape[0])
    # Simplify
    if simplify and simplify_ratio > 0 and decimate_backend == "cumesh":
        vertices, faces = _decimate_on_gpu(vertices, faces, int(faces.shape[0] * (1 - simplify_ratio)))
    elif simplify and simplify_ratio > 0:
        vertices_cpu = vertices.cpu().numpy()  # Already float32
        faces_cpu = faces.cpu().numpy()
        mesh = pv.PolyData(vertices_cpu, np.concatenate([np.full((faces_cpu.shape[0], 1), 3), faces_cpu], axis=1))
        mesh = mesh.decimate(simplify_ratio, progress_bar=verbose)
        vertices = torch.tensor(mesh.points, device='cuda')  # Will stay float32
        faces = torch.tensor(mesh.faces.reshape(-1, 4)[:, 1:], device='cuda', dtype=torch.int32)
        if verbose:
            tqdm.write(f'After decimate: {vertices.shape[0]} vertices, {faces.shape[0]} faces')
    started = mark_stage(timings, "decimate", started)
    if timings is not None:
        timings["faces_decimated"] = int(faces.shape[0])

    # Remove invisible faces - already operates on GPU
    if fill_holes:
        vertices, faces = _fill_holes(
            vertices, faces,
            max_hole_size=fill_holes_max_hole_size,
            max_hole_nbe=fill_holes_max_hole_nbe,
            resolution=fill_holes_resolution,
            num_views=fill_holes_num_views,
            debug=debug,
            verbose=verbose,
            timings=timings,
        )
        if verbose:
            tqdm.write(f'After remove invisible faces: {vertices.shape[0]} vertices, {faces.shape[0]} faces')
    mark_stage(timings, "fill_holes", started)

    return vertices, faces


def mark_stage(timings: Optional[dict], name: str, started: float) -> float:
    if timings is None:
        return started
    torch.cuda.synchronize()
    now = time.perf_counter()
    timings[name] = round(now - started, 3)
    return now


DECIMATE_START_THRESHOLD = 1e-6
DECIMATE_THRESHOLD_GROWTH = 2.0
DECIMATE_MIN_PROGRESS = 1e-2
DECIMATE_MAX_THRESHOLD = 1e3


def _decimate_on_gpu(vertices: torch.Tensor, faces: torch.Tensor, target_faces: int) -> Tuple[torch.Tensor, torch.Tensor]:
    import cumesh
    mesh = cumesh.CuMesh()
    mesh.init(vertices.float().contiguous(), faces.int().contiguous())
    remaining, threshold = faces.shape[0], DECIMATE_START_THRESHOLD
    while remaining > target_faces and threshold < DECIMATE_MAX_THRESHOLD:
        _, left = mesh.cu_mesh.simplify_step(1e-2, 1e-3, threshold, False)
        if (remaining - left) / remaining < DECIMATE_MIN_PROGRESS:
            threshold *= DECIMATE_THRESHOLD_GROWTH
        remaining = left
    vertices, faces = mesh.read()
    return vertices.float(), faces.int()


def parametrize_mesh(vertices: Union[np.ndarray, torch.Tensor], faces: Union[np.ndarray, torch.Tensor], backend="xatlas"):
    """
    Parametrize a mesh to a texture space, using xatlas.

    Args:
        vertices: Vertices of the mesh. Shape (V, 3). Can be numpy array or torch tensor.
        faces: Faces of the mesh. Shape (F, 3). Can be numpy array or torch tensor.

    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: Remapped vertices, faces, and UV coordinates
    """
    if backend == "cumesh":
        import cumesh
        mesh = cumesh.CuMesh()
        mesh.init(
            torch.as_tensor(vertices, device="cuda", dtype=torch.float32).contiguous(),
            torch.as_tensor(faces, device="cuda", dtype=torch.int32).contiguous(),
        )
        return tuple(value.cpu().numpy() for value in mesh.uv_unwrap())
    if backend != "xatlas":
        raise ValueError(f"Unknown UV backend: {backend}")

    # Convert to numpy if needed
    if isinstance(vertices, torch.Tensor):
        vertices = vertices.detach().cpu().numpy()
    if isinstance(faces, torch.Tensor):
        faces = faces.detach().cpu().numpy()
    
    # Ensure correct dtypes for xatlas
    vertices = vertices.astype(np.float32)
    faces = faces.astype(np.uint32)

    # Run parametrization
    vmapping, indices, uvs = xatlas.parametrize(vertices, faces)

    # Apply remapping
    vertices = vertices[vmapping]
    faces = indices

    return vertices, faces, uvs

GUTTER_DILATION_STEPS = 8


def _pull_push_fill(texture: torch.Tensor, covered: torch.Tensor) -> torch.Tensor:
    values = (texture * covered[..., None]).permute(2, 0, 1)[None]
    weights = covered.float()[None, None]
    for _ in range(GUTTER_DILATION_STEPS):
        value_sum = torch.nn.functional.avg_pool2d(values, 3, 1, 1)
        weight_sum = torch.nn.functional.avg_pool2d(weights, 3, 1, 1)
        reached = (weight_sum > 0) & (weights == 0)
        values = torch.where(reached, value_sum / weight_sum.clamp_min(1e-8), values)
        weights = torch.where(reached, torch.ones_like(weights), weights)
    edge_filled = values[0].permute(1, 2, 0)
    covered = weights[0, 0] > 0
    pyramid = [(values, weights)]
    while min(values.shape[-2:]) > 1:
        values = torch.nn.functional.avg_pool2d(values, 2, ceil_mode=True)
        weights = torch.nn.functional.avg_pool2d(weights, 2, ceil_mode=True)
        pyramid.append((values, weights))
    filled = values / weights.clamp_min(1e-8)
    for values, weights in reversed(pyramid[:-1]):
        coarse = torch.nn.functional.interpolate(filled, size=values.shape[-2:], mode="bilinear", align_corners=False)
        filled = torch.where(weights > 0, values / weights.clamp_min(1e-8), coarse)
    return torch.where(covered[..., None], edge_filled, filled[0].permute(1, 2, 0))


def fill_texture_gutters(texture: torch.Tensor, covered: torch.Tensor, method: str = "gpu") -> np.ndarray:
    """Fill texels outside the UV charts so bilinear and mip sampling at chart borders stays on-color."""
    if method == "cv2":
        image = np.clip(texture.cpu().numpy() * 255, 0, 255).astype(np.uint8)
        hole = (~covered).cpu().numpy().astype(np.uint8)
        return cv2.inpaint(image, hole, 3, cv2.INPAINT_TELEA)
    if method != "gpu":
        raise ValueError(f"Unknown gutter fill: {method}")
    if covered.any():
        texture = _pull_push_fill(texture.float(), covered)
    return np.clip(texture.cpu().numpy() * 255, 0, 255).astype(np.uint8)


def bake_texture(
    vertices: np.array,
    faces: np.array,
    uvs: np.array,
    observations: List[np.array],
    masks: List[np.array],
    extrinsics: List[np.array],
    intrinsics: List[np.array],
    texture_size: int = 2048,
    near: float = 0.1,
    far: float = 10.0,
    mode: Literal['fast', 'opt'] = 'opt',
    lambda_tv: float = 1e-2,
    verbose: bool = False,
    cancel_event=None,
    optimization_steps: int = 1000,
    gutter_fill: str = "gpu",
    texture_init: str = "zeros",
    texture_lr: float = 1e-2,
    timings: Optional[dict] = None,
):
    """
    Bake texture to a mesh from multiple observations.

    Args:
        vertices (np.array): Vertices of the mesh. Shape (V, 3).
        faces (np.array): Faces of the mesh. Shape (F, 3).
        uvs (np.array): UV coordinates of the mesh. Shape (V, 2).
        observations (List[np.array]): List of observations. Each observation is a 2D image. Shape (H, W, 3).
        masks (List[np.array]): List of masks. Each mask is a 2D image. Shape (H, W).
        extrinsics (List[np.array]): List of extrinsics. Shape (4, 4).
        intrinsics (List[np.array]): List of intrinsics. Shape (3, 3).
        texture_size (int): Size of the texture.
        near (float): Near plane of the camera.
        far (float): Far plane of the camera.
        mode (Literal['fast', 'opt']): Mode of texture baking.
        lambda_tv (float): Weight of total variation loss in optimization.
        verbose (bool): Whether to print progress.
    """
    vertices = torch.tensor(vertices).cuda()
    faces = torch.tensor(faces.astype(np.int32)).cuda()
    uvs = torch.tensor(uvs).cuda()
    started = time.perf_counter()
    if isinstance(observations, torch.Tensor):
        observations = list((observations.double() / 255.0).half())
        masks = list(masks)
    else:
        observations = [torch.tensor(obs / 255.0, dtype=torch.float16).cuda() for obs in observations] # Keep observations as float16 if desired for memory
        masks = [torch.tensor(m>0).bool().cuda() for m in masks]
    views = [utils3d.torch.extrinsics_to_view(torch.tensor(extr).cuda()) for extr in extrinsics]
    projections = [utils3d.torch.intrinsics_to_perspective(torch.tensor(intr).cuda(), near, far) for intr in intrinsics]

    if mode == 'fast' or texture_init == 'fast':
        texture = torch.zeros((texture_size * texture_size, 3), dtype=torch.float32).cuda()
        texture_weights = torch.zeros((texture_size * texture_size), dtype=torch.float32).cuda()
        rastctx = utils3d.torch.RastContext(backend='cuda')
        for observation, view, projection, mask_in in tqdm(zip(observations, views, projections, masks), total=len(observations), disable=not verbose, desc='Texture baking (fast)'):
            if cancel_event and cancel_event.is_set(): 
                raise CancelledException(f"Cancelled the texture baking (fast).")
            with torch.no_grad():
                rast = utils3d.torch.rasterize_triangle_faces(
                    rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection
                )
                uv_map = rast['uv'][0].detach().flip(0)
                mask = rast['mask'][0].detach().bool() & mask_in
            
            # nearest neighbor interpolation
            uv_map = (uv_map * texture_size).floor().long()
            obs = observation[mask]
            uv_map = uv_map[mask]
            idx = uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size
            texture = texture.scatter_add(0, idx.view(-1, 1).expand(-1, 3), obs.float())
            texture_weights = texture_weights.scatter_add(0, idx, torch.ones((obs.shape[0]), dtype=torch.float32, device=texture.device))

        mask = texture_weights > 0
        texture[mask] /= texture_weights[mask][:, None]
        texture = texture.reshape(texture_size, texture_size, 3)
        covered = mask.reshape(texture_size, texture_size)
        if mode == 'fast':
            return fill_texture_gutters(texture, covered, gutter_fill)
        initial_texture = texture.flip(0)

    if mode == 'opt':
        rastctx = utils3d.torch.RastContext(backend='cuda')
        observations = [observations.flip(0) for observations in observations]
        masks = [m.flip(0) for m in masks]
        _uv = []
        _uv_dr = []
        for observation, view, projection in tqdm(zip(observations, views, projections), total=len(views), disable=not verbose, desc='Texture baking (opt): UV'):
            if cancel_event and cancel_event.is_set(): 
                raise CancelledException(f"Cancelled the texture baking (opt).")
            with torch.no_grad():
                rast = utils3d.torch.rasterize_triangle_faces(
                    rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection
                )
                _uv.append(rast['uv'].detach())
                _uv_dr.append(rast['uv_dr'].detach())

        if BAKE_SYNC_FREE:
            pixels = [mask.reshape(-1).nonzero().squeeze(1) for mask in masks]
            targets = [obs.reshape(-1, 3).index_select(0, idx) for obs, idx in zip(observations, pixels)]
        initial = initial_texture[None] if texture_init == 'fast' else torch.zeros((1, texture_size, texture_size, 3))
        texture = torch.nn.Parameter(initial.to(dtype=torch.float32, device='cuda').contiguous())
        optimizer = torch.optim.Adam([texture], betas=(0.5, 0.9), lr=texture_lr, fused=TEXTURE_OPT_FUSED_ADAM or None)
        started = mark_stage(timings, "texture_prep", started)

        def exp_anealing(optimizer, step, total_steps, start_lr, end_lr):
            return start_lr * (end_lr / start_lr) ** (step / total_steps)

        def cosine_anealing(optimizer, step, total_steps, start_lr, end_lr):
            return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps))
        
        def tv_loss(texture):
            return torch.nn.functional.l1_loss(texture[:, :-1, :, :], texture[:, 1:, :, :]) + \
                   torch.nn.functional.l1_loss(texture[:, :, :-1, :], texture[:, :, 1:, :])
    
        max_mip_level = (texture_size & -texture_size).bit_length() - 1 if texture_size & (texture_size - 1) else None
        total_steps = optimization_steps
        with tqdm(total=total_steps, disable=not verbose, desc='Texture baking (opt): optimizing') as pbar:
            for step in range(total_steps):
                optimizer.zero_grad()
                selected = np.random.randint(0, len(views))
                uv, uv_dr, observation, mask = _uv[selected], _uv_dr[selected], observations[selected], masks[selected]
                render = dr.texture(texture, uv, uv_dr, max_mip_level=max_mip_level)[0]
                if BAKE_SYNC_FREE:
                    loss = torch.nn.functional.l1_loss(render.reshape(-1, 3).index_select(0, pixels[selected]), targets[selected])
                else:
                    loss = torch.nn.functional.l1_loss(render[mask], observation[mask])
                if lambda_tv > 0:
                    loss += lambda_tv * tv_loss(texture)
                loss.backward()
                optimizer.step()
                # annealing
                optimizer.param_groups[0]['lr'] = cosine_anealing(optimizer, step, total_steps, texture_lr, 1e-5)
                if verbose:
                    pbar.set_postfix({'loss': loss.item()})
                pbar.update()
                if cancel_event and cancel_event.is_set(): 
                    raise CancelledException(f"Cancelled texture optimization at step {step}/{total_steps}.")
        started = mark_stage(timings, "texture_opt", started)
        covered = utils3d.torch.rasterize_triangle_faces(
            rastctx, (uvs * 2 - 1)[None], faces, texture_size, texture_size
        )['mask'][0].detach() > 0
        texture = fill_texture_gutters(texture[0].flip(0).detach(), covered, gutter_fill)
        mark_stage(timings, "texture_gutter", started)
    else:
        raise ValueError(f'Unknown mode: {mode}')

    return texture


def to_glb(
    app_rep: Union[Strivec, Gaussian],
    mesh: MeshExtractResult,
    simplify: float = 0.95,
    fill_holes: bool = True,
    fill_holes_max_size: float = 0.04,
    texture_size: int = 1024,
    debug: bool = False,
    verbose: bool = False,
    cancel_event = None,
    fill_holes_num_views: int = 1000,
    texture_optimization_steps: int = 1000,
    uv_backend: str = "xatlas",
    texture_mode: str = "opt",
    render_views: int = 30,
    render_resolution: int = 1024,
    gutter_fill: str = "gpu",
    decimate_backend: str = "pyvista",
    timings: Optional[dict] = None,
    texture_init: str = "zeros",
    texture_lr: float = 1e-2,
) -> trimesh.Trimesh:
    """
    Convert a generated asset to a glb file.

    Args:
        app_rep (Union[Strivec, Gaussian]): Appearance representation.
        mesh (MeshExtractResult): Extracted mesh.
        simplify (float): Ratio of faces to remove in simplification.
        fill_holes (bool): Whether to fill holes in the mesh.
        fill_holes_max_size (float): Maximum area of a hole to fill.
        texture_size (int): Size of the texture.
        debug (bool): Whether to print debug information.
        verbose (bool): Whether to print progress.
    """
    vertices = mesh.vertices
    faces = mesh.faces
    
    
    vertices, faces = postprocess_mesh(
        vertices, faces,
        simplify=simplify > 0,
        simplify_ratio=simplify,
        fill_holes=fill_holes,
        fill_holes_max_hole_size=fill_holes_max_size,
        fill_holes_max_hole_nbe=int(250 * np.sqrt(1-simplify)),
        fill_holes_resolution=1024,
        fill_holes_num_views=fill_holes_num_views,
        debug=debug,
        verbose=verbose,
        decimate_backend=decimate_backend,
        timings=timings,
    )
    started = time.perf_counter()

    # parametrize mesh (converts to CPU numpy internally)
    vertices, faces, uvs = parametrize_mesh(vertices, faces, backend=uv_backend)
    started = mark_stage(timings, "uv", started)

    # bake texture
    if BAKE_SYNC_FREE:
        observations, extrinsics, intrinsics = render_multiview_gpu(app_rep, resolution=render_resolution, nviews=render_views)
        masks = (observations > 0).any(dim=-1)
    else:
        observations, extrinsics, intrinsics = render_multiview(app_rep, resolution=render_resolution, nviews=render_views)
        masks = [np.any(observation > 0, axis=-1) for observation in observations]
    extrinsics = [extrinsics[i].cpu().numpy() for i in range(len(extrinsics))]
    intrinsics = [intrinsics[i].cpu().numpy() for i in range(len(intrinsics))]
    started = mark_stage(timings, "render", started)
    texture = bake_texture(
        vertices, faces, uvs,
        observations, masks, extrinsics, intrinsics,
        texture_size=texture_size, mode=texture_mode,
        lambda_tv=0.01,
        verbose=verbose,
        cancel_event=cancel_event,
        optimization_steps=texture_optimization_steps,
        gutter_fill=gutter_fill,
        texture_init=texture_init,
        texture_lr=texture_lr,
        timings=timings,
    )
    mark_stage(timings, "texture", started)
    if timings is not None:
        timings["faces"] = int(faces.shape[0])
    texture = Image.fromarray(texture)

    # rotate mesh (from z-up to y-up)
    vertices = vertices @ np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
    material = trimesh.visual.material.PBRMaterial(
        metallicFactor=0.0,
        roughnessFactor=1.0,
        baseColorTexture=texture,
        baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8)
    )
    mesh = trimesh.Trimesh(vertices, faces, visual=trimesh.visual.TextureVisuals(uv=uvs, material=material))
    return mesh


def simplify_gs(
    gs: Gaussian,
    simplify: float = 0.95,
    verbose: bool = True,
):
    """
    Simplify 3D Gaussians with aggressive pruning
    NOTE: this function is not used in the current implementation for the unsatisfactory performance.
    
    Args:
        gs (Gaussian): 3D Gaussian.
        simplify (float): Ratio of Gaussians to remove in simplification.
    """
    # if simplify <= 0:
    #     return gs
    
    try:
        # Move everything to CUDA and ensure float32 precision
        device = torch.device('cuda')
        
        # Convert all gaussian parameters to float32
        gs._features_dc = gs._features_dc.float()
        if gs._features_rest is not None:
            gs._features_rest = gs._features_rest.float()
        gs._opacity = gs._opacity.float()
        gs._rotation = gs._rotation.float()
        gs._scaling = gs._scaling.float()
        gs._xyz = gs._xyz.float()
        
        # Get initial opacity and ensure proper dimensions
        initial_opacity = gs.get_opacity.squeeze()
        if initial_opacity.dim() == 0:
            initial_opacity = initial_opacity.unsqueeze(0)
        
        # More aggressive initial pruning
        with torch.no_grad():
            opacity_threshold = 0.1  # Increased from 0.05
            initial_mask = initial_opacity > opacity_threshold
            
            # Handle case where no points meet threshold
            if not initial_mask.any():
                num_keep = max(int(0.1 * initial_opacity.shape[0]), 1)
                _, top_indices = initial_opacity.topk(num_keep)
                initial_mask = torch.zeros_like(initial_mask, dtype=torch.bool)
                initial_mask[top_indices] = True
            
            # Apply mask and ensure at least one point remains
            if initial_mask.sum() == 0:
                max_idx = torch.argmax(initial_opacity)
                initial_mask[max_idx] = True
            
            gs._xyz = gs._xyz[initial_mask]
            gs._rotation = gs._rotation[initial_mask]
            gs._scaling = gs._scaling[initial_mask]
            gs._opacity = gs._opacity[initial_mask]
            gs._features_dc = gs._features_dc[initial_mask]
            gs._features_rest = gs._features_rest[initial_mask] if gs._features_rest is not None else None
            
            if verbose:
                print(f"Initial pruning: kept {initial_mask.sum().item()} points out of {len(initial_mask)}")
        
        # Early return if too few points
        if gs._xyz.shape[0] < 2:
            if verbose:
                print("Too few points remain after initial pruning, returning original gaussian")
            return gs
        
        # Render multiview observations with reduced views
        observations, extrinsics, intrinsics = render_multiview(gs, resolution=512, nviews=30)
        observations = [torch.tensor(obs / 255.0, dtype=torch.float32, device=device).permute(2, 0, 1) for obs in observations]
        extrinsics = [e.float() for e in extrinsics]
        intrinsics = [i.float() for i in intrinsics]
        
        # Initialize renderer with smaller resolution
        renderer = GaussianRenderer({
            "resolution": 512,
            "near": 0.8,
            "far": 1.6,
            "ssaa": 1,
            "bg_color": (0,0,0),
        })

        # Clone Gaussian parameters
        new_gs = Gaussian(**gs.init_params)
        new_gs._features_dc = gs._features_dc.clone().to(device, dtype=torch.float32)
        new_gs._features_rest = gs._features_rest.clone().to(device, dtype=torch.float32) if gs._features_rest is not None else None
        new_gs._opacity = torch.nn.Parameter(gs._opacity.clone().to(device, dtype=torch.float32))
        new_gs._rotation = torch.nn.Parameter(gs._rotation.clone().to(device, dtype=torch.float32))
        new_gs._scaling = torch.nn.Parameter(gs._scaling.clone().to(device, dtype=torch.float32))
        new_gs._xyz = torch.nn.Parameter(gs._xyz.clone().to(device, dtype=torch.float32))
        
        # Get initial point count and set target
        current_points = new_gs._xyz.shape[0]
        target_ratio = max(0.1, 1 - simplify * 1.2)  # Ensure we keep at least 10% of points
        num_target = max(int(target_ratio * current_points), 2)  # Ensure at least 2 points remain
        
        if verbose:
            print(f"Starting optimization with {current_points} points, target: {num_target}")
        
        # Optimization parameters
        start_lr = [5e-4, 5e-3, 0.025, 0.1]
        end_lr = [5e-6, 5e-5, 0.00025, 0.001]
        optimizer = torch.optim.Adam([
            {"params": new_gs._xyz, "lr": start_lr[0]},
            {"params": new_gs._rotation, "lr": start_lr[1]},
            {"params": new_gs._scaling, "lr": start_lr[2]},
            {"params": new_gs._opacity, "lr": start_lr[3]},
        ], lr=start_lr[0])
        
        def cosine_anealing(optimizer, step, total_steps, start_lr, end_lr):
            return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps))
        
        _zeta = new_gs.get_opacity.clone().detach().squeeze().float()
        if _zeta.dim() == 0:
            _zeta = _zeta.unsqueeze(0)
        _lambda = torch.zeros_like(_zeta, dtype=torch.float32)
        _delta = 1e-6
        _interval = 5
        
        total_steps = 500
        with tqdm(total=total_steps, disable=not verbose, desc='Simplifying Gaussian') as pbar:
            for i in range(total_steps):
                try:
                    # More frequent pruning
                    if i % 50 == 0 and new_gs._xyz.shape[0] > 2:  # Only prune if we have enough points
                        with torch.cuda.amp.autocast(enabled=False):
                            opacity = new_gs.get_opacity.squeeze()
                            if opacity.dim() == 0:
                                opacity = opacity.unsqueeze(0)
                            
                            mask = opacity > opacity_threshold
                            if not mask.any():  # If all would be pruned, keep top points
                                num_keep = max(int(0.1 * len(mask)), 2)
                                _, top_indices = opacity.topk(min(num_keep, len(mask)))
                                mask = torch.zeros_like(mask, dtype=torch.bool)
                                mask[top_indices] = True
                            
                            # Ensure we keep at least 2 points
                            if mask.sum() < 2:
                                _, top_indices = opacity.topk(2)
                                mask = torch.zeros_like(mask, dtype=torch.bool)
                                mask[top_indices] = True
                            
                            # Apply mask
                            new_gs._xyz = torch.nn.Parameter(new_gs._xyz[mask].float())
                            new_gs._rotation = torch.nn.Parameter(new_gs._rotation[mask].float())
                            new_gs._scaling = torch.nn.Parameter(new_gs._scaling[mask].float())
                            new_gs._opacity = torch.nn.Parameter(new_gs._opacity[mask].float())
                            new_gs._features_dc = new_gs._features_dc[mask].float()
                            new_gs._features_rest = new_gs._features_rest[mask].float() if new_gs._features_rest is not None else None
                            
                            # Update optimization variables
                            _zeta = _zeta[mask].float()
                            _lambda = _lambda[mask].float()
                            
                            # Update optimizer state
                            optimizer = torch.optim.Adam([
                                {"params": new_gs._xyz, "lr": start_lr[0]},
                                {"params": new_gs._rotation, "lr": start_lr[1]},
                                {"params": new_gs._scaling, "lr": start_lr[2]},
                                {"params": new_gs._opacity, "lr": start_lr[3]},
                            ], lr=start_lr[0])

                    with torch.cuda.amp.autocast(enabled=False):
                        opacity = new_gs.get_opacity.squeeze().float()
                        if opacity.dim() == 0:
                            opacity = opacity.unsqueeze(0)
                        current_points = opacity.shape[0]
                        
                        # Sparsification
                        if i % _interval == 0 and current_points > 2:
                            _zeta = _lambda + opacity.detach()
                            if current_points > num_target:
                                k = min(num_target, current_points - 2)  # Keep at least 2 points
                                if k > 0:
                                    index = _zeta.topk(k)[1]
                                    _m = torch.ones_like(_zeta, dtype=torch.bool)
                                    _m[index] = 0
                                    _zeta[_m] = 0
                            _lambda = _lambda + opacity.detach() - _zeta
                        
                        # Sample random view
                        view_idx = np.random.randint(len(observations))
                        observation = observations[view_idx].float()
                        extrinsic = extrinsics[view_idx].float()
                        intrinsic = intrinsics[view_idx].float()
                        
                        # Render and compute loss
                        color = renderer.render(new_gs, extrinsic, intrinsic)['color'].float()
                        rgb_loss = torch.nn.functional.l1_loss(color, observation)
                        loss = rgb_loss + _delta * torch.sum(torch.pow(_lambda + opacity - _zeta, 2))
                        
                        # Optimization step
                        optimizer.zero_grad()
                        loss.backward()
                        optimizer.step()
                        
                        # Update learning rates
                        for j in range(len(optimizer.param_groups)):
                            optimizer.param_groups[j]['lr'] = cosine_anealing(optimizer, i, total_steps, start_lr[j], end_lr[j])
                        
                        # Update progress bar
                        if not torch.isnan(rgb_loss).any():
                            pbar.set_postfix({
                                'loss': rgb_loss.item(),
                                'num': current_points,
                                'lambda': _lambda.mean().item()
                            })
                        
                    pbar.update()
                    
                except RuntimeError as e:
                    if "out of memory" in str(e):
                        torch.cuda.empty_cache()
                        continue
                    else:
                        raise e

        # Final pruning with safety check
        with torch.no_grad():
            opacity = new_gs.get_opacity.squeeze()
            if opacity.dim() == 0:
                opacity = opacity.unsqueeze(0)
            
            final_mask = opacity > opacity_threshold
            if not final_mask.any() or final_mask.sum() < 2:
                num_keep = max(int(0.1 * len(opacity)), 2)
                _, top_indices = opacity.topk(min(num_keep, len(opacity)))
                final_mask = torch.zeros_like(final_mask, dtype=torch.bool)
                final_mask[top_indices] = True
            
            new_gs._xyz = new_gs._xyz.data[final_mask].float()
            new_gs._rotation = new_gs._rotation.data[final_mask].float()
            new_gs._scaling = new_gs._scaling.data[final_mask].float()
            new_gs._opacity = new_gs._opacity.data[final_mask].float()
            new_gs._features_dc = new_gs._features_dc[final_mask].float()
            new_gs._features_rest = new_gs._features_rest[final_mask].float() if new_gs._features_rest is not None else None
            
            if verbose:
                print(f"Final number of points: {final_mask.sum().item()}")
        
        return new_gs
        
    except Exception as e:
        print(f"Error in simplify_gs: {str(e)}")
        print(f"Error details: {str(e.__class__.__name__)}")
        import traceback
        traceback.print_exc()
        return gs  # Return original gaussian if simplification fails


def _visible_outer_faces(visibility: torch.Tensor, components: List[torch.Tensor]) -> torch.Tensor:
    """Per component, mark faces seen more often than its 75th visibility percentile clamped to [0.25, 0.5]."""
    sizes = torch.tensor([len(c) for c in components], device=visibility.device)
    members = torch.cat(components).long()
    component = torch.repeat_interleave(torch.arange(len(components), device=visibility.device), sizes)
    order = torch.argsort(component.double() * 2 + visibility[members].double())
    ranked = visibility[members][order]
    starts = torch.cumsum(sizes, 0) - sizes
    position = 0.75 * (sizes - 1).float()
    low = position.floor().long()
    high = position.ceil().long()
    quantile = ranked[starts + low] + (ranked[starts + high] - ranked[starts + low]) * (position - low.float())
    threshold = quantile.clamp(min=0.25).clamp(max=0.5)
    outer = torch.zeros_like(visibility, dtype=torch.bool)
    outer[members] = visibility[members] > threshold[component]
    return outer


@torch.no_grad()
def _fill_holes(
    verts,
    faces,
    max_hole_size=0.04,
    max_hole_nbe=32,
    resolution=128,
    num_views=500,
    debug=False,
    verbose=False,
    timings=None,
):
    """
    Rasterize a mesh from multiple views and remove invisible faces.
    Also includes postprocessing to:
        1. Remove connected components that are have low visibility.
        2. Mincut to remove faces at the inner side of the mesh connected to the outer side with a small hole.

    Args:
        verts (torch.Tensor): Vertices of the mesh. Shape (V, 3).
        faces (torch.Tensor): Faces of the mesh. Shape (F, 3).
        max_hole_size (float): Maximum area of a hole to fill.
        resolution (int): Resolution of the rasterization.
        num_views (int): Number of views to rasterize the mesh.
        verbose (bool): Whether to print progress.
    """
    started = time.perf_counter()
    # Construct cameras
    yaws = []
    pitchs = []
    for i in range(num_views):
        y, p = sphere_hammersley_sequence(i, num_views)
        yaws.append(y)
        pitchs.append(p)
    yaws = torch.tensor(yaws).cuda()
    pitchs = torch.tensor(pitchs).cuda()
    radius = 2.0
    fov = torch.deg2rad(torch.tensor(40)).cuda()
    projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3)
    views = []
    if BAKE_SYNC_FREE:
        origs = torch.stack([torch.sin(yaws) * torch.cos(pitchs), torch.cos(yaws) * torch.cos(pitchs), torch.sin(pitchs)], dim=-1).float() * radius
    else:
        origs = [torch.tensor([
            torch.sin(yaw) * torch.cos(pitch),
            torch.cos(yaw) * torch.cos(pitch),
            torch.sin(pitch),
        ]).cuda().float() * radius for (yaw, pitch) in zip(yaws, pitchs)]
    for orig in origs:
        view = utils3d.torch.view_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
        views.append(view)
    views = torch.stack(views, dim=0)

    # Rasterize
    visblity = torch.zeros(faces.shape[0], dtype=torch.int32, device=verts.device)
    rastctx = utils3d.torch.RastContext(backend='cuda')
    for i in tqdm(range(views.shape[0]), total=views.shape[0], disable=not verbose, desc='Rasterizing'):
        view = views[i]
        buffers = utils3d.torch.rasterize_triangle_faces(
            rastctx, verts[None], faces, resolution, resolution, view=view, projection=projection
        )
        if BAKE_SYNC_FREE:
            face_id = torch.where(buffers['mask'][0] > 0.95, buffers['face_id'][0], 0).long().reshape(-1)
            seen = torch.zeros(faces.shape[0] + 1, dtype=torch.bool, device=verts.device)
            seen[face_id] = True
            visblity += seen[1:]
            continue
        face_id = buffers['face_id'][0][buffers['mask'][0] > 0.95] - 1
        face_id = torch.unique(face_id).long()
        visblity[face_id] += 1
    visblity = visblity.float() / num_views
    started = mark_stage(timings, "holes_raster", started)
    
    # Mincut
    ## construct outer faces
    edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces)
    boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1)
    connected_components = utils3d.torch.compute_connected_components(faces, edges, face2edge)
    outer_face_indices = _visible_outer_faces(visblity, connected_components).nonzero().reshape(-1)
    
    ## construct inner faces
    inner_face_indices = torch.nonzero(visblity == 0).reshape(-1)
    if verbose:
        tqdm.write(f'Found {inner_face_indices.shape[0]} invisible faces')
    if inner_face_indices.shape[0] == 0:
        return verts, faces
    
    ## Construct dual graph (faces as nodes, edges as edges)
    dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(face2edge)
    dual_edge2edge = edges[dual_edge2edge]
    dual_edges_weights = torch.norm(verts[dual_edge2edge[:, 0]] - verts[dual_edge2edge[:, 1]], dim=1)
    if verbose:
        tqdm.write(f'Dual graph: {dual_edges.shape[0]} edges')

    ## solve mincut problem
    ### construct main graph
    g = igraph.Graph()
    g.add_vertices(faces.shape[0])
    g.add_edges(dual_edges.cpu().numpy())
    g.es['weight'] = dual_edges_weights.cpu().numpy()
    
    ### source and target
    g.add_vertex('s')
    g.add_vertex('t')
    
    ### connect invisible faces to source
    source, target = faces.shape[0], faces.shape[0] + 1
    inner = inner_face_indices.cpu().numpy()
    g.add_edges(np.stack([inner, np.full_like(inner, source)], axis=1), attributes={'weight': np.ones(len(inner), dtype=np.float32)})
    
    ### connect outer faces to target
    outer = outer_face_indices.cpu().numpy()
    g.add_edges(np.stack([outer, np.full_like(outer, target)], axis=1), attributes={'weight': np.ones(len(outer), dtype=np.float32)})
                
    ### solve mincut
    cut = g.mincut('s', 't', (np.array(g.es['weight']) * 1000).tolist())
    remove_face_indices = torch.tensor([v for v in cut.partition[0] if v < faces.shape[0]], dtype=torch.long, device=faces.device)
    started = mark_stage(timings, "holes_mincut", started)
    if verbose:
        tqdm.write(f'Mincut solved, start checking the cut')
    
    ### check if the cut is valid with each connected component
    to_remove_cc = utils3d.torch.compute_connected_components(faces[remove_face_indices])
    if debug:
        tqdm.write(f'Number of connected components of the cut: {len(to_remove_cc)}')
    valid_remove_cc = []
    cutting_edges = []
    for cc in to_remove_cc:
        #### check if the connected component has low visibility
        visblity_median = visblity[remove_face_indices[cc]].median()
        if debug:
            tqdm.write(f'visblity_median: {visblity_median}')
        if visblity_median > 0.25:
            continue
        
        #### check if the cuting loop is small enough
        cc_edge_indices, cc_edges_degree = torch.unique(face2edge[remove_face_indices[cc]], return_counts=True)
        cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1]
        cc_new_boundary_edge_indices = cc_boundary_edge_indices[~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)]
        if len(cc_new_boundary_edge_indices) > 0:
            cc_new_boundary_edge_cc = utils3d.torch.compute_edge_connected_components(edges[cc_new_boundary_edge_indices])
            cc_new_boundary_edges_cc_center = [verts[edges[cc_new_boundary_edge_indices[edge_cc]]].mean(dim=1).mean(dim=0) for edge_cc in cc_new_boundary_edge_cc]
            cc_new_boundary_edges_cc_area = []
            for i, edge_cc in enumerate(cc_new_boundary_edge_cc):
                _e1 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]] - cc_new_boundary_edges_cc_center[i]
                _e2 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]] - cc_new_boundary_edges_cc_center[i]
                cc_new_boundary_edges_cc_area.append(torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum() * 0.5)
            if debug:
                cutting_edges.append(cc_new_boundary_edge_indices)
                tqdm.write(f'Area of the cutting loop: {cc_new_boundary_edges_cc_area}')
            if any([l > max_hole_size for l in cc_new_boundary_edges_cc_area]):
                continue
            
        valid_remove_cc.append(cc)
        
    if debug:
        face_v = verts[faces].mean(dim=1).cpu().numpy()
        vis_dual_edges = dual_edges.cpu().numpy()
        vis_colors = np.zeros((faces.shape[0], 3), dtype=np.uint8)
        vis_colors[inner_face_indices.cpu().numpy()] = [0, 0, 255]
        vis_colors[outer_face_indices.cpu().numpy()] = [0, 255, 0]
        vis_colors[remove_face_indices.cpu().numpy()] = [255, 0, 255]
        if len(valid_remove_cc) > 0:
            vis_colors[remove_face_indices[torch.cat(valid_remove_cc)].cpu().numpy()] = [255, 0, 0]
        utils3d.io.write_ply('dbg_dual.ply', face_v, edges=vis_dual_edges, vertex_colors=vis_colors)
        
        vis_verts = verts.cpu().numpy()
        vis_edges = edges[torch.cat(cutting_edges)].cpu().numpy()
        utils3d.io.write_ply('dbg_cut.ply', vis_verts, edges=vis_edges)
        
    
    if len(valid_remove_cc) > 0:
        remove_face_indices = remove_face_indices[torch.cat(valid_remove_cc)]
        mask = torch.ones(faces.shape[0], dtype=torch.bool, device=faces.device)
        mask[remove_face_indices] = 0
        faces = faces[mask]
        faces, verts = utils3d.torch.remove_unreferenced_vertices(faces, verts)
        if verbose:
            tqdm.write(f'Removed {(~mask).sum()} faces by mincut')
    else:
        if verbose:
            tqdm.write(f'Removed 0 faces by mincut')
            
    if timings is not None:
        timings["faces_cut"] = int(faces.shape[0])
    started = mark_stage(timings, "holes_check", started)
    mesh = _meshfix.PyTMesh()
    mesh.load_array(verts.cpu().numpy(), faces.cpu().numpy())
    mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True)
    verts, faces = mesh.return_arrays()
    verts, faces = torch.tensor(verts, device='cuda', dtype=torch.float32), torch.tensor(faces, device='cuda', dtype=torch.int32)
    mark_stage(timings, "holes_meshfix", started)

    return verts, faces