"""MLX mesh renderer for Hunyuan3D texture pipeline.

Port of MeshRender's rendering and texture baking pipeline to MLX,
using mlx_arsenal.rasterize as the Metal rasterization backend.
"""

import mlx.core as mx
import numpy as np

from mlx_arsenal.rasterize import interpolate as mlx_interpolate
from mlx_arsenal.rasterize import rasterize_triangles

from .camera_utils_mlx import (
    get_mv_matrix,
    get_orthographic_projection_matrix,
    get_perspective_projection_matrix,
    transform_pos,
)


# ---------------------------------------------------------------------------
# Rasterizer adapter
# ---------------------------------------------------------------------------


class MLXRasterizer:
    """Adapter wrapping mlx_arsenal.rasterize to match MeshRender's interface."""

    @staticmethod
    def rasterize(pos, tri, resolution, clamp_depth=None, use_depth_prior=0):
        """Rasterize clip-space triangles.

        Args:
            pos: (1, N, 4) or (N, 4) mx.array clip-space vertices.
            tri: (F, 3) mx.array int32 face indices.
            resolution: (height, width) tuple.

        Returns:
            findices: (H, W) int32, 1-indexed face IDs.
            barycentric: (H, W, 3) float32.
        """
        vertices = pos[0] if pos.ndim == 3 else pos
        height, width = int(resolution[0]), int(resolution[1])

        depth_prior = None
        if clamp_depth is not None and use_depth_prior:
            depth_prior = clamp_depth

        findices, barycentric = rasterize_triangles(
            vertices.astype(mx.float32),
            tri.astype(mx.int32),
            width,
            height,
            depth_prior=depth_prior,
        )
        return findices, barycentric

    @staticmethod
    def interpolate(col, findices, barycentric, tri):
        """Interpolate per-vertex attributes at rasterized pixels.

        Args:
            col: (1, N, C) or (N, C) mx.array per-vertex data.
            findices: (H, W) int32.
            barycentric: (H, W, 3) float32.
            tri: (F, 3) int32.

        Returns:
            (1, H, W, C) mx.array.
        """
        attributes = col[0] if col.ndim == 3 else col
        result = mlx_interpolate(
            attributes.astype(mx.float32),
            findices,
            barycentric,
            tri.astype(mx.int32),
        )
        return mx.expand_dims(result, axis=0)


# ---------------------------------------------------------------------------
# Mesh renderer
# ---------------------------------------------------------------------------


class MeshRenderMLX:
    """MLX mesh renderer for normal / position / alpha rendering.

    Drop-in replacement for MeshRender's core render pipeline.
    All internal computation uses MLX arrays; public render methods
    return numpy arrays for downstream compatibility (PIL, cv2, etc.).
    """

    def __init__(
        self,
        camera_distance=1.45,
        camera_type="orth",
        default_resolution=1024,
        texture_size=1024,
        shader_type="face",
        scale_factor=1.15,
        bake_angle_thres=75,
        boundary_scale=2,
        bake_mode="back_sample",
        raster_mode="mlx",
    ):
        self.camera_distance = camera_distance
        self.shader_type = shader_type
        self._default_scale_factor = scale_factor
        self.bake_angle_thres = bake_angle_thres
        self.raster = MLXRasterizer()

        if isinstance(default_resolution, int):
            self.default_resolution = (default_resolution, default_resolution)
        else:
            self.default_resolution = tuple(default_resolution)

        if isinstance(texture_size, int):
            self.texture_size = (texture_size, texture_size)
        else:
            self.texture_size = tuple(texture_size)

        self.bake_unreliable_kernel_size = int(
            (boundary_scale / 512)
            * max(self.default_resolution[0], self.default_resolution[1])
        )

        if camera_type == "orth":
            self.set_orth_scale(1.2)
        elif camera_type == "perspective":
            self.camera_proj_mat = get_perspective_projection_matrix(
                49.13,
                self.default_resolution[1] / self.default_resolution[0],
                0.01,
                100.0,
            )
        else:
            raise ValueError(f"Unknown camera type: {camera_type}")

        # Geometry (set by set_mesh)
        self.vtx_pos = None
        self.pos_idx = None
        self.vtx_uv = None
        self.uv_idx = None
        self.scale_factor = 1.0
        self.mesh_normalize_scale_factor = 1.0
        self.mesh_normalize_scale_center = np.array([[0, 0, 0]])

        # Texture-space data (set by extract_textiles)
        self.tex_position = None
        self.tex_normal = None
        self.tex_grid = None
        self.texture_indices = None

    # ------------------------------------------------------------------
    # Projection
    # ------------------------------------------------------------------

    def set_orth_scale(self, ortho_scale):
        self.ortho_scale = ortho_scale
        self.camera_proj_mat = get_orthographic_projection_matrix(
            left=-ortho_scale * 0.5,
            right=ortho_scale * 0.5,
            bottom=-ortho_scale * 0.5,
            top=ortho_scale * 0.5,
            near=0.1,
            far=100,
        )

    # ------------------------------------------------------------------
    # Mesh loading
    # ------------------------------------------------------------------

    def set_mesh(self, vtx_pos, pos_idx, vtx_uv=None, uv_idx=None,
                 scale_factor=None, auto_center=True):
        """Load mesh geometry from numpy arrays.

        Applies the same coordinate transform and normalization
        as the original MeshRender.set_mesh().
        """
        if scale_factor is None:
            scale_factor = self._default_scale_factor

        vp = np.array(vtx_pos, dtype=np.float32)
        pi = np.array(pos_idx, dtype=np.int32)

        # Coordinate transform: flip X/Y, swap Y/Z
        vp[:, [0, 1]] = -vp[:, [0, 1]]
        vp[:, [1, 2]] = vp[:, [2, 1]].copy()

        if auto_center:
            max_bb = vp.max(axis=0)
            min_bb = vp.min(axis=0)
            center = (max_bb + min_bb) / 2.0
            scale = np.linalg.norm(vp - center, axis=1).max() * 2.0
            vp = (vp - center) * (scale_factor / scale)
            self.scale_factor = scale_factor
            self.mesh_normalize_scale_factor = scale_factor / scale
            self.mesh_normalize_scale_center = center[None, :]
        else:
            self.scale_factor = 1.0
            self.mesh_normalize_scale_factor = 1.0
            self.mesh_normalize_scale_center = np.array([[0, 0, 0]])

        self.vtx_pos = mx.array(vp)
        self.pos_idx = mx.array(pi)

        if vtx_uv is not None and uv_idx is not None:
            uv = np.array(vtx_uv, dtype=np.float32)
            uv[:, 1] = 1.0 - uv[:, 1]
            self.vtx_uv = mx.array(uv)
            self.uv_idx = mx.array(np.array(uv_idx, dtype=np.int32))
            self.extract_textiles()
        else:
            self.vtx_uv = None
            self.uv_idx = None

    def load_mesh(self, mesh, scale_factor=1.15, auto_center=True):
        """Load mesh from file path or trimesh object."""
        import trimesh

        if isinstance(mesh, str):
            mesh = trimesh.load(mesh, process=False, force="mesh")

        vtx_pos = np.array(mesh.vertices, dtype=np.float32)
        pos_idx = np.array(mesh.faces, dtype=np.int32)

        vtx_uv, uv_idx = None, None
        if hasattr(mesh.visual, "uv") and mesh.visual.uv is not None:
            vtx_uv = np.array(mesh.visual.uv, dtype=np.float32)
            uv_idx = pos_idx  # same topology

        self.set_mesh(
            vtx_pos, pos_idx, vtx_uv=vtx_uv, uv_idx=uv_idx,
            scale_factor=scale_factor, auto_center=auto_center,
        )

    # ------------------------------------------------------------------
    # Internal pipeline
    # ------------------------------------------------------------------

    def _create_view_state(self, elev, azim, camera_distance=None,
                           center=None, resolution=None):
        """Build clip-space vertices and camera-space positions."""
        if camera_distance is None:
            camera_distance = self.camera_distance
        if resolution is None:
            resolution = self.default_resolution

        mv = get_mv_matrix(elev, azim, camera_distance, center)
        pos_camera = transform_pos(mv, self.vtx_pos, keepdim=True)  # (N, 4)
        pos_clip = transform_pos(self.camera_proj_mat, pos_camera)   # (1, N, 4)
        return pos_camera, pos_clip, resolution

    def _rasterize(self, pos_clip, resolution):
        """Run rasterization, return packed rast_out (1, H, W, 4)."""
        findices, bary = self.raster.rasterize(
            pos_clip, self.pos_idx, resolution
        )
        # Pack [bary(3) | findices(1)] -> (1, H, W, 4)
        rast_out = mx.concatenate(
            [bary, mx.expand_dims(findices.astype(mx.float32), axis=-1)],
            axis=-1,
        )
        return mx.expand_dims(rast_out, axis=0)

    def _rasterize_tiled(self, pos_clip, resolution, tiles=4):
        """Tile rasterization across clip-space quadrants.

        At 4096² the single-dispatch raster exceeds Metal's command-buffer
        budget. Split clip space into ``tiles`` × ``tiles`` regions; for
        each region, scale vertex clip coords so the region maps to
        [-1, 1]² and rasterize at (H/tiles, W/tiles). Concatenate.

        Assumes ``tiles`` divides H and W evenly. Preserves the row-to-
        clip-y convention of the underlying rasterizer (verified by
        comparing tiled output to single-dispatch output at 2048²).
        """
        H, W = int(resolution[0]), int(resolution[1])
        if H % tiles != 0 or W % tiles != 0:
            return self._rasterize(pos_clip, resolution)
        tH, tW = H // tiles, W // tiles

        verts = pos_clip[0] if pos_clip.ndim == 3 else pos_clip
        vx, vy, vz, vw = verts[:, 0], verts[:, 1], verts[:, 2], verts[:, 3]

        rows_fi, rows_bary = [], []
        for i in range(tiles):
            cols_fi, cols_bary = [], []
            # Tile i covers rows [i*tH, (i+1)*tH). Row 0 = first tile output.
            # Matches the full-raster convention (tiled(0,0) == full[0:tH, 0:tW]).
            y_lo = -1.0 + 2.0 * i / tiles
            y_hi = -1.0 + 2.0 * (i + 1) / tiles
            mid_y = 0.5 * (y_lo + y_hi)
            span_y = 0.5 * (y_hi - y_lo)
            for j in range(tiles):
                x_lo = -1.0 + 2.0 * j / tiles
                x_hi = -1.0 + 2.0 * (j + 1) / tiles
                mid_x = 0.5 * (x_lo + x_hi)
                span_x = 0.5 * (x_hi - x_lo)
                # Homogeneous-safe clip transform: NDC' = (NDC - mid)/span
                # ⇒ clip' = (clip - w*mid)/span (w unchanged).
                new_verts = mx.stack([
                    (vx - vw * mid_x) / span_x,
                    (vy - vw * mid_y) / span_y,
                    vz,
                    vw,
                ], axis=-1)
                fi, bary = self.raster.rasterize(
                    new_verts, self.pos_idx, (tH, tW),
                )
                mx.synchronize()
                cols_fi.append(fi)
                cols_bary.append(bary)
            rows_fi.append(mx.concatenate(cols_fi, axis=1))
            rows_bary.append(mx.concatenate(cols_bary, axis=1))
        findices = mx.concatenate(rows_fi, axis=0)
        bary = mx.concatenate(rows_bary, axis=0)
        rast_out = mx.concatenate(
            [bary, mx.expand_dims(findices.astype(mx.float32), axis=-1)],
            axis=-1,
        )
        return mx.expand_dims(rast_out, axis=0)

    def _compute_face_normals(self, triangles):
        """Face normals from (F, 3, 3) triangle vertices."""
        e1 = triangles[:, 1, :] - triangles[:, 0, :]
        e2 = triangles[:, 2, :] - triangles[:, 0, :]
        # Cross product
        n = mx.stack([
            e1[:, 1] * e2[:, 2] - e1[:, 2] * e2[:, 1],
            e1[:, 2] * e2[:, 0] - e1[:, 0] * e2[:, 2],
            e1[:, 0] * e2[:, 1] - e1[:, 1] * e2[:, 0],
        ], axis=-1)
        norm = mx.sqrt((n * n).sum(axis=-1, keepdims=True) + 1e-12)
        return n / norm

    def _get_normals(self, pos_camera, pos_clip, resolution, use_abs_coor=False):
        """Compute per-pixel normals via face or vertex shading."""
        if use_abs_coor:
            mesh_tris = self.vtx_pos[self.pos_idx]  # (F, 3, 3)
        else:
            pc3 = pos_camera[:, :3] / pos_camera[:, 3:4]
            mesh_tris = pc3[self.pos_idx]

        face_normals = self._compute_face_normals(mesh_tris)
        rast_out = self._rasterize(pos_clip, resolution)

        if self.shader_type == "vertex":
            vertex_normals = self._vertex_normals_from_faces(face_normals)
            normal = self.raster.interpolate(
                mx.expand_dims(vertex_normals, 0),
                rast_out[0, ..., -1].astype(mx.int32),
                rast_out[0, ..., :-1],
                self.pos_idx,
            )
        else:
            # Face shader: direct face-normal lookup per pixel
            tri_ids = rast_out[..., 3].astype(mx.int32)  # (1, H, W)
            mask = (tri_ids > 0).astype(mx.float32)
            tri_ids_0 = mx.maximum(tri_ids - 1, 0)  # 0-indexed
            normal = face_normals[tri_ids_0[0]]  # (H, W, 3)
            normal = normal * mask[0, ..., None]
            normal = mx.expand_dims(normal, 0)  # (1, H, W, 3)

        return normal, rast_out

    def _vertex_normals_from_faces(self, face_normals):
        """Compute vertex normals as average of adjacent face normals."""
        num_verts = self.vtx_pos.shape[0]
        fi = np.array(self.pos_idx)
        fn = np.array(face_normals)
        vn = np.zeros((num_verts, 3), dtype=np.float32)
        counts = np.zeros((num_verts, 1), dtype=np.float32)
        for c in range(3):
            np.add.at(vn, fi[:, c], fn)
            np.add.at(counts, fi[:, c], 1.0)
        counts = np.maximum(counts, 1.0)
        vn = vn / counts
        norms = np.maximum(np.linalg.norm(vn, axis=-1, keepdims=True), 1e-12)
        return mx.array(vn / norms)

    # ------------------------------------------------------------------
    # Public render API
    # ------------------------------------------------------------------

    @staticmethod
    def _format_output(image_np, return_type="np"):
        """Convert numpy output to requested format."""
        if return_type == "pl":
            from PIL import Image as PILImage
            img = np.clip(image_np, 0, 1)
            return PILImage.fromarray((img * 255).astype(np.uint8))
        return image_np

    def render_normal(self, elev, azim, resolution=None, bg_color=(1, 1, 1),
                      use_abs_coor=False, camera_distance=None, center=None,
                      return_type="np"):
        """Render surface normals from a given viewpoint.

        Returns:
            (H, W, 3) numpy array (return_type="np") or PIL Image ("pl").
        """
        pos_camera, pos_clip, res = self._create_view_state(
            elev, azim, camera_distance, center, resolution,
        )
        normal, rast_out = self._get_normals(
            pos_camera, pos_clip, res, use_abs_coor,
        )
        visible = mx.clip(rast_out[..., -1:], 0, 1)
        bg = mx.array(list(bg_color), dtype=mx.float32).reshape(1, 1, 1, 3)
        result = normal * visible + bg * (1.0 - visible)
        result = (result + 1.0) * 0.5
        mx.synchronize()
        return self._format_output(np.array(result[0]), return_type)

    def render_position(self, elev, azim, resolution=None, bg_color=(1, 1, 1),
                        camera_distance=None, center=None, return_type="np"):
        """Render world-space positions from a given viewpoint.

        Returns:
            (H, W, 3) numpy array (return_type="np") or PIL Image ("pl").
        """
        _, pos_clip, res = self._create_view_state(
            elev, azim, camera_distance, center, resolution,
        )
        rast_out = self._rasterize(pos_clip, res)
        findices = rast_out[0, ..., -1].astype(mx.int32)
        bary = rast_out[0, ..., :-1]

        tex_pos = 0.5 - self.vtx_pos[:, :3] / self.scale_factor
        position = self.raster.interpolate(
            mx.expand_dims(tex_pos, 0), findices, bary, self.pos_idx,
        )

        visible = mx.clip(rast_out[..., -1:], 0, 1)
        bg = mx.array(list(bg_color), dtype=mx.float32).reshape(1, 1, 1, 3)
        result = position * visible + bg * (1.0 - visible)
        mx.synchronize()
        return self._format_output(np.array(result[0]), return_type)

    def render_alpha(self, elev, azim, resolution=None,
                     camera_distance=None, center=None, return_type="np"):
        """Render face-index map from a given viewpoint.

        Returns face indices (1-indexed, 0=background) matching the original
        MeshRender behavior. Shape: (1, H, W, 1) for return_type="np".
        """
        _, pos_clip, res = self._create_view_state(
            elev, azim, camera_distance, center, resolution,
        )
        rast_out = self._rasterize(pos_clip, res)
        # Return face indices (int), same as original: rast_out[..., -1:].long()
        face_ids = np.array(rast_out[..., -1:])[0].astype(np.int64)
        # Shape (1, H, W, 1) to match original pipeline expectations
        return face_ids[None, ...]

    # ------------------------------------------------------------------
    # Mesh & texture management
    # ------------------------------------------------------------------

    def set_default_render_resolution(self, default_resolution):
        if isinstance(default_resolution, int):
            default_resolution = (default_resolution, default_resolution)
        self.default_resolution = default_resolution

    def set_boundary_unreliable_scale(self, scale):
        self.bake_unreliable_kernel_size = int(
            (scale / 512)
            * max(self.default_resolution[0], self.default_resolution[1])
        )

    def get_face_areas(self, from_one_index=False):
        """Compute area of each triangle face.

        Returns:
            numpy array of face areas. If from_one_index, prepends a zero
            so areas[face_id] works with 1-indexed face IDs.
        """
        vp = np.array(self.vtx_pos)
        fi = np.array(self.pos_idx)
        v0 = vp[fi[:, 0]]
        v1 = vp[fi[:, 1]]
        v2 = vp[fi[:, 2]]
        cross = np.cross(v1 - v0, v2 - v0)
        areas = np.linalg.norm(cross, axis=-1) * 0.5
        if from_one_index:
            areas = np.insert(areas, 0, 0.0)
        return areas

    def set_texture(self, tex, force_set=False):
        """Store diffuse texture (numpy uint8 or float, or PIL Image)."""
        from PIL import Image as PILImage
        if isinstance(tex, PILImage.Image):
            tex = np.array(tex).astype(np.float32) / 255.0
        elif isinstance(tex, np.ndarray) and tex.dtype == np.uint8:
            tex = tex.astype(np.float32) / 255.0
        self.tex = tex

    def get_texture(self):
        """Return diffuse texture as numpy float32 array."""
        return self.tex if hasattr(self, "tex") and self.tex is not None else None

    def set_texture_mr(self, mr, force_set=False):
        """Store metallic-roughness texture."""
        from PIL import Image as PILImage
        if isinstance(mr, PILImage.Image):
            mr = np.array(mr).astype(np.float32) / 255.0
        elif isinstance(mr, np.ndarray) and mr.dtype == np.uint8:
            mr = mr.astype(np.float32) / 255.0
        self.tex_mr = mr

    def get_texture_mr(self):
        metallic, roughness = None, None
        if hasattr(self, "tex_mr") and self.tex_mr is not None:
            mr = self.tex_mr
            metallic = np.repeat(mr[:, :, 0:1], 3, axis=2)
            roughness = np.repeat(mr[:, :, 1:2], 3, axis=2)
        return metallic, roughness

    def get_texture_normal(self):
        if hasattr(self, "tex_normalMap") and self.tex_normalMap is not None:
            return self.tex_normalMap
        return None

    def get_mesh(self, normalize=True):
        """Return mesh geometry as numpy arrays.

        Applies inverse coordinate transform to restore original space.
        """
        vtx_pos = np.array(self.vtx_pos)
        pos_idx = np.array(self.pos_idx)

        if not normalize:
            vtx_pos = vtx_pos / self.mesh_normalize_scale_factor
            vtx_pos = vtx_pos + self.mesh_normalize_scale_center

        # Inverse of set_mesh transform: undo swap Y/Z, then undo flip X/Y
        vtx_pos[:, [1, 2]] = vtx_pos[:, [2, 1]].copy()
        vtx_pos[:, [0, 1]] = -vtx_pos[:, [0, 1]]

        vtx_uv, uv_idx = None, None
        if self.vtx_uv is not None:
            vtx_uv = np.array(self.vtx_uv)
            vtx_uv[:, 1] = 1.0 - vtx_uv[:, 1]
            uv_idx = np.array(self.uv_idx)

        return vtx_pos, pos_idx, vtx_uv, uv_idx

    def save_mesh(self, mesh_path, downsample=False):
        """Save textured mesh to OBJ file."""
        import cv2
        import trimesh

        vtx_pos, pos_idx, vtx_uv, uv_idx = self.get_mesh(normalize=False)
        texture = self.get_texture()

        if texture is not None and downsample:
            h, w = texture.shape[0] // 2, texture.shape[1] // 2
            texture = cv2.resize(texture, (w, h))

        mesh = trimesh.Trimesh(
            vertices=vtx_pos,
            faces=pos_idx,
            process=False,
        )
        if vtx_uv is not None and texture is not None:
            from PIL import Image as PILImage
            tex_img = PILImage.fromarray(
                (np.clip(texture, 0, 1) * 255).astype(np.uint8)
            )
            # Use PBRMaterial so GLB exports tag the baseColorTexture as
            # sRGB. SimpleMaterial doesn't carry color-space metadata and
            # GLB viewers assumed the texture was linear, applying an
            # unwanted gamma curve that darkened / desaturated the mesh
            # relative to the atlas PNG.
            mr_tex_img = None
            mr_atlas = getattr(self, "tex_mr", None)
            if mr_atlas is not None:
                # Repack for glTF: the diffusion MR atlas uses R=metallic,
                # G=roughness (PT convention, see MeshRender.get_texture_mr),
                # but glTF metallicRoughnessTexture expects B=metallic,
                # G=roughness, R ignored (occlusion in some pipelines).
                # Without this swap, glTF viewers read metallic=0 and the
                # mesh looks uniformly dielectric.
                mr_u8 = (np.clip(mr_atlas, 0, 1) * 255).astype(np.uint8)
                if mr_u8.ndim == 2:
                    mr_u8 = np.stack([mr_u8, mr_u8, mr_u8], axis=-1)
                metallic_ch = mr_u8[..., 0]
                roughness_ch = mr_u8[..., 1]
                mr_gltf = np.stack(
                    [np.zeros_like(metallic_ch), roughness_ch, metallic_ch],
                    axis=-1,
                )
                mr_tex_img = PILImage.fromarray(mr_gltf)

            kwargs = dict(
                name="paint_pbr",
                baseColorTexture=tex_img,
            )
            if mr_tex_img is not None:
                kwargs["metallicRoughnessTexture"] = mr_tex_img
                kwargs["metallicFactor"] = 1.0
                kwargs["roughnessFactor"] = 1.0
            else:
                kwargs["metallicFactor"] = 0.0
                kwargs["roughnessFactor"] = 1.0
            material = trimesh.visual.material.PBRMaterial(**kwargs)
            # doubleSided stops backface culling that was making the
            # silhouettes see-through.
            material.doubleSided = True
            visuals = trimesh.visual.TextureVisuals(
                uv=vtx_uv, material=material,
            )
            mesh.visual = visuals

        mesh.export(mesh_path)
        # Expose the in-memory trimesh so callers can export GLB directly.
        # Going OBJ -> trimesh.load -> GLB would drop the PBRMaterial
        # (OBJ/MTL can't carry metallicRoughnessTexture, doubleSided, or
        # sRGB hints) and silently erase our PBR work.
        self._last_exported_mesh = mesh

    # ------------------------------------------------------------------
    # Texture-space operations
    # ------------------------------------------------------------------

    def extract_textiles(self):
        """Rasterize mesh in UV space to build texture-space geometry maps.

        Populates self.tex_position, tex_normal, tex_grid, texture_indices.
        Required by back_project(method="back_sample").
        """
        if self.vtx_uv is None or self.uv_idx is None:
            return

        # Build clip-space coords from UVs: (u, v, 0, 1) * 2 - 1
        vnum = self.vtx_uv.shape[0]
        vtx_uv_clip = mx.concatenate([
            self.vtx_uv,
            mx.zeros((vnum, 1)),
            mx.ones((vnum, 1)),
        ], axis=1) * 2.0 - 1.0
        vtx_uv_clip = mx.expand_dims(vtx_uv_clip, 0)  # (1, V, 4)

        # Rasterize in UV space. At ≥4096² the single-dispatch raster
        # exceeds Metal's command-buffer budget, so tile it.
        th, tw = self.texture_size
        max_dim = max(th, tw)
        tiles = 1
        if max_dim >= 4096:
            tiles = 4
        elif max_dim >= 3072:
            tiles = 2
        if tiles > 1:
            rast_out = self._rasterize_tiled(vtx_uv_clip, self.texture_size, tiles=tiles)
        else:
            rast_out = self._rasterize(vtx_uv_clip, self.texture_size)
        fi = rast_out[0, ..., -1].astype(mx.int32)
        bary = rast_out[0, ..., :-1]

        # Interpolate world positions in UV space
        position = self.raster.interpolate(
            mx.expand_dims(self.vtx_pos, 0), fi, bary, self.pos_idx,
        )[0]  # (H, W, 3)

        # Compute face normals and per-pixel normals (face shader)
        face_normals = self._compute_face_normals(self.vtx_pos[self.pos_idx])
        tri_ids = rast_out[0, ..., 3].astype(mx.int32)
        mask = tri_ids > 0
        tri_ids_0 = mx.maximum(tri_ids - 1, 0)
        position_normal = face_normals[tri_ids_0]  # (H, W, 3)
        position_normal = position_normal * mask[..., None].astype(mx.float32)

        visible_mask = mx.clip(rast_out[0, ..., -1], 0, 1)  # (H, W)
        mx.synchronize()

        # Extract visible pixels to flat arrays (numpy for indexing)
        vis_np = np.array(visible_mask).reshape(-1)
        mask_flat = vis_np > 0

        pos_np = np.array(position).reshape(-1, 3)
        norm_np = np.array(position_normal).reshape(-1, 3)

        th, tw = self.texture_size
        row, col = np.meshgrid(np.arange(th), np.arange(tw), indexing="ij")
        grid_np = np.stack([row, col], axis=-1).reshape(-1, 2)

        pos_visible = pos_np[mask_flat]
        norm_visible = norm_np[mask_flat]
        grid_visible = grid_np[mask_flat]

        # Add homogeneous w=1
        pos_visible = np.concatenate(
            [pos_visible, np.ones((pos_visible.shape[0], 1), dtype=np.float32)],
            axis=1,
        )

        # Build reverse index: texture pixel -> flat visible index
        texture_indices = np.full(th * tw, -1, dtype=np.int64)
        flat_idx = grid_visible[:, 0] * tw + grid_visible[:, 1]
        texture_indices[flat_idx] = np.arange(len(grid_visible))

        self.tex_position = pos_visible     # (K, 4) float32
        self.tex_normal = norm_visible      # (K, 3) float32
        self.tex_grid = grid_visible        # (K, 2) int
        self.texture_indices = texture_indices.reshape(th, tw)  # (H, W)

        # Per-texel face id (1-indexed, 0 = empty), for per-face bake merging
        self.tex_face_id = np.array(tri_ids).astype(np.int32)  # (H, W)

    def uv_feature_map(self, vert_feat):
        """Map per-vertex features to UV texture space.

        Args:
            vert_feat: (N, C) mx.array of per-vertex features.

        Returns:
            (H, W, C) mx.array feature map.
        """
        vtx_uv_clip = mx.concatenate([
            self.vtx_uv * 2 - 1,
            mx.zeros((self.vtx_uv.shape[0], 1)),
            mx.ones((self.vtx_uv.shape[0], 1)),
        ], axis=1)
        vtx_uv_clip = mx.expand_dims(vtx_uv_clip, 0)

        rast_out = self._rasterize(vtx_uv_clip, self.texture_size)
        fi = rast_out[0, ..., -1].astype(mx.int32)
        bary = rast_out[0, ..., :-1]

        feat_map = self.raster.interpolate(
            mx.expand_dims(vert_feat, 0), fi, bary, self.uv_idx,
        )
        return feat_map[0]  # (H, W, C)

    # ------------------------------------------------------------------
    # Texture baking
    # ------------------------------------------------------------------

    @staticmethod
    def _render_sketch_from_depth(depth_image_np):
        """Edge detection on depth map (CPU, cv2)."""
        import cv2
        depth_u8 = (depth_image_np * 255).astype(np.uint8)
        edges = cv2.Canny(depth_u8, 30, 80)
        return (edges.astype(np.float32) / 255.0)[..., None]  # (H, W, 1)

    @staticmethod
    def _erode_boundary(mask_np, sketch_np, kernel_size):
        """Erode visible mask near silhouette boundaries (CPU, cv2)."""
        import cv2
        if kernel_size <= 0:
            return mask_np
        ks = kernel_size * 2 + 1
        kernel = np.ones((ks, ks), dtype=np.float32)

        # Erode visible mask: pixels near background become unreliable
        inv = 1.0 - mask_np[..., 0]
        eroded = cv2.filter2D(inv, -1, kernel)
        mask_np = mask_np * (eroded <= 0).astype(np.float32)[..., None]

        # Also exclude pixels near depth edges
        sketch_dilated = cv2.filter2D(sketch_np[..., 0], -1, kernel)
        mask_np = mask_np * (sketch_dilated < 0.5).astype(np.float32)[..., None]
        return mask_np

    def back_project(self, image, elev, azim, camera_distance=None,
                     center=None):
        """Back-project a rendered image onto UV texture space.

        Uses the "back_sample" method: projects texture-space positions
        into image space and bilinearly samples colors.

        Args:
            image: (H, W, C) numpy array or PIL Image, values in [0, 1].
            elev: Camera elevation in degrees.
            azim: Camera azimuth in degrees.

        Returns:
            texture: (H, W, C) numpy array in UV space.
            cos_map: (H, W, 1) numpy array cosine weights.
            boundary_map: (H, W, 1) numpy array boundary mask.
        """
        from PIL import Image as PILImage

        if isinstance(image, PILImage.Image):
            image = np.array(image).astype(np.float32) / 255.0
        elif not isinstance(image, np.ndarray):
            image = np.array(image)
        if image.ndim == 2:
            image = image[..., None]
        image = image.astype(np.float32)

        resolution = image.shape[:2]  # (H, W)
        channel = image.shape[2]
        th, tw = self.texture_size

        if camera_distance is None:
            camera_distance = self.camera_distance

        # --- Camera setup ---
        proj = self.camera_proj_mat
        mv = get_mv_matrix(elev, azim, camera_distance, center)
        pos_camera = transform_pos(mv, self.vtx_pos, keepdim=True)  # (N, 4)
        pos_clip = transform_pos(self.camera_proj_mat, pos_camera)   # (1, N, 4)

        pc3 = pos_camera[:, :3] / pos_camera[:, 3:4]

        # --- Face normals in camera space ---
        mesh_tris = pc3[self.pos_idx]  # (F, 3, 3)
        face_normals = self._compute_face_normals(mesh_tris)

        # --- Rasterize from camera view ---
        rast_out = self._rasterize(pos_clip, resolution)
        fi = rast_out[0, ..., -1].astype(mx.int32)
        bary = rast_out[0, ..., :-1]

        visible = mx.clip(rast_out[0, ..., -1:], 0, 1)  # (H, W, 1)

        # Per-pixel normals (face shader)
        tri_ids = fi
        tri_mask = (tri_ids > 0).astype(mx.float32)
        tri_ids_0 = mx.maximum(tri_ids - 1, 0)
        normal = face_normals[tri_ids_0] * tri_mask[..., None]  # (H, W, 3)

        # Depth
        depth_attr = pc3[:, 2:3].reshape(1, -1, 1)
        depth = self.raster.interpolate(depth_attr, fi, bary, self.pos_idx)
        depth = depth[0]  # (H, W, 1)
        mx.synchronize()

        visible_np = np.array(visible)
        normal_np = np.array(normal)
        depth_np = np.array(depth)

        # Depth normalization for sketch
        vis_flat = visible_np.reshape(-1) > 0
        if vis_flat.sum() == 0:
            return (np.zeros((*self.texture_size, channel), dtype=np.float32),
                    np.zeros((*self.texture_size, 1), dtype=np.float32),
                    np.zeros((*self.texture_size, 1), dtype=np.float32))
        d_vis = depth_np.reshape(-1)[vis_flat]
        d_max, d_min = d_vis.max(), d_vis.min()
        depth_norm = (depth_np - d_min) / max(d_max - d_min, 1e-8)
        depth_image = depth_norm * visible_np

        sketch_np = self._render_sketch_from_depth(depth_image[..., 0])

        # Cosine weighting (view angle vs surface normal)
        lookat = np.array([0, 0, -1], dtype=np.float32)
        normal_flat = normal_np.reshape(-1, 3)
        cos_flat = (normal_flat * lookat).sum(axis=-1)
        cos_np = cos_flat.reshape(resolution[0], resolution[1], 1)
        cos_thres = np.cos(np.radians(self.bake_angle_thres))
        cos_np[cos_np < cos_thres] = 0.0

        # Boundary erosion
        visible_np = self._erode_boundary(
            visible_np, sketch_np, self.bake_unreliable_kernel_size,
        )
        cos_np[visible_np[..., 0] == 0] = 0.0

        # --- Back-sample: project tex_position to image space ---
        if self.tex_position is None:
            raise RuntimeError("extract_textiles() must be called first")

        img_proj = np.array([
            [proj[0, 0], 0, 0, 0],
            [0, proj[1, 1], 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ], dtype=np.float32)

        mv_f = mv.astype(np.float32)
        tex_pos = self.tex_position  # (K, 4)
        v_proj = tex_pos @ mv_f.T @ img_proj  # (K, 4)

        # Check which tex pixels project inside the image
        inner = (
            (v_proj[:, 0] >= -1.0) & (v_proj[:, 0] <= 1.0)
            & (v_proj[:, 1] >= -1.0) & (v_proj[:, 1] <= 1.0)
        )
        inner_idx = np.where(inner)[0]

        img_x = np.clip(
            ((np.clip(v_proj[:, 0], -1, 1) * 0.5 + 0.5) * resolution[0]).astype(np.int64),
            0, resolution[0] - 1,
        )
        img_y = np.clip(
            ((np.clip(v_proj[:, 1], -1, 1) * 0.5 + 0.5) * resolution[1]).astype(np.int64),
            0, resolution[1] - 1,
        )

        indices = img_y * resolution[0] + img_x
        depth_flat = depth_np.reshape(-1)
        vis_mask_flat = visible_np.reshape(-1)
        cos_flat = cos_np.reshape(-1)

        sampled_z = depth_flat[indices]
        sampled_m = vis_mask_flat[indices]
        sampled_w = cos_flat[indices]
        v_z = v_proj[:, 2]

        depth_thres = 3e-3
        valid_mask = (np.abs(v_z - sampled_z) < depth_thres) & (sampled_m * sampled_w > 0)
        valid_idx = np.where(valid_mask)[0]
        valid_idx = np.intersect1d(valid_idx, inner_idx)

        # Bilinear sampling from image
        wx = ((v_proj[:, 0] * 0.5 + 0.5) * resolution[0] - img_x)[valid_idx].reshape(-1, 1)
        wy = ((v_proj[:, 1] * 0.5 + 0.5) * resolution[1] - img_y)[valid_idx].reshape(-1, 1)
        ix = img_x[valid_idx]
        iy = img_y[valid_idx]
        ix_r = np.clip(ix + 1, 0, resolution[0] - 1)
        iy_r = np.clip(iy + 1, 0, resolution[1] - 1)

        rgb = image.reshape(-1, channel)
        i00 = iy * resolution[0] + ix
        i10 = iy * resolution[0] + ix_r
        i01 = iy_r * resolution[0] + ix
        i11 = iy_r * resolution[0] + ix_r
        sampled_rgb = (
            (rgb[i00] * (1 - wx) + rgb[i10] * wx) * (1 - wy)
            + (rgb[i01] * (1 - wx) + rgb[i11] * wx) * wy
        )

        sampled_b = sketch_np.reshape(-1)[indices[valid_idx]]
        sampled_w_valid = sampled_w[valid_idx]

        # Write to texture space
        texture = np.zeros((th * tw, channel), dtype=np.float32)
        cos_map = np.zeros(th * tw, dtype=np.float32)
        boundary_map = np.zeros(th * tw, dtype=np.float32)

        tex_grid = self.tex_grid[valid_idx]
        tex_flat_idx = tex_grid[:, 0] * tw + tex_grid[:, 1]
        texture[tex_flat_idx] = sampled_rgb
        cos_map[tex_flat_idx] = sampled_w_valid
        boundary_map[tex_flat_idx] = sampled_b

        return (
            texture.reshape(th, tw, channel),
            cos_map.reshape(th, tw, 1),
            boundary_map.reshape(th, tw, 1),
        )

    def fast_bake_texture(self, textures, cos_maps, mode="weighted"):
        """Merge multiple view textures into a single UV atlas.

        Three modes (PT parity is default):
          - ``"weighted"`` (default, PyTorch parity): cosine-weighted
            average of all views. Smooth across face boundaries; matches
            MeshRender.fast_bake_texture exactly.
          - ``"face_wta"``: pick one view per UV face. Sharper but creates
            visible discontinuities at mesh-face seams where adjacent
            faces pick different winning views.
          - ``"wta"``: per-texel argmax. Fastest, stripiest.

        Args:
            textures: list of (H, W, C) numpy arrays.
            cos_maps: list of (H, W, 1) numpy arrays (already weighted).
            mode: "wta" or "weighted".

        Returns:
            texture_merge: (H, W, C) numpy array.
            trust_map: (H, W, 1) boolean numpy array.
        """
        if not textures:
            raise ValueError("fast_bake_texture: no textures provided")

        if mode == "face_wta":
            # Per-face WTA: aggregate cos per (face, view), pick best view per
            # face, then bake all of that face's texels from that view.
            if getattr(self, "tex_face_id", None) is None:
                # Fall back to per-texel WTA if face ids aren't available
                mode = "wta"
            else:
                stack_t = np.stack(textures, axis=0)              # (V, H, W, C)
                stack_c = np.stack([c[..., 0] for c in cos_maps], axis=0)  # (V, H, W)
                face_ids = self.tex_face_id                        # (H, W) 1-indexed
                V = stack_c.shape[0]
                n_faces = int(face_ids.max())
                # Sum cos per (view, face)
                flat_face = face_ids.reshape(-1)
                valid = flat_face > 0
                face_idx = (flat_face[valid] - 1).astype(np.int64)
                stack_c_flat = stack_c.reshape(V, -1)[:, valid]    # (V, K)
                cos_per_face = np.zeros((V, n_faces), dtype=np.float32)
                for v in range(V):
                    np.add.at(cos_per_face[v], face_idx, stack_c_flat[v])
                best_view_per_face = cos_per_face.argmax(axis=0)   # (n_faces,)
                # Per-texel best view from face lookup
                best_per_texel = np.zeros_like(face_ids, dtype=np.int64)
                best_per_texel[face_ids > 0] = best_view_per_face[face_ids[face_ids > 0] - 1]
                # Painted = any view contributed at this texel
                best_c = stack_c.max(axis=0)
                H, W = face_ids.shape
                ii, jj = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
                merged = stack_t[best_per_texel, ii, jj]
                merged[best_c == 0] = 0
                return merged, (best_c > 0)[..., None]

        if mode == "wta":
            stack_t = np.stack(textures, axis=0)
            stack_c = np.stack([c[..., 0] for c in cos_maps], axis=0)
            best = stack_c.argmax(axis=0)
            best_c = stack_c.max(axis=0)
            H, W = best.shape
            ii, jj = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
            merged = stack_t[best, ii, jj]
            merged[best_c == 0] = 0
            return merged, (best_c > 0)[..., None]

        # PyTorch-parity weighted average path
        channel = textures[0].shape[-1]
        th, tw = self.texture_size
        texture_merge = np.zeros((th, tw, channel), dtype=np.float32)
        trust_map = np.zeros((th, tw, 1), dtype=np.float32)
        for texture, cos_map in zip(textures, cos_maps):
            view_sum = (cos_map > 0).sum()
            if view_sum == 0:
                continue
            painted_sum = ((cos_map > 0) & (trust_map > 0)).sum()
            if painted_sum / view_sum > 0.99:
                continue
            texture_merge += texture * cos_map
            trust_map += cos_map
        texture_merge = texture_merge / np.maximum(trust_map, 1e-8)
        return texture_merge, trust_map > 1e-8

    def bake_texture(self, colors, elevs, azims, camera_distance=None,
                     center=None, exp=6, weights=None):
        """Bake multiple views into a single UV texture.

        Args:
            colors: list of images (numpy or PIL).
            elevs: list of elevation angles.
            azims: list of azimuth angles.
            exp: exponent for cosine weighting.
            weights: optional per-view weights.

        Returns:
            texture_merge: (H, W, C) numpy array.
            trust_map: (H, W, 1) boolean numpy array.
        """
        from PIL import Image as PILImage

        if weights is None:
            weights = [1.0] * len(colors)

        textures = []
        cos_maps = []
        for color, elev, azim, weight in zip(colors, elevs, azims, weights):
            if isinstance(color, PILImage.Image):
                color = np.array(color).astype(np.float32) / 255.0
            texture, cos_map, _ = self.back_project(
                color, elev, azim, camera_distance, center,
            )
            cos_map = weight * (cos_map ** exp)
            textures.append(texture)
            cos_maps.append(cos_map)

        return self.fast_bake_texture(textures, cos_maps)

    def uv_inpaint(self, texture, mask, method="NS", vertex_inpaint=True):
        """Inpaint missing regions in UV texture.

        Two-pass pipeline matching PyTorch's bake flow:
          1. Mesh-aware vertex-color propagation + face-bary rasterization
             (fills most of the UV islands using 3D adjacency).
          2. Gutter fill for the remaining empty texels.

        For step 2 the default ``"edt"`` mode copies each gutter texel from
        its nearest painted texel via ``scipy.ndimage.distance_transform_edt``.
        It's deterministic and never blends colors across distant UV islands —
        unlike ``"NS"`` (cv2.INPAINT_NS) which diffuses and tends to mix
        unrelated islands together when they happen to be 2D-close in atlas
        space.

        Args:
            texture: (H, W, C) numpy array in [0, 1].
            mask: (H, W) or (H, W, 1) numpy uint8, 255=keep, 0=inpaint.
            method: "edt" (default, EDT nearest-fill) or "NS" (cv2 NS).
            vertex_inpaint: Run the mesh-aware propagation pass first.

        Returns:
            (H, W, C) numpy uint8 array.
        """
        if isinstance(texture, mx.array):
            texture = np.array(texture)

        if mask.ndim == 3:
            mask = mask[..., 0]
        if mask.dtype != np.uint8:
            mask = (mask * 255).astype(np.uint8)

        tex_f = np.clip(texture, 0, 1).astype(np.float32)

        if vertex_inpaint and self.vtx_uv is not None and self.uv_idx is not None:
            from .mesh_inpaint_py import mesh_vertex_inpaint
            tex_f, mask = mesh_vertex_inpaint(
                tex_f, mask,
                np.asarray(self.vtx_pos), np.asarray(self.vtx_uv),
                np.asarray(self.pos_idx), np.asarray(self.uv_idx),
            )

        if method == "edt":
            from scipy.ndimage import distance_transform_edt
            painted = mask > 0
            if painted.any():
                idx = distance_transform_edt(
                    ~painted, return_distances=False, return_indices=True,
                )
                tex_f = tex_f[idx[0], idx[1]]
            return (np.clip(tex_f, 0, 1) * 255).astype(np.uint8)

        # cv2.INPAINT_NS for near-field fill (PT parity), then an explicit
        # EDT edge-padding pass for UV gutter. NS's 3-px radius is smaller
        # than typical UV-island spacing in a 2048^2 atlas, so 3D viewers
        # doing bilinear texture sampling across island boundaries read
        # the NS-smoothed (near-black) gutter and display thin wireframe
        # lines on the mesh. EDT pads every unpainted texel with the
        # nearest painted color — this is the standard post-bake step in
        # production UV pipelines and never affects texels that are
        # visible in 3D (painted).
        import cv2
        from scipy.ndimage import distance_transform_edt
        texture_u8 = (np.clip(tex_f, 0, 1) * 255).astype(np.uint8)
        inpainted = cv2.inpaint(texture_u8, 255 - mask, 3, cv2.INPAINT_NS)
        painted = mask > 0
        if painted.any() and (~painted).any():
            idx = distance_transform_edt(
                ~painted, return_distances=False, return_indices=True,
            )
            inpainted = inpainted[idx[0], idx[1]]
        return inpainted
