renderer.py

# Renderer - Handles all drawing operations for the neural network visualization

import pygame
import numpy as np
import math
from config import *
import random

class NetworkRenderer:
    def __init__(self, screen, layer_sizes):
        self.screen = screen
        self.layer_sizes = layer_sizes
        self.num_layers = len(layer_sizes)
     
        # Calculate neuron postitions
        self.neuron_positions = self._calculate_positions()
     
        # Fonts
        pygame.font.init()
        self.font_large = pygame.font.Font(None, FONT_SIZE_LARGE)
        self.font_medium = pygame.font.Font(None, FONT_SIZE_MEDIUM)
        self.font_small = pygame.font.Font(None, FONT_SIZE_SMALL)
     
        # Animation State
        self.current_activations = None
        self.target_activations = None
        self.animation_progress = 1.0
     
    def _calculate_positions(self):
        # Calculate screen positions for all neurons
        positions = []
        max_neurons = max(self.layer_sizes)
     
        # Calculate available space
        total_width = NETWORK_PANEL_WIDTH - 100 # MARGINS
        layer_spacing = total_width / (self.num_layers - 1) if self.num_layers > 1 else 0
        
        # Use consistent vertical spacing based on max layer size
        available_height = WINDOW_HEIGHT - 200
     
        for layer_idx, layer_size in enumerate(self.layer_sizes):
            layer_positions = []
         
            # X position
            x = 50 + layer_idx * layer_spacing
            if layer_size > 1:
                spacing = available_height / (layer_size -1)
            else:
                spacing = 0
            # Center this layer vertically
            start_y = (WINDOW_HEIGHT - (layer_size - 1) * spacing) / 2
         
            for neuron_idx in range(layer_size):
                y = start_y + neuron_idx * spacing
                layer_positions.append((int(x), int(y)))
            
            positions.append(layer_positions)
         
        return positions
     
    def draw_connections(self, activations=None, weights=None, strength_threshold=0.2):
        # Draw connections between Neurons - Only show strong weights
        if weights is None:
            return # Can't draw connections without weights

        # Single surface for all connections
        surface = pygame.Surface((NETWORK_PANEL_WIDTH, WINDOW_HEIGHT), pygame.SRCALPHA)
        pulse = abs(math.sin(pygame.time.get_ticks() / 800.0))
        connection_color = get_color("connection")

        for layer_idx in range(self.num_layers -1):
            for i, (x1, y1) in enumerate(self.neuron_positions[layer_idx]):
                for j, (x2, y2) in enumerate(self.neuron_positions[layer_idx + 1]):

                    # Get the actual weight value for this connection
                    weight_value = abs(weights[layer_idx][i, j])
                    # Only draw if weight is above threshold
                    if weight_value < strength_threshold:
                        continue

                    # Calculate connection strength based on weight magnitude
                    normalized_weight = min(weight_value / 2.0, 1.0)

                    # Add subtle pulsing to connections
                    alpha = int(CONNECTION_MIN_ALPHA +
                                (CONNECTION_MAX_ALPHA - CONNECTION_MIN_ALPHA) * normalized_weight * (0.8 + 0.2 * pulse))

                    color = (*connection_color, alpha)
                    pygame.draw.line(surface, color, (x1, y1), (x2, y2), CONNECTION_WIDTH)

        self.screen.blit(surface, (0, 0))
     
    def draw_neurons(self, activations=None):
        # Draw all neurons with activation coloring and pulsing
        # Add time-based pulsing
        pulse = abs(math.sin(pygame.time.get_ticks() / 500.0)) # Oscillates 0-1
     
        for layer_idx, layer_positions in enumerate(self.neuron_positions):
            for neuron_idx, (x, y) in enumerate(layer_positions):
             
                # Get activation value
                activation = 0.0
                if activations is not None and layer_idx < len(activations):
                    if len(activations[layer_idx].shape) > 1:
                        activation = np.mean(activations[layer_idx][:, neuron_idx])
                    else:
                        activation = activations[layer_idx][neuron_idx]
                     
                     
                    # Normalize per layer type
                    if layer_idx == 0:  # Input layer - map standardized range [-3, 3] to [0, 1]
                        activation = (activation + 3) / 6
                    elif layer_idx == len(self.neuron_positions) - 1:  # Output layer
                        activation = activation ** 0.7

                    activation = np.clip(activation, 0, 1)
             
                # Interpolate color based on activation
                inactive = get_color("neuron_inactive")
                active = get_color("neuron_active")
             
                color = (
                    int(inactive[0] + (active[0] - inactive[0]) * activation),
                    int(inactive[1] + (active[1] - inactive[1]) * activation),
                    int(inactive[2] + (active[2] - inactive[2]) * activation)
                )
                # Calculate pulsing radius - active neurons pulse more
                base_radius = NEURON_RADIUS
                pulse_amount = activation * 3 * pulse # Active neurons pulse more
                current_radius = int(base_radius + pulse_amount)
                
                # Draw glow if enabled and neuron is active
                if GLOW_ENABLED and activation > 0.3:
                    
                    glow_radius = int(current_radius + GLOW_RADIUS * activation * (0.8 + 0.4 * pulse))
                    
                    glow_surface = pygame.Surface((glow_radius * 2, glow_radius * 2), pygame.SRCALPHA)
                    
                    glow_alpha = int(100 * activation * (0.7 + 0.3 * pulse))
                    pygame.draw.circle(glow_surface, (*active, glow_alpha),
                                      (glow_radius, glow_radius), glow_radius)
                 
                    self.screen.blit(glow_surface,
                                    (x - glow_radius, y - glow_radius))
                
                # Draw neuron
                pygame.draw.circle(self.screen, color, (x, y), current_radius)
                pygame.draw.circle(self.screen, NEURON_BORDER_COLOR, (x, y),
                                 current_radius, NEURON_BORDER_WIDTH)
      
    def draw_layer_labels(self):
        # Draw labels for each layer
        layer_names = ["Input"] + [f"Hidden {i+1}" for i in range(self.num_layers - 2)] + ["Output"]
        
        for layer_idx, name in enumerate(layer_names):
            if self.neuron_positions[layer_idx]:
                x = self.neuron_positions[layer_idx][0][0]
                y = 30
             
                # Draw layer name
                text = self.font_medium.render(name, True, TITLE_COLOR)
                text_rect = text.get_rect(center=(x, y))
                self.screen.blit(text, text_rect)
             
                # Draw neuron count
                count_text = self.font_small.render(f"{self.layer_sizes[layer_idx]} neurons",
                                                    True, TEXT_COLOR)
                
                count_rect = count_text.get_rect(center=(x, y + 25))
                self.screen.blit(count_text, count_rect)
     
    def draw_title(self, title="Neural Network Visualization"):
        # Draw main title
        text = self.font_large.render(title, True, TITLE_COLOR)
        self.screen.blit(text, (20, WINDOW_HEIGHT - 40))

class StatsRenderer:
    def __init__(self, screen):
        self.screen = screen
        self.x = STATS_PANEL_X
        self.y = STATS_PANEL_Y
        self.width = STATS_PANEL_WIDTH
        self.height = WINDOW_HEIGHT
     
        pygame.font.init()
        self.font_large = pygame.font.Font(None, FONT_SIZE_LARGE)
        self.font_medium = pygame.font.Font(None, FONT_SIZE_MEDIUM)
        self.font_small = pygame.font.Font(None, FONT_SIZE_SMALL)
     
        # Graph data
        self.loss_history = []
        self.accuracy_history = []
     
    def draw_background(self):
        # Draw stats panel background
        pygame.draw.rect(self.screen, STATS_PANEL_COLOR,
                        (self.x, self.y, self.width, self.height))
     
        # Draw seperator line
        pygame.draw.line(self.screen, get_color("accent"),
                        (self.x, 0), (self.x, self.height), 3)
     
    def draw_metrics(self, epoch, loss, accuracy, learning_rate):
        # Draw current training metrics
        y_offset = 50
     
        # Title
        title = self.font_large.render("Training Stats", True, TITLE_COLOR)
        self.screen.blit(title, (self.x + 20, y_offset))
        y_offset += 60
     
        # Metrics
        metrics = [
                f"Epoch: {epoch}",
                f"Loss: {loss:.6f}",
                f"Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)",
                f"Learning Rate: {learning_rate}"
        ]
        for metric in metrics:
            text = self.font_medium.render(metric, True, TEXT_COLOR)
            self.screen.blit(text, (self.x + 20, y_offset))
            y_offset += 35
         
    def draw_graphs(self):
        # Draw loss and accuracy graphs
        if len(self.loss_history) < 2:
            return
     
        graph_y = 300
        graph_height = 150
        graph_width = self.width - 60
     
        # Draw loss graph
        self._draw_line_graph(
            self.loss_history[-GRAPH_HISTORY_LENGTH:],
            self.x + 30, graph_y,
            graph_width, graph_height,
            "Loss", (255, 100, 100)
        )
        # Draw accuracy graph
        self._draw_line_graph(
            self.accuracy_history[-GRAPH_HISTORY_LENGTH:],
            self.x + 30, graph_y + graph_height + 80,
            graph_width, graph_height,
            "Accuracy", (100, 255, 100),
            y_range=(0,1)
        )
    def _draw_line_graph(self, data, x, y, width, height, title, color, y_range=None):
        # Draw a line graph
        # Background
        pygame.draw.rect(self.screen, (30, 30, 40), (x, y, width, height))
        pygame.draw.rect(self.screen, (60, 60, 70), (x, y, width, height), 1)
     
        # Title
        title_text = self.font_medium.render(title, True, TITLE_COLOR)
        self.screen.blit(title_text, (x + 5, y - 25))
     
        if len(data) < 2:
            return
     
        # Calculate scale
        if y_range:
            min_val, max_val = y_range
        else:
            min_val = min(data)
            max_val = max(data)
          
        value_range = max_val - min_val if max_val != min_val else 1
     
        # Draw points
        points = []
        for i, value in enumerate(data):
            px = x + (i / (len(data) -1)) * width
            py = y + height - ((value - min_val) / value_range) * height
            points.append((px, py))
     
        # Draw line
        if len(points) > 1:
            pygame.draw.lines(self.screen, color, False, points, 2)
        
        # Draw current value
        current_val = data[-1]
        val_text = self.font_small.render(f"{current_val:.4f}", True, color)
        self.screen.blit(val_text, (x + width - 80, y + 5))
     
    def update_data(self, loss, accuracy):
        # Add new data point
        self.loss_history.append(loss)
        self.accuracy_history.append(accuracy)