"""
Neural Simulation Engine for IIT Analysis
Modular architecture supporting various connectivity patterns and state dynamics
"""

import math
import json
import random
import hashlib
import datetime
from typing import Dict, List, Optional, Tuple, Any
import itertools
import collections


class Node:
    """Individual neural node with state and connectivity"""
    
    def __init__(self, node_id: int, threshold: float = 0.5):
        self.id = node_id
        self.state = 0.0
        self.threshold = threshold
        self.inputs = []
        self.outputs = []
        self.weights = {}
        
    def update_state(self, external_input: float = 0.0) -> float:
        """Update node state based on inputs and threshold"""
        weighted_sum = external_input
        for source_id, weight in self.weights.items():
            weighted_sum += weight * 0.0  # Placeholder for source state
        
        # Sigmoid activation
        self.state = 1.0 / (1.0 + math.exp(-weighted_sum + self.threshold))
        return self.state


class ConnectivityPattern:
    """Defines network connectivity patterns"""
    
    @staticmethod
    def fully_connected(n_nodes: int) -> List[Tuple[int, int, float]]:
        """Create fully connected network"""
        connections = []
        for i in range(n_nodes):
            for j in range(n_nodes):
                if i != j:
                    weight = random.uniform(-1.0, 1.0)
                    connections.append((i, j, weight))
        return connections
    
    @staticmethod
    def small_world(n_nodes: int, k: int = 4, p: float = 0.3) -> List[Tuple[int, int, float]]:
        """Create small-world network (Watts-Strogatz)"""
        connections = []
        
        # Create regular ring lattice
        for i in range(n_nodes):
            for j in range(1, k // 2 + 1):
                target = (i + j) % n_nodes
                weight = random.uniform(-1.0, 1.0)
                connections.append((i, target, weight))
        
        # Rewire with probability p
        for idx in range(len(connections)):
            if random.random() < p:
                source, _, weight = connections[idx]
                new_target = random.randint(0, n_nodes - 1)
                while new_target == source:
                    new_target = random.randint(0, n_nodes - 1)
                connections[idx] = (source, new_target, weight)
        
        return connections
    
    @staticmethod
    def scale_free(n_nodes: int, m: int = 2) -> List[Tuple[int, int, float]]:
        """Create scale-free network (Barabási-Albert)"""
        connections = []
        degrees = [0] * n_nodes
        
        # Start with small complete graph
        initial_nodes = min(m + 1, n_nodes)
        for i in range(initial_nodes):
            for j in range(i + 1, initial_nodes):
                weight = random.uniform(-1.0, 1.0)
                connections.append((i, j, weight))
                connections.append((j, i, weight))
                degrees[i] += 1
                degrees[j] += 1
        
        # Add remaining nodes with preferential attachment
        for new_node in range(initial_nodes, n_nodes):
            # Select m existing nodes based on degree
            targets = []
            for _ in range(m):
                total_degree = sum(degrees[:new_node])
                target = 0  # Initialize target
                if total_degree == 0:
                    target = random.randint(0, new_node - 1)
                else:
                    r = random.uniform(0, total_degree)
                    cumsum = 0
                    for i in range(new_node):
                        cumsum += degrees[i]
                        if cumsum >= r:
                            target = i
                            break
                
                weight = random.uniform(-1.0, 1.0)
                connections.append((new_node, target, weight))
                connections.append((target, new_node, weight))
                degrees[new_node] += 1
                degrees[target] += 1
                targets.append(target)
        
        return connections


class NeuralNetwork:
    """Main neural network simulation class"""
    
    def __init__(self, n_nodes: int, pattern_type: str = "fully_connected"):
        self.n_nodes = n_nodes
        self.nodes = [Node(i) for i in range(n_nodes)]
        self.time = 0
        self.state_history = []
        self.pattern_type = pattern_type
        
        # Build connectivity
        if pattern_type == "fully_connected":
            connections = ConnectivityPattern.fully_connected(n_nodes)
        elif pattern_type == "small_world":
            connections = ConnectivityPattern.small_world(n_nodes)
        elif pattern_type == "scale_free":
            connections = ConnectivityPattern.scale_free(n_nodes)
        else:
            connections = ConnectivityPattern.fully_connected(n_nodes)
        
        # Apply connections
        for source, target, weight in connections:
            self.nodes[source].outputs.append(target)
            self.nodes[target].inputs.append(source)
            self.nodes[target].weights[source] = weight
    
    def step(self, external_inputs: Optional[List[float]] = None) -> List[float]:
        """Perform one simulation step"""
        if external_inputs is None:
            external_inputs = [0.0] * self.n_nodes
        
        # Calculate new states
        new_states = []
        for i, node in enumerate(self.nodes):
            # Get weighted sum from sources
            weighted_sum = external_inputs[i]
            for source_id in node.inputs:
                weighted_sum += node.weights[source_id] * self.nodes[source_id].state
            
            # Update with sigmoid
            new_state = 1.0 / (1.0 + math.exp(-weighted_sum + node.threshold))
            new_states.append(new_state)
        
        # Apply new states
        for i, state in enumerate(new_states):
            self.nodes[i].state = state
        
        self.time += 1
        self.state_history.append(new_states.copy())
        return new_states
    
    def get_state_vector(self) -> List[float]:
        """Get current state vector"""
        return [node.state for node in self.nodes]
    
    def get_connectivity_matrix(self) -> List[List[float]]:
        """Get adjacency matrix of weights"""
        matrix = [[0.0] * self.n_nodes for _ in range(self.n_nodes)]
        for i, node in enumerate(self.nodes):
            for j, weight in node.weights.items():
                matrix[i][j] = weight
        return matrix
    
    def reset(self):
        """Reset network to initial state"""
        for node in self.nodes:
            node.state = 0.0
        self.time = 0
        self.state_history = []