"""
HyperNEAT (Hypercube-based NeuroEvolution of Augmenting Topologies) implementation.

This module implements the HyperNEAT algorithm, which extends NEAT to evolve
large-scale neural networks with geometric regularities by using a CPPN
(Compositional Pattern Producing Network) to generate the connectivity patterns.
"""

import random
import numpy as np
from typing import List, Dict, Tuple, Optional, Any, Union, Callable
from collections import defaultdict

from neurenix.tensor import Tensor
from neurenix.nn.module import Module
from neurenix.nn.functional import sigmoid, tanh, relu

from .neat import NEAT, NEATConfig, NEATGenome, NEATNetwork, NodeType, NodeGene, ConnectionGene


class Substrate:
    """Represents a substrate in HyperNEAT, which defines the layout of nodes in n-dimensional space."""
    
    def __init__(self, input_coords: List[List[float]], output_coords: List[List[float]], 
                 hidden_coords: Optional[List[List[float]]] = None):
        """Initialize a substrate with node coordinates.
        
        Args:
            input_coords: List of coordinates for input nodes
            output_coords: List of coordinates for output nodes
            hidden_coords: Optional list of coordinates for hidden nodes
        """
        self.input_coords = input_coords
        self.output_coords = output_coords
        self.hidden_coords = hidden_coords or []
        
        if input_coords and output_coords:
            dim = len(input_coords[0])
            for coords in [input_coords, output_coords, hidden_coords or []]:
                for coord in coords:
                    if len(coord) != dim:
                        raise ValueError(f"All coordinates must have the same dimensionality: {dim}")
        
        self.dimensionality = len(input_coords[0]) if input_coords else 0
        
    def get_connection_inputs(self, source_idx: int, target_idx: int, 
                             source_coords: List[float], target_coords: List[float]) -> List[float]:
        """Get the inputs to the CPPN for a connection between two nodes.
        
        Args:
            source_idx: Index of the source node
            target_idx: Index of the target node
            source_coords: Coordinates of the source node
            target_coords: Coordinates of the target node
            
        Returns:
            List of inputs to the CPPN
        """
        inputs = []
        inputs.extend(source_coords)
        inputs.extend(target_coords)
        
        
        return inputs


class CPPN(Module):
    """Compositional Pattern Producing Network used in HyperNEAT."""
    
    def __init__(self, genome: NEATGenome):
        """Initialize a CPPN from a NEAT genome.
        
        Args:
            genome: NEAT genome to use as the CPPN
        """
        super().__init__()
        self.network = NEATNetwork(genome)
        
    def forward(self, x: Tensor) -> Tensor:
        """Forward pass through the CPPN.
        
        Args:
            x: Input tensor of shape (batch_size, num_inputs)
            
        Returns:
            Output tensor of shape (batch_size, num_outputs)
        """
        return self.network.forward(x)
    
    def query_connection(self, inputs: List[float]) -> float:
        """Query the CPPN for a connection weight.
        
        Args:
            inputs: List of inputs to the CPPN
            
        Returns:
            Connection weight
        """
        x = Tensor([inputs])
        output = self.forward(x)
        return output.item()


class HyperNEATNetwork(Module):
    """Neural network generated by HyperNEAT."""
    
    def __init__(self, substrate: Substrate, cppn: CPPN, 
                weight_threshold: float = 0.2, activation: str = 'sigmoid'):
        """Initialize a HyperNEAT-generated neural network.
        
        Args:
            substrate: Substrate defining the layout of nodes
            cppn: CPPN used to generate connection weights
            weight_threshold: Minimum absolute weight value for a connection to be created
            activation: Activation function for nodes
        """
        super().__init__()
        self.substrate = substrate
        self.cppn = cppn
        self.weight_threshold = weight_threshold
        self.activation = activation
        
        self.weights = {}
        self._generate_network()
        
        self.activation_functions = {
            'sigmoid': sigmoid,
            'tanh': tanh,
            'relu': relu,
        }
        
    def _generate_network(self) -> None:
        """Generate the network by querying the CPPN for connection weights."""
        if self.substrate.hidden_coords:
            for i, source_coord in enumerate(self.substrate.input_coords):
                for j, target_coord in enumerate(self.substrate.hidden_coords):
                    inputs = self.substrate.get_connection_inputs(
                        i, j, source_coord, target_coord
                    )
                    weight = self.cppn.query_connection(inputs)
                    
                    if abs(weight) > self.weight_threshold:
                        self.weights[(i, j + len(self.substrate.input_coords))] = weight
        
            for i, source_coord in enumerate(self.substrate.hidden_coords):
                for j, target_coord in enumerate(self.substrate.output_coords):
                    inputs = self.substrate.get_connection_inputs(
                        i + len(self.substrate.input_coords), 
                        j + len(self.substrate.input_coords) + len(self.substrate.hidden_coords),
                        source_coord, target_coord
                    )
                    weight = self.cppn.query_connection(inputs)
                    
                    if abs(weight) > self.weight_threshold:
                        self.weights[(i + len(self.substrate.input_coords), 
                                     j + len(self.substrate.input_coords) + len(self.substrate.hidden_coords))] = weight
        else:
            for i, source_coord in enumerate(self.substrate.input_coords):
                for j, target_coord in enumerate(self.substrate.output_coords):
                    inputs = self.substrate.get_connection_inputs(
                        i, j + len(self.substrate.input_coords),
                        source_coord, target_coord
                    )
                    weight = self.cppn.query_connection(inputs)
                    
                    if abs(weight) > self.weight_threshold:
                        self.weights[(i, j + len(self.substrate.input_coords))] = weight
    
    def forward(self, x: Tensor) -> Tensor:
        """Forward pass through the network.
        
        Args:
            x: Input tensor of shape (batch_size, num_inputs)
            
        Returns:
            Output tensor of shape (batch_size, num_outputs)
        """
        batch_size = x.shape[0]
        
        values = {}
        
        for i in range(len(self.substrate.input_coords)):
            if i < x.shape[1]:
                values[i] = x[:, i]
        
        if self.substrate.hidden_coords:
            for i in range(len(self.substrate.hidden_coords)):
                node_idx = i + len(self.substrate.input_coords)
                incoming_sum = None
                
                for j in range(len(self.substrate.input_coords)):
                    if (j, node_idx) in self.weights:
                        if incoming_sum is None:
                            incoming_sum = values[j] * self.weights[(j, node_idx)]
                        else:
                            incoming_sum = incoming_sum + values[j] * self.weights[(j, node_idx)]
                
                if incoming_sum is not None:
                    activation_fn = self.activation_functions.get(self.activation, sigmoid)
                    values[node_idx] = activation_fn(incoming_sum)
                else:
                    values[node_idx] = Tensor([[0.0]] * batch_size)
        
        outputs = []
        for i in range(len(self.substrate.output_coords)):
            node_idx = i + len(self.substrate.input_coords)
            if self.substrate.hidden_coords:
                node_idx += len(self.substrate.hidden_coords)
                
            incoming_sum = None
            
            if self.substrate.hidden_coords:
                for j in range(len(self.substrate.hidden_coords)):
                    source_idx = j + len(self.substrate.input_coords)
                    if (source_idx, node_idx) in self.weights:
                        if source_idx in values:
                            if incoming_sum is None:
                                incoming_sum = values[source_idx] * self.weights[(source_idx, node_idx)]
                            else:
                                incoming_sum = incoming_sum + values[source_idx] * self.weights[(source_idx, node_idx)]
            else:
                for j in range(len(self.substrate.input_coords)):
                    if (j, node_idx) in self.weights:
                        if incoming_sum is None:
                            incoming_sum = values[j] * self.weights[(j, node_idx)]
                        else:
                            incoming_sum = incoming_sum + values[j] * self.weights[(j, node_idx)]
            
            if incoming_sum is not None:
                activation_fn = self.activation_functions.get(self.activation, sigmoid)
                outputs.append(activation_fn(incoming_sum).unsqueeze(1))
            else:
                outputs.append(Tensor([[0.0]] * batch_size).unsqueeze(1))
        
        if not outputs:
            return Tensor([[0.0]] * batch_size)
            
        return Tensor.cat(outputs, dim=1)


class HyperNEAT:
    """Implementation of the HyperNEAT algorithm."""
    
    def __init__(self, substrate: Substrate, neat_config: Optional[NEATConfig] = None,
                weight_threshold: float = 0.2, activation: str = 'sigmoid'):
        """Initialize the HyperNEAT algorithm.
        
        Args:
            substrate: Substrate defining the layout of nodes
            neat_config: Configuration for the NEAT algorithm
            weight_threshold: Minimum absolute weight value for a connection to be created
            activation: Activation function for nodes
        """
        self.substrate = substrate
        self.neat = NEAT(neat_config)
        self.weight_threshold = weight_threshold
        self.activation = activation
        self.best_network = None
        
    def initialize(self, population_size: int) -> None:
        """Initialize the population of CPPNs.
        
        Args:
            population_size: Size of the population
        """
        cppn_inputs = self.substrate.dimensionality * 2  # Source and target coordinates
        cppn_outputs = 1  # Weight
        
        self.neat.initialize(population_size, cppn_inputs, cppn_outputs)
        
    def evolve(self, fitness_function: Callable[[HyperNEATNetwork], float], generations: int = 100,
              callback: Callable[[int, List[HyperNEATNetwork]], None] = None) -> None:
        """Evolve the population of CPPNs.
        
        Args:
            fitness_function: Function to evaluate fitness of networks
            generations: Number of generations to evolve
            callback: Function called after each generation
        """
        def neat_fitness_function(genome: NEATGenome) -> float:
            cppn = CPPN(genome)
            network = HyperNEATNetwork(
                self.substrate, cppn, 
                self.weight_threshold, self.activation
            )
            return fitness_function(network)
        
        for generation in range(generations):
            self.neat.evolve(neat_fitness_function)
            
            if callback:
                networks = []
                for genome in self.neat.genomes:
                    cppn = CPPN(genome)
                    network = HyperNEATNetwork(
                        self.substrate, cppn,
                        self.weight_threshold, self.activation
                    )
                    networks.append(network)
                    
                callback(generation, networks)
                
        best_genome = self.neat.get_best_genome()
        if best_genome:
            cppn = CPPN(best_genome)
            self.best_network = HyperNEATNetwork(
                self.substrate, cppn,
                self.weight_threshold, self.activation
            )
            
    def get_best_network(self) -> HyperNEATNetwork:
        """Get the best network found by the algorithm.
        
        Returns:
            The best network
        """
        return self.best_network
