virtual_gpu/render.py

"""
Render Module - Software Raster Pipeline

This module implements the software raster pipeline for drawing primitives
and images onto framebuffers stored in VRAM.
"""

import numpy as np
from typing import Tuple, Optional, Any, Dict
import time


class Renderer:
    """
    Software-based renderer that implements basic drawing operations.
    
    This renderer operates on framebuffers stored in VRAM and provides
    functions for drawing primitives like rectangles, lines, and pixels.
    """
    
    def __init__(self, vram=None):
        self.vram = vram
        self.current_shader = None
        
        # Rendering statistics
        self.pixels_drawn = 0
        self.draw_calls = 0
        self.render_time = 0.0
        
    def set_vram(self, vram):
        """Set the VRAM reference."""
        self.vram = vram
        
    def set_shader(self, shader):
        """Set the current shader for rendering operations."""
        self.current_shader = shader
        
    def clear(self, framebuffer_id: str, color: Tuple[int, int, int] = (0, 0, 0)) -> bool:
        """Clear a framebuffer with the specified color."""
        if not self.vram:
            return False
            
        start_time = time.time()
        
        framebuffer = self.vram.get_framebuffer(framebuffer_id)
        if not framebuffer:
            return False
            
        try:
            framebuffer.clear(color)
            self.pixels_drawn += framebuffer.width * framebuffer.height
            self.draw_calls += 1
            self.render_time += time.time() - start_time
            return True
        except Exception as e:
            print(f"Error clearing framebuffer {framebuffer_id}: {e}")
            return False
            
    def draw_pixel(self, framebuffer_id: str, x: int, y: int, 
                   color: Tuple[int, int, int] = (255, 255, 255)) -> bool:
        """Draw a single pixel on the framebuffer."""
        if not self.vram:
            return False
            
        start_time = time.time()
        
        framebuffer = self.vram.get_framebuffer(framebuffer_id)
        if not framebuffer:
            return False
            
        try:
            # Apply shader if available
            final_color = color
            if self.current_shader:
                final_color = self.current_shader.process_pixel(x, y, color)
                
            framebuffer.set_pixel(x, y, final_color)
            self.pixels_drawn += 1
            self.draw_calls += 1
            self.render_time += time.time() - start_time
            return True
        except Exception as e:
            print(f"Error drawing pixel at ({x}, {y}): {e}")
            return False
            
    def draw_rect(self, framebuffer_id: str, x: int, y: int, width: int, height: int,
                  color: Tuple[int, int, int] = (255, 255, 255)) -> bool:
        """Draw a filled rectangle on the framebuffer."""
        if not self.vram:
            return False
            
        start_time = time.time()
        
        framebuffer = self.vram.get_framebuffer(framebuffer_id)
        if not framebuffer:
            return False
            
        try:
            # Clamp rectangle to framebuffer bounds
            x1 = max(0, x)
            y1 = max(0, y)
            x2 = min(framebuffer.width, x + width)
            y2 = min(framebuffer.height, y + height)
            
            if x2 <= x1 or y2 <= y1:
                return True  # Nothing to draw
                
            # Use NumPy for efficient rectangle filling
            if self.current_shader:
                # Apply shader to each pixel (slower but more flexible)
                for py in range(y1, y2):
                    for px in range(x1, x2):
                        final_color = self.current_shader.process_pixel(px, py, color)
                        framebuffer.pixel_buffer[py, px] = final_color[:framebuffer.channels]
            else:
                # Direct fill (faster)
                framebuffer.pixel_buffer[y1:y2, x1:x2] = color[:framebuffer.channels]
                
            pixels_affected = (x2 - x1) * (y2 - y1)
            self.pixels_drawn += pixels_affected
            self.draw_calls += 1
            self.render_time += time.time() - start_time
            return True
            
        except Exception as e:
            print(f"Error drawing rectangle at ({x}, {y}, {width}, {height}): {e}")
            return False
            
    def draw_line(self, framebuffer_id: str, x1: int, y1: int, x2: int, y2: int,
                  color: Tuple[int, int, int] = (255, 255, 255)) -> bool:
        """Draw a line using Bresenham's algorithm."""
        if not self.vram:
            return False
            
        start_time = time.time()
        
        framebuffer = self.vram.get_framebuffer(framebuffer_id)
        if not framebuffer:
            return False
            
        try:
            # Bresenham's line algorithm
            dx = abs(x2 - x1)
            dy = abs(y2 - y1)
            sx = 1 if x1 < x2 else -1
            sy = 1 if y1 < y2 else -1
            err = dx - dy
            
            x, y = x1, y1
            pixels_drawn = 0
            
            while True:
                # Draw pixel if within bounds
                if 0 <= x < framebuffer.width and 0 <= y < framebuffer.height:
                    final_color = color
                    if self.current_shader:
                        final_color = self.current_shader.process_pixel(x, y, color)
                    framebuffer.set_pixel(x, y, final_color)
                    pixels_drawn += 1
                    
                if x == x2 and y == y2:
                    break
                    
                e2 = 2 * err
                if e2 > -dy:
                    err -= dy
                    x += sx
                if e2 < dx:
                    err += dx
                    y += sy
                    
            self.pixels_drawn += pixels_drawn
            self.draw_calls += 1
            self.render_time += time.time() - start_time
            return True
            
        except Exception as e:
            print(f"Error drawing line from ({x1}, {y1}) to ({x2}, {y2}): {e}")
            return False
            
    def draw_circle(self, framebuffer_id: str, center_x: int, center_y: int, radius: int,
                    color: Tuple[int, int, int] = (255, 255, 255), filled: bool = False) -> bool:
        """Draw a circle using the midpoint circle algorithm."""
        if not self.vram:
            return False
            
        start_time = time.time()
        
        framebuffer = self.vram.get_framebuffer(framebuffer_id)
        if not framebuffer:
            return False
            
        try:
            pixels_drawn = 0
            
            if filled:
                # Draw filled circle
                for y in range(center_y - radius, center_y + radius + 1):
                    for x in range(center_x - radius, center_x + radius + 1):
                        if (x - center_x) ** 2 + (y - center_y) ** 2 <= radius ** 2:
                            if 0 <= x < framebuffer.width and 0 <= y < framebuffer.height:
                                final_color = color
                                if self.current_shader:
                                    final_color = self.current_shader.process_pixel(x, y, color)
                                framebuffer.set_pixel(x, y, final_color)
                                pixels_drawn += 1
            else:
                # Draw circle outline using midpoint algorithm
                x = 0
                y = radius
                d = 1 - radius
                
                def draw_circle_points(cx, cy, x, y):
                    points = [
                        (cx + x, cy + y), (cx - x, cy + y),
                        (cx + x, cy - y), (cx - x, cy - y),
                        (cx + y, cy + x), (cx - y, cy + x),
                        (cx + y, cy - x), (cx - y, cy - x)
                    ]
                    drawn = 0
                    for px, py in points:
                        if 0 <= px < framebuffer.width and 0 <= py < framebuffer.height:
                            final_color = color
                            if self.current_shader:
                                final_color = self.current_shader.process_pixel(px, py, color)
                            framebuffer.set_pixel(px, py, final_color)
                            drawn += 1
                    return drawn
                
                pixels_drawn += draw_circle_points(center_x, center_y, x, y)
                
                while x < y:
                    if d < 0:
                        d += 2 * x + 3
                    else:
                        d += 2 * (x - y) + 5
                        y -= 1
                    x += 1
                    pixels_drawn += draw_circle_points(center_x, center_y, x, y)
                    
            self.pixels_drawn += pixels_drawn
            self.draw_calls += 1
            self.render_time += time.time() - start_time
            return True
            
        except Exception as e:
            print(f"Error drawing circle at ({center_x}, {center_y}) with radius {radius}: {e}")
            return False
            
    def draw_image(self, framebuffer_id: str, x: int, y: int, texture_id: str,
                   scale_x: float = 1.0, scale_y: float = 1.0) -> bool:
        """Draw an image/texture onto the framebuffer."""
        if not self.vram:
            return False
            
        start_time = time.time()
        
        framebuffer = self.vram.get_framebuffer(framebuffer_id)
        texture = self.vram.get_texture(texture_id)
        
        if not framebuffer or texture is None:
            return False
            
        try:
            # Get texture dimensions
            if len(texture.shape) == 3:
                tex_height, tex_width, tex_channels = texture.shape
            else:
                tex_height, tex_width = texture.shape
                tex_channels = 1
                
            # Calculate scaled dimensions
            scaled_width = int(tex_width * scale_x)
            scaled_height = int(tex_height * scale_y)
            
            pixels_drawn = 0
            
            # Simple nearest-neighbor scaling and blitting
            for dy in range(scaled_height):
                for dx in range(scaled_width):
                    # Calculate destination pixel
                    dest_x = x + dx
                    dest_y = y + dy
                    
                    # Check bounds
                    if (dest_x < 0 or dest_x >= framebuffer.width or 
                        dest_y < 0 or dest_y >= framebuffer.height):
                        continue
                        
                    # Calculate source pixel (nearest neighbor)
                    src_x = int(dx / scale_x)
                    src_y = int(dy / scale_y)
                    
                    # Clamp source coordinates
                    src_x = min(src_x, tex_width - 1)
                    src_y = min(src_y, tex_height - 1)
                    
                    # Get source pixel color
                    if tex_channels == 1:
                        color = (texture[src_y, src_x], texture[src_y, src_x], texture[src_y, src_x])
                    else:
                        color = tuple(texture[src_y, src_x, :min(3, tex_channels)])
                        
                    # Apply shader if available
                    final_color = color
                    if self.current_shader:
                        final_color = self.current_shader.process_pixel(dest_x, dest_y, color)
                        
                    # Set pixel
                    framebuffer.set_pixel(dest_x, dest_y, final_color)
                    pixels_drawn += 1
                    
            self.pixels_drawn += pixels_drawn
            self.draw_calls += 1
            self.render_time += time.time() - start_time
            return True
            
        except Exception as e:
            print(f"Error drawing image {texture_id} at ({x}, {y}): {e}")
            return False
            
    def get_stats(self) -> Dict[str, Any]:
        """Get rendering statistics."""
        return {
            "pixels_drawn": self.pixels_drawn,
            "draw_calls": self.draw_calls,
            "total_render_time": self.render_time,
            "avg_render_time": self.render_time / max(1, self.draw_calls),
            "pixels_per_second": self.pixels_drawn / max(0.001, self.render_time)
        }
        
    def reset_stats(self) -> None:
        """Reset rendering statistics."""
        self.pixels_drawn = 0
        self.draw_calls = 0
        self.render_time = 0.0


if __name__ == "__main__":
    # Test the renderer
    from vram import VRAM
    
    # Create VRAM and renderer
    vram = VRAM(memory_size_gb=1)
    renderer = Renderer(vram)
    
    # Create a test framebuffer
    fb_id = vram.create_framebuffer(800, 600, 3)
    
    # Test rendering operations
    print("Testing renderer...")
    
    # Clear screen
    renderer.clear(fb_id, (64, 128, 255))
    
    # Draw some rectangles
    renderer.draw_rect(fb_id, 100, 100, 200, 150, (255, 0, 0))
    renderer.draw_rect(fb_id, 200, 200, 100, 100, (0, 255, 0))
    
    # Draw some lines
    renderer.draw_line(fb_id, 0, 0, 799, 599, (255, 255, 255))
    renderer.draw_line(fb_id, 799, 0, 0, 599, (255, 255, 255))
    
    # Draw a circle
    renderer.draw_circle(fb_id, 400, 300, 50, (255, 255, 0), filled=True)
    
    # Draw some pixels
    for i in range(100):
        renderer.draw_pixel(fb_id, 50 + i, 50, (255, 0, 255))
        
    # Print statistics
    stats = renderer.get_stats()
    print(f"Renderer stats: {stats}")
    
    # Get framebuffer and check a pixel
    fb = vram.get_framebuffer(fb_id)
    if fb:
        pixel = fb.get_pixel(100, 100)
        print(f"Pixel at (100, 100): {pixel}")
        
    print("Renderer test completed!")