neat\network.py

"""Neural network implementation for BackpropNEAT."""

import jax
import jax.numpy as jnp
import numpy as np
from typing import Dict, List, Optional, Tuple, Union
from .genome import Genome
import copy
import random

class Network:
    """Neural network for NEAT implementation.
    Implements a strictly feed-forward network following original NEAT principles:
    1. Start minimal - direct input-output connections only
    2. Complexify gradually through structural mutations
    3. Protect innovation through speciation
    4. No recurrent connections (as per requirements)
    """
    def __init__(self, genome: Genome):
        """Initialize network from genome."""
        # Store genome and sizes
        self.genome = genome
        
        # Verify genome sizes match volleyball requirements
        if genome.input_size != 12 or genome.output_size != 3:
            print(f"Warning: Genome size mismatch. Expected 12 inputs, 3 outputs. Got {genome.input_size} inputs, {genome.output_size} outputs")
            genome.input_size = 12
            genome.output_size = 3
        
        self.input_size = 12  # Fixed for volleyball
        self.output_size = 3  # Fixed for volleyball
        
        # Deep copy to avoid shared references
        self.node_genes = {}
        self.connection_genes = []
        
        # Create input nodes (0 to 11)
        for i in range(12):
            self.node_genes[i] = NodeGene(i, 'input', 'linear')
        
        # Create bias node (12)
        self.node_genes[12] = NodeGene(12, 'bias', 'linear')
        
        # Create output nodes (13, 14, 15)
        for i in range(3):
            node_id = 13 + i
            self.node_genes[node_id] = NodeGene(node_id, 'output', 'sigmoid')
            
            # Connect to bias with appropriate weight based on action type
            if i < 2:  # Left/Right actions: encourage movement
                self.connection_genes.append(
                    ConnectionGene(12, node_id, random.uniform(0.0, 1.0), True)
                )
            else:  # Jump action: neutral bias
                self.connection_genes.append(
                    ConnectionGene(12, node_id, random.uniform(-0.5, 0.5), True)
                )
            
            # Connect to relevant inputs with larger weights
            if i == 0:  # Left action: connect to ball x position and velocity
                self.connection_genes.append(
                    ConnectionGene(0, node_id, random.uniform(0.5, 1.5), True)  # ball x
                )
                self.connection_genes.append(
                    ConnectionGene(2, node_id, random.uniform(0.5, 1.5), True)  # ball vx
                )
            elif i == 1:  # Right action: connect to ball x position and velocity
                self.connection_genes.append(
                    ConnectionGene(0, node_id, random.uniform(-1.5, -0.5), True)  # ball x
                )
                self.connection_genes.append(
                    ConnectionGene(2, node_id, random.uniform(-1.5, -0.5), True)  # ball vx
                )
            else:  # Jump action: connect to ball y position and velocity
                self.connection_genes.append(
                    ConnectionGene(1, node_id, random.uniform(-1.5, -0.5), True)  # ball y
                )
                self.connection_genes.append(
                    ConnectionGene(3, node_id, random.uniform(-1.0, 0.0), True)  # ball vy
                )
        
        # Copy existing nodes (if any)
        for node_id, node in genome.node_genes.items():
            if node_id not in self.node_genes:  # Skip I/O nodes
                self.node_genes[node_id] = NodeGene(
                    node_id,
                    node.node_type,
                    node.activation
                )
        
        # Copy connections
        if genome.connection_genes:
            # Clear initial connections if genome has its own
            self.connection_genes = []
            for conn in genome.connection_genes:
                # Verify connection nodes exist
                if conn.source not in self.node_genes or conn.target not in self.node_genes:
                    print(f"Warning: Connection {conn.source}->{conn.target} references missing nodes")
                    continue
                self.connection_genes.append(ConnectionGene(
                    conn.source,
                    conn.target,
                    conn.weight,
                    conn.enabled
                ))
        
        # Verify output connections (13, 14, 15)
        for output_id in [13, 14, 15]:
            has_connection = False
            for conn in self.connection_genes:
                if conn.enabled and conn.target == output_id:
                    has_connection = True
                    break
            
            if not has_connection:
                print(f"Adding missing connections for output {output_id}")
                # Connect to bias
                self.connection_genes.append(
                    ConnectionGene(12, output_id, random.uniform(-1.0, 1.0), True)
                )
                # Connect to random input
                input_id = random.randint(0, 11)
                self.connection_genes.append(
                    ConnectionGene(input_id, output_id, random.uniform(-1.0, 1.0), True)
                )
        
        # Build evaluation order
        self.node_evals = {}
        self._build_feed_forward_order()
        
        # Verify all outputs are properly connected
        self._verify_outputs()
    
    def _verify_outputs(self):
        """Verify all outputs have valid connections and evaluations."""
        output_ids = {13, 14, 15}  # Fixed output IDs
        
        # Check node evaluations
        for output_id in output_ids:
            if output_id not in self.node_evals:
                print(f"Adding missing evaluation for output {output_id}")
                bias_id = 12
                self.node_evals[output_id] = {
                    'inputs': [bias_id],
                    'weights': [1.0],
                    'activation': 'sigmoid'
                }
                # Add connection if needed
                if not any(c.target == output_id and c.enabled for c in self.connection_genes):
                    self.connection_genes.append(
                        ConnectionGene(bias_id, output_id, 1.0, True)
                    )
    
    def _create_minimal_connections(self):
        """Create minimal initial connections for a new network."""
        bias_id = 12
        output_start = bias_id + 1
        
        # Connect each output to bias and one random input
        for i in range(self.output_size):
            output_id = output_start + i
            
            # Connect to bias
            self.connection_genes.append(ConnectionGene(
                bias_id, output_id,
                random.uniform(-1.0, 1.0),
                True
            ))
            
            # Connect to random input
            input_id = random.randint(0, self.input_size - 1)
            self.connection_genes.append(ConnectionGene(
                input_id, output_id,
                random.uniform(-1.0, 1.0),
                True
            ))
    
    def _build_feed_forward_order(self):
        """Build evaluation order ensuring feed-forward only topology."""
        try:
            # Fixed node sets for volleyball
            input_nodes = set(range(12))  # 0-11
            bias_node = {12}  # Bias node
            output_nodes = {13, 14, 15}  # Output nodes
            
            # Create adjacency lists
            connections = {}
            for conn in self.connection_genes:
                if not conn.enabled:
                    continue
                if conn.source not in connections:
                    connections[conn.source] = []
                connections[conn.source].append(conn.target)
            
            # Start with inputs and bias evaluated
            evaluated = input_nodes | bias_node
            eval_order = []
            
            # Helper function to check if a node can be evaluated
            def can_evaluate(node_id):
                if node_id in connections:
                    return all(dep in evaluated for dep in connections[node_id])
                return True
            
            # Keep trying to evaluate nodes until we can't anymore
            while True:
                ready_nodes = set()
                for node_id in self.node_genes:
                    if node_id not in evaluated and can_evaluate(node_id):
                        ready_nodes.add(node_id)
                
                if not ready_nodes:
                    break
                
                # Add nodes to evaluation order
                for node_id in sorted(ready_nodes):
                    incoming = []
                    incoming_weights = []
                    for conn in self.connection_genes:
                        if conn.enabled and conn.target == node_id:
                            incoming.append(conn.source)
                            incoming_weights.append(conn.weight)
                    
                    if incoming:  # Only add if node has inputs
                        self.node_evals[node_id] = {
                            'inputs': incoming,
                            'weights': incoming_weights,
                            'activation': self.node_genes[node_id].activation
                        }
                        eval_order.append(node_id)
                    
                    evaluated.add(node_id)
            
            # Ensure all outputs have evaluations
            for output_id in output_nodes:
                if output_id not in self.node_evals:
                    print(f"Adding default evaluation for output {output_id}")
                    # Connect to bias by default
                    self.node_evals[output_id] = {
                        'inputs': [12],  # Bias node
                        'weights': [1.0],
                        'activation': 'sigmoid'
                    }
                    # Add connection if needed
                    if not any(c.target == output_id and c.enabled for c in self.connection_genes):
                        self.connection_genes.append(
                            ConnectionGene(12, output_id, 1.0, True)
                        )
            
        except Exception as e:
            print(f"Error in feed-forward build: {e}")
            # Create minimal fallback evaluations
            self.node_evals = {}
            for i in range(3):  # 3 outputs
                output_id = 13 + i
                self.node_evals[output_id] = {
                    'inputs': [12],  # Bias node
                    'weights': [1.0],
                    'activation': 'sigmoid'
                }
    
    def forward(self, inputs: jnp.ndarray) -> jnp.ndarray:
        """Forward pass through the network."""
        try:
            # Only use first 8 inputs like original network
            inputs = inputs[:8]
            
            # Handle input shape
            original_shape = inputs.shape
            if len(inputs.shape) == 1:
                inputs = inputs.reshape(1, -1)
            batch_size = inputs.shape[0]
            
            # Get max node ID for activation array
            max_node_id = max(node.id for node in self.node_genes.values())
            
            # Initialize activations array
            activations = jnp.zeros((batch_size, max_node_id + 1))
            
            # Set input values (0-7)
            for i in range(8):
                if i < len(inputs):
                    activations = activations.at[:, i].set(inputs[:, i])
                else:
                    activations = activations.at[:, i].set(0.0)
            
            # Initialize recurrent nodes (8-11) with previous outputs
            # For now just use zeros, in the future we could store previous outputs
            for i in range(8, 12):
                activations = activations.at[:, i].set(0.0)
            
            # Evaluate nodes in order (hidden then output)
            for node_id, eval_info in self.node_evals.items():
                try:
                    # Skip input and recurrent nodes
                    if node_id < 12:
                        continue
                        
                    # Get weighted sum of inputs
                    act = jnp.zeros(batch_size)
                    for conn_source, conn_weight in zip(eval_info['inputs'], eval_info['weights']):
                        act += activations[:, conn_source] * conn_weight
                    
                    # Apply activation function
                    if eval_info['activation'] == 'tanh':
                        act = jnp.tanh(act)
                    elif eval_info['activation'] == 'sigmoid':
                        act = jax.nn.sigmoid(act)
                    elif eval_info['activation'] == 'relu':
                        act = jax.nn.relu(act)
                    
                    # Apply threshold like original network for output nodes
                    if node_id >= 20:  # Output nodes
                        act = jnp.where(act > 0.75, 1.0, 0.0)
                    
                    activations = activations.at[:, node_id].set(act)
                except Exception as e:
                    print(f"Error at node {node_id}: {e}")
            
            # Get output node activations
            output = activations[:, -3:]
            
            # Update recurrent nodes for next time step
            # (In a real implementation, we'd need to store these)
            for i in range(8, 12):
                act = jnp.zeros(batch_size)
                for conn_source, conn_weight in zip(eval_info['inputs'], eval_info['weights']):
                    if conn_source >= 20:  # Only use output nodes
                        act += activations[:, conn_source] * conn_weight
                activations = activations.at[:, i].set(jnp.tanh(act))
            
            # Return to original shape
            if len(original_shape) == 1:
                output = output.reshape(-1)
            
            return output
        except Exception as e:
            print(f"Error in forward pass: {e}")
            return jnp.zeros(3)
    
    def predict(self, inputs: jnp.ndarray) -> jnp.ndarray:
        """Make a prediction for the given inputs.
        
        Args:
            inputs: Input array of shape (input_size,) or (batch_size, input_size)
            
        Returns:
            Predictions of shape (3,) for single input or (batch_size, 3) for batch
        """
        outputs = self.forward(inputs)
        
        # Ensure correct output shape for volleyball (always 3 outputs)
        if len(outputs.shape) == 1:
            # Single input case - ensure shape (3,)
            if outputs.shape[0] != 3:
                print(f"Adjusting output shape from {outputs.shape} to (3,)")
                return jnp.pad(outputs, (0, max(0, 3 - outputs.shape[0])))
            return outputs
        else:
            # Batch case - ensure shape (batch_size, 3)
            if outputs.shape[1] != 3:
                print(f"Adjusting output shape from {outputs.shape} to (batch_size, 3)")
                return jnp.pad(outputs, ((0, 0), (0, max(0, 3 - outputs.shape[1]))))
            return outputs
    
    def clone(self) -> 'Network':
        """Create a copy of this network with a cloned genome."""
        return Network(self.genome.clone())
        
    def mutate(self, config: Dict):
        """Mutate the network's genome."""
        self.genome.mutate(config)
        # Rebuild evaluation order after mutation
        self._build_feed_forward_order()
    
    def to_genome(self) -> Genome:
        """Convert network back to genome representation."""
        genome = Genome(self.input_size, self.output_size)
        genome.node_genes = copy.deepcopy(self.node_genes)
        genome.connection_genes = copy.deepcopy(self.connection_genes)
        return genome

class BaseNetwork:
    """Base Network class for NEAT."""

    def __init__(self, n_inputs: int, n_outputs: int):
        self.input_size = n_inputs
        self.output_size = n_outputs
        self.fitness = float('-inf')
        
        # Initialize weights and biases with JAX
        key = jax.random.PRNGKey(0)
        # Use larger initial weights to encourage exploration
        self.weights = jax.random.normal(key, (n_outputs, n_inputs)) * 0.5
        # Add small positive bias to encourage some initial movement
        self.bias = jnp.ones(n_outputs) * 0.1
    
    def forward(self, x: jnp.ndarray) -> jnp.ndarray:
        """Forward pass through the network."""
        if x.ndim > 1:
            # Batched input
            h = jnp.dot(x, self.weights.T) + self.bias[None, :]
        else:
            # Single input
            h = jnp.dot(x, self.weights.T) + self.bias
        return jnp.tanh(h)
    
    def get_params(self) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """Get network parameters."""
        return self.weights, self.bias
    
    def set_params(self, params: Tuple[jnp.ndarray, jnp.ndarray]):
        """Set network parameters."""
        self.weights, self.bias = params
        
    def get_weights_numpy(self) -> np.ndarray:
        """Get weights as numpy array for visualization."""
        return np.array(self.weights)

class NodeGene:
    """Node gene containing node information."""
    def __init__(self, node_id: int, node_type: str, activation: str = 'tanh'):
        """Initialize node gene.
        
        Args:
            node_id: Node ID
            node_type: Type of node ('input', 'hidden', or 'output')
            activation: Activation function ('tanh', 'sigmoid', or 'relu')
        """
        self.id = node_id
        self.type = node_type
        self.activation = activation
        # Initialize with larger random bias for hidden/output nodes
        if node_type in ['hidden', 'output']:
            key = jax.random.PRNGKey(node_id)  # Use node_id as seed for reproducibility
            self.bias = jax.random.normal(key, ()) * 0.5  # Increased from 0.1
        else:
            self.bias = 0.0  # No bias for input nodes

class ConnectionGene:
    """Gene representing a connection between nodes."""
    def __init__(self, source: int, target: int, weight: float = None, enabled: bool = True):
        self.source = source
        self.target = target
        # Initialize with larger weights if not provided
        if weight is None:
            key = jax.random.PRNGKey(hash((source, target)) % 2**32)
            self.weight = jax.random.uniform(key, (), minval=-2.0, maxval=2.0)
        else:
            self.weight = weight
        self.enabled = enabled
        self.innovation = None  # Will be set by NEAT