components.py

"""
Optical rendering panel backed by Mitsuba.
"""
from typing import Dict, Any
import numpy as np
from witwin_server import Panel, panel
from witwin_server.core.components import button, bool_field, int_field, image_field
from witwin_server.utils.logging import get_logger
from witwin_server.utils.mitsuba_utils import (
    MITSUBA_AVAILABLE,
    scene_to_mitsuba,
)

logger = get_logger("Optical Renderer")

if MITSUBA_AVAILABLE:
    import mitsuba as mi


@panel(
    id="optical_renderer",
    display_name="Optical Renderer",
    icon="image",
    menu_path="Window/Optical Renderer"
)
class MitsubaPanel(Panel):
    """Optical rendering panel backed by Mitsuba."""

    # Render settings
    use_viewport_resolution = bool_field(True, title="Render Settings", description="Use viewport resolution")
    width = int_field(512, min=64, max=4096, hide_if="use_viewport_resolution")
    height = int_field(512, min=64, max=4096, hide_if="use_viewport_resolution")
    spp = int_field(64, min=1, max=1024, description="Samples per pixel")
    transparent_background = bool_field(True, description="Transparent background")

    # Preview image
    preview = image_field(title="Render Preview", hide_label=True)

    def _get_value(self, val):
        """Helper to convert tensor/array to Python value."""
        if hasattr(val, 'item'):
            return val.item()
        if hasattr(val, 'tolist'):
            return val.tolist()
        return val

    @button(display_name="Render", description="Render scene using Mitsuba")
    def render(self, camera: Dict[str, Any] = None):
        """Render the scene using Mitsuba.

        Args:
            camera: Optional camera data from frontend viewport:
                - world_matrix: 4x4 camera world matrix (16 floats, column-major)
                - fov: Field of view in degrees
                - viewport_width: Viewport width in pixels
                - viewport_height: Viewport height in pixels
        """
        if not MITSUBA_AVAILABLE:
            return "Mitsuba not available - install with: pip install mitsuba drjit"

        # Get scene from server
        if not self._server or not hasattr(self._server, 'scene'):
            return "Server not available"

        app_scene = self._server.scene
        if not app_scene:
            return "No scene available"

        logger.info("Building scene...")

        try:
            # Get render settings
            camera_data = camera or {}
            use_viewport = self._get_value(self.use_viewport_resolution)
            use_transparent_bg = self._get_value(self.transparent_background)
            spp_val = int(self._get_value(self.spp))

            # Determine resolution
            if use_viewport and camera_data.get('viewport_width') and camera_data.get('viewport_height'):
                width_val = int(camera_data['viewport_width'])
                height_val = int(camera_data['viewport_height'])
                logger.info(f"Using viewport resolution: {width_val}x{height_val}")
            else:
                width_val = int(self._get_value(self.width))
                height_val = int(self._get_value(self.height))
                logger.info(f"Using custom resolution: {width_val}x{height_val}")

            # Build Mitsuba scene using utility function
            mi_scene = scene_to_mitsuba(
                app_scene,
                camera=camera_data,
                include_lights=True,
                spp=spp_val,
                width=width_val,
                height=height_val,
                transparent_background=use_transparent_bg
            )

            logger.info("Rendering...")
            image = mi.render(mi_scene)

            # Convert to numpy (RGBA)
            image_np = np.array(image)
            logger.info(f"Raw image shape: {image_np.shape}")

            # Tone mapping (simple clamp and gamma) for RGB channels
            rgb = image_np[..., :3]
            rgb = np.clip(rgb, 0, 1)
            rgb = np.power(rgb, 1/2.2)  # Gamma correction
            rgb = (rgb * 255).astype(np.uint8)

            # Alpha channel
            if image_np.shape[-1] >= 4:
                alpha = image_np[..., 3]
                alpha = np.clip(alpha, 0, 1)
                alpha = (alpha * 255).astype(np.uint8)
            else:
                # If no alpha, create one (fully opaque where there's color)
                alpha = np.where(np.sum(rgb, axis=-1) > 0, 255, 0).astype(np.uint8)

            # Combine to RGBA
            image_np = np.dstack([rgb, alpha])

            logger.info(f"Final image shape: {image_np.shape}")

            # Update preview (auto-converts numpy array to base64)
            self.preview = image_np

            logger.info("Render complete")
            return f"Rendered {width_val}x{height_val} @ {spp_val} spp"

        except Exception as e:
            import traceback
            logger.info(f"Render error: {e}")
            traceback.print_exc()
            return f"Render failed: {str(e)}"