import torch
import torch.nn as nn
from typing import Dict, List, Optional, Any, Tuple
import json
from .base import BaseEvolvableModule

class OrganismManager:
    """
    The Central Nervous System (Brain Map).
    Manages the Directed Graph of modules (nodes) and their connections (edges).
    """
    def get_topology_summary(self) -> Dict[str, Any]:
        """Returns a serializable summary of the graph topology."""
        return {
            "nodes": [m.get_metadata() for m in self.registry.values()],
            "edges": self.adjacency_list,
            "inputs": self.input_nodes,
            "outputs": self.output_nodes,
            "step": self.step_counter,
            "history": self.genome_history
        }

    # --- Phase 5: History & Reporting ---
    
    def log_event(self, event_type: str, details: str):
        """Logs an evolutionary event."""
        # Ideally, this should be called by Growth/Survival engines
        if not hasattr(self, 'genome_history'):
             self.genome_history = []
        
        event = {
            "step": self.step_counter,
            "type": event_type,
            "details": details
        }
        self.genome_history.append(event)

    def summarize_evolution(self) -> str:
        """Returns a text report of the organism's life."""
        if not hasattr(self, 'genome_history'):
             self.genome_history = []
             
        # Count events
        births = sum(1 for e in self.genome_history if e['type'] == 'BIRTH')
        deaths = sum(1 for e in self.genome_history if e['type'] == 'DEATH')
        pruned = sum(1 for e in self.genome_history if e['type'] == 'PRUNE')
        
        report = []
        report.append(f"📜 Evolution Report (Age: {self.step_counter} steps)")
        report.append(f"   • Total Births (Neurogenesis): {births}")
        report.append(f"   • Total Deaths (Garbage Collection): {deaths}")
        report.append(f"   • Synaptic Pruning Events: {pruned}")
        report.append(f"   • Current active layers: {len(self.registry)}")
        
        return "\n".join(report)

    def __init__(self):
        # Nodes: UUID -> BaseEvolvableModule
        self.registry: Dict[str, BaseEvolvableModule] = {}
        
        # Edges: adjacency_list[uuid] = [next_uuid1, next_uuid2, ...]
        self.adjacency_list: Dict[str, List[str]] = {}
        
        # Reverse lookup for backtracking: incoming_edges[uuid] = [prev_uuid1, ...]
        self.incoming_edges: Dict[str, List[str]] = {}
        
        # Inputs/Outputs tracking
        self.input_nodes: List[str] = []
        self.output_nodes: List[str] = []
        
        self.step_counter = 0
        self.genome_history: List[Dict] = [] # Phase 5

    def register_module(self, module: BaseEvolvableModule, parents: List[str] = None):
        """
        Registers a new module into the organism graph.
        
        Args:
            module: The BaseEvolvableModule to register.
            parents: List of UUIDs that connect TO this module. 
                     If None or empty, this is treated as an INPUT node (connected to external input).
        """
        uuid = module.dynamic_id
        if uuid in self.registry:
            raise ValueError(f"Module with UUID {uuid} already exists.")
        
        self.registry[uuid] = module
        self.adjacency_list[uuid] = [] # Initialize children list
        self.incoming_edges[uuid] = [] # Initialize parents list
        
        if not parents:
            self.input_nodes.append(uuid)
        else:
            for p_uuid in parents:
                if p_uuid not in self.registry:
                    raise ValueError(f"Parent UUID {p_uuid} does not exist.")
                
                # Connect Parent -> Child
                self.adjacency_list[p_uuid].append(uuid)
                # Connect Child <- Parent
                self.incoming_edges[uuid].append(p_uuid)
                
                # If parent was an output node, it might not be anymore (unless it branches)
                # But typically we explicitly define output nodes. For now, let's keep it dynamic.
                if p_uuid in self.output_nodes:
                    self.output_nodes.remove(p_uuid)

        # By default, a new leaf node is an output node until something connects to it
        self.output_nodes.append(uuid)
    
    def get_module(self, uuid: str) -> BaseEvolvableModule:
        return self.registry.get(uuid)

    def get_forward_candidates(self, current_uuid: str) -> List[str]:
        """Returns valid next steps from the current node."""
        return self.adjacency_list.get(current_uuid, [])

    def get_backward_candidates(self, current_uuid: str) -> List[str]:
        """Returns parents for feedback loops."""
        return self.incoming_edges.get(current_uuid, [])

    def get_topology_summary(self) -> Dict[str, Any]:
        """Returns a serializable summary of the graph topology."""
        return {
            "nodes": [m.get_metadata() for m in self.registry.values()],
            "edges": self.adjacency_list,
            "inputs": self.input_nodes,
            "outputs": self.output_nodes,
            "step": self.step_counter
        }
