# cognitive_net/network.py
import torch
import torch.nn as nn
import torch.optim as optim
import math
from typing import Dict, List, Optional
from .node import CognitiveNode

class DynamicCognitiveNet(nn.Module):
    """Self-organizing neural architecture with structural plasticity"""
    def __init__(self, input_size: int, output_size: int):
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        
        # Neural population
        self.input_nodes = nn.ModuleList([
            CognitiveNode(i, 1) for i in range(input_size)
        ])
        self.output_nodes = nn.ModuleList([
            CognitiveNode(input_size + i, 1) for i in range(output_size)
        ])
        
        # Structural configuration
        self.connections: Dict[str, nn.Parameter] = nn.ParameterDict()
        self._init_base_connections()
        
        # Meta-learning components
        self.emotional_state = nn.Parameter(torch.tensor(0.0))
        self.optimizer = optim.AdamW(self.parameters(), lr=0.001)
        self.loss_fn = nn.MSELoss()

    def _init_base_connections(self):
        """Initialize sparse input-output connectivity"""
        for i, in_node in enumerate(self.input_nodes):
            for j, out_node in enumerate(self.output_nodes):
                conn_id = f"{in_node.id}->{out_node.id}"
                self.connections[conn_id] = nn.Parameter(
                    torch.randn(1) * 0.1
                )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Input processing
        activations = {}
        for i, node in enumerate(self.input_nodes):
            activations[node.id] = node(x[i].unsqueeze(0))
        
        # Output integration
        outputs = []
        for out_node in self.output_nodes:
            integrated = []
            for in_node in self.input_nodes:
                conn_id = f"{in_node.id}->{out_node.id}"
                weight = torch.sigmoid(self.connections[conn_id])
                integrated.append(activations[in_node.id] * weight)
            
            if integrated:
                combined = sum(integrated) / math.sqrt(len(integrated))
                outputs.append(out_node(combined))
        
        return torch.cat(outputs)

    def structural_update(self, global_reward: float):
        """Evolutionary architecture modification"""
        # Connection strength adaptation
        for conn_id, weight in self.connections.items():
            if global_reward > 0:
                new_weight = weight + 0.1 * global_reward
            else:
                new_weight = weight * 0.95
            self.connections[conn_id].data = new_weight.clamp(-1, 1)
        
        # Structural neurogenesis
        if global_reward < -0.5:
            new_conn = self._generate_connection()
            if new_conn not in self.connections:
                self.connections[new_conn] = nn.Parameter(
                    torch.randn(1) * 0.1
                )

    def _generate_connection(self) -> str:
        """Create new input-output connection based on activity"""
        input_act = {n.id: np.mean(n.recent_activations) 
                    for n in self.input_nodes if n.recent_activations}
        output_act = {n.id: np.mean(n.recent_activations)
                     for n in self.output_nodes if n.recent_activations}
        
        if not input_act or not output_act:
            return ""
            
        src = min(input_act, key=input_act.get)  # type: ignore
        tgt = min(output_act, key=output_act.get)  # type: ignore
        return f"{src}->{tgt}"

    def train_step(self, x: torch.Tensor, y: torch.Tensor) -> float:
        self.optimizer.zero_grad()
        pred = self(x)
        loss = self.loss_fn(pred, y)
        
        # Structural regularization
        reg_loss = sum(p.abs().mean() for p in self.connections.values())
        total_loss = loss + 0.01 * reg_loss
        
        total_loss.backward()
        self.optimizer.step()
        
        # Emotional state adaptation
        self.emotional_state.data = torch.sigmoid(
            self.emotional_state + (0.5 - loss.item()) * 0.1
        )
        
        # Structural reorganization
        self.structural_update(0.5 - loss.item())
        
        return total_loss.item()