fluencyCAD/src/fluency/rendering/pygfx_renderer.py

"""
pygfx-based renderer for Fluency CAD.

This module provides a modern WebGPU-based renderer using pygfx,
offering a smaller dependency footprint than VTK while providing
excellent 3D visualization capabilities.
"""

from typing import List, Tuple, Optional, Callable, Any
import numpy as np
from dataclasses import dataclass

from fluency.rendering.base import (
    Renderer,
    RenderObject,
    RenderColor,
)


def compute_normals(vertices: np.ndarray, faces: np.ndarray) -> np.ndarray:
    """Compute vertex normals from positions and face indices."""
    vertices = np.asarray(vertices, dtype=np.float32)
    faces = np.asarray(faces, dtype=np.int32)

    normals = np.zeros_like(vertices)

    for face in faces:
        v0, v1, v2 = vertices[face[0]], vertices[face[1]], vertices[face[2]]

        edge1 = v1 - v0
        edge2 = v2 - v0

        face_normal = np.cross(edge1, edge2)

        length = np.linalg.norm(face_normal)
        if length > 1e-10:
            face_normal = face_normal / length

        normals[face[0]] += face_normal
        normals[face[1]] += face_normal
        normals[face[2]] += face_normal

    lengths = np.linalg.norm(normals, axis=1, keepdims=True)
    lengths = np.where(lengths > 1e-10, lengths, 1)
    normals = normals / lengths

    return normals.astype(np.float32)


@dataclass
class PygfxRenderObject(RenderObject):
    """pygfx render object wrapper."""

    scene_node: Any = None
    geometry: Any = None
    material: Any = None

    def __init__(
        self,
        scene_node: Any = None,
        geometry: Any = None,
        material: Any = None,
        name: Optional[str] = None,
    ):
        super().__init__(name=name)
        self.scene_node = scene_node
        self.geometry = geometry
        self.material = material


class PygfxRenderer(Renderer):
    """
    pygfx-based renderer implementation.

    This renderer uses pygfx (WebGPU-based) for 3D visualization,
    providing modern rendering with a small dependency footprint.
    """

    def __init__(self) -> None:
        self._canvas: Any = None
        self._renderer: Any = None
        self._scene: Any = None
        self._camera: Any = None
        self._controller: Any = None
        self._objects: List[PygfxRenderObject] = []
        self._pick_callback: Optional[Callable[[Any], None]] = None
        self._camera_change_callback: Optional[Callable[[], None]] = None
        self._background_color: Tuple[float, float, float] = (0.1, 0.1, 0.15)
        self._initialized: bool = False

    def initialize(self, parent_widget: Any) -> bool:
        """Initialize pygfx with Qt widget."""
        try:
            import pygfx as gfx
            from rendercanvas.qt import RenderWidget
            from PySide6.QtWidgets import QVBoxLayout

            self._canvas = RenderWidget(parent=parent_widget)
            self._renderer = gfx.renderers.WgpuRenderer(self._canvas)
            self._scene = gfx.Scene()

            self._camera = gfx.PerspectiveCamera(50, 16 / 9)
            self._camera.local.position = (100, 100, 100)

            self._controller = gfx.OrbitController(self._camera)
            self._controller.register_events(self._renderer)

            self._setup_lighting()
            self._add_grid()

            layout = QVBoxLayout(parent_widget)
            layout.setContentsMargins(0, 0, 0, 0)
            layout.addWidget(self._canvas)

            self._canvas.request_draw(self._animate)

            self._setup_picking()
            self._initialized = True

            return True

        except Exception as e:
            print(f"Failed to initialize pygfx: {e}")
            import traceback

            traceback.print_exc()
            return False

    def _animate(self):
        """Animation callback for the canvas."""
        if self._initialized:
            self._renderer.render(self._scene, self._camera)

    def _setup_lighting(self) -> None:
        """Setup scene lighting."""
        import pygfx as gfx

        ambient = gfx.AmbientLight(intensity=0.3)
        self._scene.add(ambient)

        directional = gfx.DirectionalLight(intensity=1.0)
        directional.local.position = (100, 100, 100)
        self._scene.add(directional)

        fill = gfx.DirectionalLight(intensity=0.5)
        fill.local.position = (-100, 50, 50)
        self._scene.add(fill)

    def _add_grid(self) -> None:
        """Add reference grid."""
        import pygfx as gfx

        grid = gfx.GridHelper(
            size=200, divisions=20, color1=(0.3, 0.3, 0.3, 1), color2=(0.2, 0.2, 0.2, 1)
        )
        self._scene.add(grid)

    def _setup_picking(self) -> None:
        """Setup mesh picking."""
        pass

    def shutdown(self) -> None:
        """Clean up renderer resources."""
        self._objects.clear()
        self._initialized = False

    def add_mesh(
        self,
        vertices: np.ndarray,
        faces: np.ndarray,
        color: Tuple[float, float, float] = (0.2, 0.4, 0.8),
        name: Optional[str] = None,
    ) -> str:
        """Add a mesh to the scene. Returns the mesh ID."""
        import pygfx as gfx
        import uuid

        vertices = np.asarray(vertices, dtype=np.float32)
        faces = np.asarray(faces, dtype=np.int32)

        normals = compute_normals(vertices, faces)

        geometry = gfx.Geometry(positions=vertices, indices=faces, normals=normals)

        material = gfx.MeshPhongMaterial(color=color, flat_shading=False)

        mesh = gfx.Mesh(geometry, material)
        self._scene.add(mesh)

        mesh_id = name or f"mesh_{uuid.uuid4().hex[:8]}"
        obj = PygfxRenderObject(name=mesh_id, scene_node=mesh, geometry=geometry, material=material)
        obj.color = RenderColor(*color)
        self._objects.append(obj)

        return mesh_id

    def add_wireframe(
        self,
        vertices: np.ndarray,
        edges: np.ndarray,
        color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
        line_width: float = 1.0,
        name: Optional[str] = None,
    ) -> str:
        """Add a wireframe to the scene. Returns the wireframe ID."""
        import pygfx as gfx
        import uuid

        positions: List[List[float]] = []
        for edge in edges:
            positions.append(vertices[edge[0]])
            positions.append(vertices[edge[1]])

        positions_arr = np.array(positions, dtype=np.float32)

        geometry = gfx.Geometry(positions=positions_arr)
        material = gfx.LineMaterial(color=color, thickness=line_width)

        lines = gfx.Line(geometry, material)
        self._scene.add(lines)

        wireframe_id = name or f"wireframe_{uuid.uuid4().hex[:8]}"
        obj = PygfxRenderObject(
            name=wireframe_id, scene_node=lines, geometry=geometry, material=material
        )
        obj.color = RenderColor(*color)
        self._objects.append(obj)

        return wireframe_id

    def add_points(
        self,
        points: np.ndarray,
        color: Tuple[float, float, float] = (1.0, 0.0, 0.0),
        size: float = 5.0,
        name: Optional[str] = None,
    ) -> str:
        """Add points to the scene. Returns the points ID."""
        import pygfx as gfx
        import uuid

        points_arr = np.asarray(points, dtype=np.float32)

        geometry = gfx.Geometry(positions=points_arr)
        material = gfx.PointsMaterial(color=color, size=size)

        points_obj = gfx.Points(geometry, material)
        self._scene.add(points_obj)

        points_id = name or f"points_{uuid.uuid4().hex[:8]}"
        obj = PygfxRenderObject(
            name=points_id, scene_node=points_obj, geometry=geometry, material=material
        )
        obj.color = RenderColor(*color)
        self._objects.append(obj)

        return points_id

    def add_lines(
        self,
        start_points: np.ndarray,
        end_points: np.ndarray,
        color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
        line_width: float = 1.0,
        name: Optional[str] = None,
    ) -> str:
        """Add line segments to the scene. Returns the lines ID."""
        import pygfx as gfx
        import uuid

        positions = np.hstack([start_points, end_points]).flatten()
        positions = positions.reshape(-1, 3).astype(np.float32)

        geometry = gfx.Geometry(positions=positions)
        material = gfx.LineMaterial(color=color, thickness=line_width)

        lines = gfx.Line(geometry, material)
        self._scene.add(lines)

        lines_id = name or f"lines_{uuid.uuid4().hex[:8]}"
        obj = PygfxRenderObject(
            name=lines_id, scene_node=lines, geometry=geometry, material=material
        )
        obj.color = RenderColor(*color)
        self._objects.append(obj)

        return lines_id

    def remove_object(self, obj: RenderObject) -> bool:
        """Remove an object from the scene."""
        if isinstance(obj, PygfxRenderObject) and obj in self._objects:
            if obj.scene_node is not None:
                self._scene.remove(obj.scene_node)
            self._objects.remove(obj)
            return True
        return False

    def remove_mesh(self, mesh_id: str) -> bool:
        """Remove a mesh by its ID."""
        for obj in self._objects:
            if obj.name == mesh_id:
                return self.remove_object(obj)
        return False

    def clear_scene(self) -> None:
        """Remove all objects from the scene."""
        for obj in self._objects[:]:
            if obj.scene_node is not None:
                self._scene.remove(obj.scene_node)
        self._objects.clear()

    def update_mesh(self, obj: RenderObject, vertices: np.ndarray, faces: np.ndarray) -> bool:
        """Update mesh geometry."""
        if not isinstance(obj, PygfxRenderObject):
            return False

        import pygfx as gfx

        vertices = np.asarray(vertices, dtype=np.float32)
        faces = np.asarray(faces, dtype=np.int32)

        normals = compute_normals(vertices, faces)

        geometry = gfx.Geometry(positions=vertices, indices=faces, normals=normals)

        obj.geometry = geometry
        if obj.scene_node is not None:
            obj.scene_node.geometry = geometry

        return True

    def set_object_color(self, obj: RenderObject, color: Tuple[float, float, float]) -> None:
        """Set the color of an object."""
        if isinstance(obj, PygfxRenderObject):
            obj.color = RenderColor(*color)
            if obj.material is not None:
                obj.material.color = color

    def set_object_visible(self, obj: RenderObject, visible: bool) -> None:
        """Set the visibility of an object."""
        if isinstance(obj, PygfxRenderObject):
            obj.visible = visible
            if obj.scene_node is not None:
                obj.scene_node.visible = visible

    def set_camera_position(
        self,
        position: Tuple[float, float, float],
        target: Tuple[float, float, float] = (0, 0, 0),
        up: Tuple[float, float, float] = (0, 0, 1),
    ) -> None:
        """Set camera position and orientation."""
        self._camera.local.position = position
        self._camera.look_at(target)

    def get_camera_position(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Get camera position, target, and up vector."""
        pos = np.array(self._camera.local.position)
        target = np.array([0, 0, 0])
        up = np.array(self._camera.local.up)
        return pos, target, up

    def fit_camera(self, padding: float = 1.1) -> None:
        """Fit camera to show all objects."""
        if not self._objects:
            return

        all_positions: List[np.ndarray] = []
        for obj in self._objects:
            if obj.geometry is not None and hasattr(obj.geometry, "positions"):
                positions = obj.geometry.positions.data
                all_positions.append(positions)

        if all_positions:
            positions = np.vstack(all_positions)
            min_pos = positions.min(axis=0)
            max_pos = positions.max(axis=0)
            center = (min_pos + max_pos) / 2
            size = np.linalg.norm(max_pos - min_pos) * padding

            self._camera.local.position = (center[0] + size, center[1] + size, center[2] + size)
            self._camera.look_at(tuple(center))

    def set_camera_perspective(
        self, fov: float = 50.0, near: float = 0.1, far: float = 10000.0
    ) -> None:
        """Set camera perspective parameters."""
        self._camera.fov = fov
        self._camera.near = near
        self._camera.far = far

    def set_camera_orthographic(
        self, width: float = 100.0, near: float = 0.1, far: float = 10000.0
    ) -> None:
        """Set camera orthographic parameters."""
        import pygfx as gfx

        self._camera = gfx.OrthographicCamera(width=width, near=near, far=far)
        self._controller.camera = self._camera

    def render(self) -> None:
        """Trigger a render."""
        if self._initialized:
            self._canvas.request_draw()

    def on_pick(self, callback: Callable[[Any], None]) -> None:
        """Register a callback for picking/selection."""
        self._pick_callback = callback

    def on_camera_change(self, callback: Callable[[], None]) -> None:
        """Register a callback for camera changes."""
        self._camera_change_callback = callback

    def set_background_color(self, color: Tuple[float, float, float]) -> None:
        """Set the background color."""
        self._background_color = color
        self._scene.background = color

    def add_grid(
        self,
        size: float = 100.0,
        divisions: int = 10,
        color: Tuple[float, float, float] = (0.3, 0.3, 0.3),
    ) -> str:
        """Add a reference grid. Returns the grid ID."""
        import pygfx as gfx

        grid = gfx.GridHelper(
            size=size, divisions=divisions, color1=(*color, 1), color2=(*color, 0.5)
        )
        self._scene.add(grid)

        obj = PygfxRenderObject(name="grid", scene_node=grid)
        self._objects.append(obj)
        return "grid"

    def add_axes(self, size: float = 10.0, visible: bool = True) -> str:
        """Add coordinate axes. Returns the axes ID."""
        import pygfx as gfx

        axes = gfx.AxesHelper(size=size)
        axes.visible = visible
        self._scene.add(axes)

        obj = PygfxRenderObject(name="axes", scene_node=axes)
        self._objects.append(obj)
        return "axes"

    def get_screen_size(self) -> Tuple[int, int]:
        """Get the screen size in pixels."""
        if self._canvas is not None:
            return self._canvas.get_physical_size()
        return (800, 600)

    def project_to_screen(self, point: Tuple[float, float, float]) -> Tuple[int, int]:
        """Project a 3D point to screen coordinates."""
        return (0, 0)

    def unproject_from_screen(
        self, screen_x: int, screen_y: int, depth: float = 0.0
    ) -> Tuple[float, float, float]:
        """Unproject screen coordinates to 3D."""
        return (0.0, 0.0, 0.0)

    def take_screenshot(self) -> np.ndarray:
        """Take a screenshot of the current view."""
        return np.zeros((100, 100, 4), dtype=np.uint8)

    def save_screenshot(self, filepath: str) -> bool:
        """Save a screenshot to file."""
        try:
            img = self.take_screenshot()
            from PIL import Image

            Image.fromarray(img).save(filepath)
            return True
        except Exception as e:
            print(f"Screenshot error: {e}")
            return False