virtual_gpu/shader.py

"""
Shader Module - Programmable Shader Logic

This module provides a mechanism for simulating programmable shader logic,
allowing custom functions to be applied to pixels or vertices during rendering.
"""

import numpy as np
from typing import Callable, Dict, Any, Tuple, Optional
from abc import ABC, abstractmethod
import math


class Shader(ABC):
    """Abstract base class for all shaders."""
    
    @abstractmethod
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Process a single pixel and return the modified color."""
        pass
        
    @abstractmethod
    def process_vertex(self, x: float, y: float, z: float = 0.0, 
                      **kwargs) -> Tuple[float, float, float]:
        """Process a single vertex and return the modified position."""
        pass


class PixelShader(Shader):
    """Base class for pixel shaders that only modify pixel colors."""
    
    def process_vertex(self, x: float, y: float, z: float = 0.0, 
                      **kwargs) -> Tuple[float, float, float]:
        """Default vertex processing (no change)."""
        return (x, y, z)


class VertexShader(Shader):
    """Base class for vertex shaders that only modify vertex positions."""
    
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Default pixel processing (no change)."""
        return color


class ColorTintShader(PixelShader):
    """Shader that applies a color tint to all pixels."""
    
    def __init__(self, tint_color: Tuple[float, float, float], strength: float = 0.5):
        self.tint_color = tint_color
        self.strength = strength
        
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Apply color tint to the pixel."""
        r, g, b = color
        tr, tg, tb = self.tint_color
        
        # Blend original color with tint
        new_r = int(r * (1 - self.strength) + tr * 255 * self.strength)
        new_g = int(g * (1 - self.strength) + tg * 255 * self.strength)
        new_b = int(b * (1 - self.strength) + tb * 255 * self.strength)
        
        # Clamp values
        new_r = max(0, min(255, new_r))
        new_g = max(0, min(255, new_g))
        new_b = max(0, min(255, new_b))
        
        return (new_r, new_g, new_b)


class GrayscaleShader(PixelShader):
    """Shader that converts colors to grayscale."""
    
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Convert pixel to grayscale."""
        r, g, b = color
        
        # Use luminance formula for better grayscale conversion
        gray = int(0.299 * r + 0.587 * g + 0.114 * b)
        gray = max(0, min(255, gray))
        
        return (gray, gray, gray)


class SepiaShader(PixelShader):
    """Shader that applies a sepia tone effect."""
    
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Apply sepia tone to the pixel."""
        r, g, b = color
        
        # Sepia transformation matrix
        new_r = int(0.393 * r + 0.769 * g + 0.189 * b)
        new_g = int(0.349 * r + 0.686 * g + 0.168 * b)
        new_b = int(0.272 * r + 0.534 * g + 0.131 * b)
        
        # Clamp values
        new_r = max(0, min(255, new_r))
        new_g = max(0, min(255, new_g))
        new_b = max(0, min(255, new_b))
        
        return (new_r, new_g, new_b)


class InvertShader(PixelShader):
    """Shader that inverts pixel colors."""
    
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Invert pixel colors."""
        r, g, b = color
        return (255 - r, 255 - g, 255 - b)


class BrightnessShader(PixelShader):
    """Shader that adjusts pixel brightness."""
    
    def __init__(self, brightness: float = 0.0):
        """
        Initialize brightness shader.
        
        Args:
            brightness: Brightness adjustment (-1.0 to 1.0)
                       -1.0 = completely dark, 0.0 = no change, 1.0 = completely bright
        """
        self.brightness = max(-1.0, min(1.0, brightness))
        
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Adjust pixel brightness."""
        r, g, b = color
        
        if self.brightness >= 0:
            # Brighten
            new_r = int(r + (255 - r) * self.brightness)
            new_g = int(g + (255 - g) * self.brightness)
            new_b = int(b + (255 - b) * self.brightness)
        else:
            # Darken
            new_r = int(r * (1 + self.brightness))
            new_g = int(g * (1 + self.brightness))
            new_b = int(b * (1 + self.brightness))
            
        # Clamp values
        new_r = max(0, min(255, new_r))
        new_g = max(0, min(255, new_g))
        new_b = max(0, min(255, new_b))
        
        return (new_r, new_g, new_b)


class ContrastShader(PixelShader):
    """Shader that adjusts pixel contrast."""
    
    def __init__(self, contrast: float = 0.0):
        """
        Initialize contrast shader.
        
        Args:
            contrast: Contrast adjustment (-1.0 to 1.0)
                     -1.0 = no contrast, 0.0 = no change, 1.0 = maximum contrast
        """
        self.contrast = max(-1.0, min(1.0, contrast))
        self.factor = (259 * (self.contrast * 255 + 255)) / (255 * (259 - self.contrast * 255))
        
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Adjust pixel contrast."""
        r, g, b = color
        
        new_r = int(self.factor * (r - 128) + 128)
        new_g = int(self.factor * (g - 128) + 128)
        new_b = int(self.factor * (b - 128) + 128)
        
        # Clamp values
        new_r = max(0, min(255, new_r))
        new_g = max(0, min(255, new_g))
        new_b = max(0, min(255, new_b))
        
        return (new_r, new_g, new_b)


class CheckerboardShader(PixelShader):
    """Shader that creates a checkerboard pattern overlay."""
    
    def __init__(self, size: int = 8, color1: Tuple[int, int, int] = (255, 255, 255),
                 color2: Tuple[int, int, int] = (0, 0, 0), blend: float = 0.5):
        self.size = size
        self.color1 = color1
        self.color2 = color2
        self.blend = blend
        
    def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                     **kwargs) -> Tuple[int, int, int]:
        """Apply checkerboard pattern."""
        # Determine which checker square we're in
        checker_x = x // self.size
        checker_y = y // self.size
        
        # Determine checker color
        if (checker_x + checker_y) % 2 == 0:
            checker_color = self.color1
        else:
            checker_color = self.color2
            
        # Blend with original color
        r, g, b = color
        cr, cg, cb = checker_color
        
        new_r = int(r * (1 - self.blend) + cr * self.blend)
        new_g = int(g * (1 - self.blend) + cg * self.blend)
        new_b = int(b * (1 - self.blend) + cb * self.blend)
        
        return (new_r, new_g, new_b)


class WaveDistortionShader(VertexShader):
    """Shader that applies wave distortion to vertices."""
    
    def __init__(self, amplitude: float = 10.0, frequency: float = 0.1, time: float = 0.0):
        self.amplitude = amplitude
        self.frequency = frequency
        self.time = time
        
    def process_vertex(self, x: float, y: float, z: float = 0.0, 
                      **kwargs) -> Tuple[float, float, float]:
        """Apply wave distortion to vertex position."""
        # Apply sine wave distortion
        offset_x = self.amplitude * math.sin(y * self.frequency + self.time)
        offset_y = self.amplitude * math.sin(x * self.frequency + self.time)
        
        return (x + offset_x, y + offset_y, z)
        
    def update_time(self, time: float):
        """Update the time parameter for animation."""
        self.time = time


class ShaderManager:
    """Manages shader instances and provides shader registry functionality."""
    
    def __init__(self):
        self.shaders: Dict[str, Shader] = {}
        self.shader_counter = 0
        
        # Register built-in shaders
        self._register_builtin_shaders()
        
    def _register_builtin_shaders(self):
        """Register built-in shader types."""
        self.register_shader("grayscale", GrayscaleShader())
        self.register_shader("sepia", SepiaShader())
        self.register_shader("invert", InvertShader())
        self.register_shader("red_tint", ColorTintShader((1.0, 0.0, 0.0), 0.3))
        self.register_shader("blue_tint", ColorTintShader((0.0, 0.0, 1.0), 0.3))
        self.register_shader("bright", BrightnessShader(0.3))
        self.register_shader("dark", BrightnessShader(-0.3))
        self.register_shader("high_contrast", ContrastShader(0.5))
        self.register_shader("checkerboard", CheckerboardShader())
        
    def register_shader(self, name: str, shader: Shader) -> None:
        """Register a shader with a given name."""
        self.shaders[name] = shader
        
    def get_shader(self, name: str) -> Optional[Shader]:
        """Get a shader by name."""
        return self.shaders.get(name)
        
    def create_custom_shader(self, pixel_func: Optional[Callable] = None,
                           vertex_func: Optional[Callable] = None,
                           name: Optional[str] = None) -> str:
        """Create a custom shader from functions."""
        if name is None:
            name = f"custom_shader_{self.shader_counter}"
            self.shader_counter += 1
            
        class CustomShader(Shader):
            def __init__(self, pf, vf):
                self.pixel_func = pf
                self.vertex_func = vf
                
            def process_pixel(self, x: int, y: int, color: Tuple[int, int, int], 
                            **kwargs) -> Tuple[int, int, int]:
                if self.pixel_func:
                    return self.pixel_func(x, y, color, **kwargs)
                return color
                
            def process_vertex(self, x: float, y: float, z: float = 0.0, 
                             **kwargs) -> Tuple[float, float, float]:
                if self.vertex_func:
                    return self.vertex_func(x, y, z, **kwargs)
                return (x, y, z)
                
        shader = CustomShader(pixel_func, vertex_func)
        self.register_shader(name, shader)
        return name
        
    def list_shaders(self) -> list:
        """Get a list of all registered shader names."""
        return list(self.shaders.keys())
        
    def get_stats(self) -> Dict[str, Any]:
        """Get shader manager statistics."""
        return {
            "total_shaders": len(self.shaders),
            "shader_names": self.list_shaders()
        }


if __name__ == "__main__":
    # Test the shader system
    print("Testing Shader System...")
    
    # Create shader manager
    shader_manager = ShaderManager()
    
    # Test built-in shaders
    test_color = (128, 64, 192)
    test_x, test_y = 100, 50
    
    print(f"Original color: {test_color}")
    
    # Test each built-in shader
    for shader_name in shader_manager.list_shaders():
        shader = shader_manager.get_shader(shader_name)
        if shader:
            result_color = shader.process_pixel(test_x, test_y, test_color)
            print(f"{shader_name}: {result_color}")
            
    # Test custom shader
    def rainbow_pixel(x, y, color, **kwargs):
        """Custom shader that creates a rainbow effect based on position."""
        r, g, b = color
        
        # Create rainbow effect based on x position
        hue = (x % 360) / 360.0
        
        # Simple HSV to RGB conversion for rainbow effect
        if hue < 1/6:
            new_r, new_g, new_b = 255, int(255 * hue * 6), 0
        elif hue < 2/6:
            new_r, new_g, new_b = int(255 * (2/6 - hue) * 6), 255, 0
        elif hue < 3/6:
            new_r, new_g, new_b = 0, 255, int(255 * (hue - 2/6) * 6)
        elif hue < 4/6:
            new_r, new_g, new_b = 0, int(255 * (4/6 - hue) * 6), 255
        elif hue < 5/6:
            new_r, new_g, new_b = int(255 * (hue - 4/6) * 6), 0, 255
        else:
            new_r, new_g, new_b = 255, 0, int(255 * (1 - hue) * 6)
            
        # Blend with original color
        blend = 0.7
        final_r = int(r * (1 - blend) + new_r * blend)
        final_g = int(g * (1 - blend) + new_g * blend)
        final_b = int(b * (1 - blend) + new_b * blend)
        
        return (final_r, final_g, final_b)
        
    # Register custom shader
    custom_name = shader_manager.create_custom_shader(pixel_func=rainbow_pixel, name="rainbow")
    custom_shader = shader_manager.get_shader(custom_name)
    
    if custom_shader:
        for x in range(0, 360, 60):
            result = custom_shader.process_pixel(x, 0, test_color)
            print(f"Rainbow shader at x={x}: {result}")
            
    # Test vertex shader
    wave_shader = WaveDistortionShader(amplitude=5.0, frequency=0.1)
    test_vertex = (100.0, 50.0, 0.0)
    distorted_vertex = wave_shader.process_vertex(*test_vertex)
    print(f"Original vertex: {test_vertex}")
    print(f"Distorted vertex: {distorted_vertex}")
    
    # Print statistics
    stats = shader_manager.get_stats()
    print(f"Shader Manager stats: {stats}")
    
    print("Shader system test completed!")