import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from typing import Dict, List, Set, Any
from .organism import OrganismManager
from .base import BaseEvolvableModule

class SurvivalEngine:
    """
    The Immune System & Garbage Collector.
    Monitors Information Gain (KL Div) and removes inefficient layers.
    """
    def __init__(self, manager: OrganismManager, config: Any, debug: bool = False):
        self.manager = manager
        self.config = config
        self.utility_threshold = config.utility_threshold
        self.debug = debug

    def calculate_information_gain(self, module: BaseEvolvableModule, input_tensor: torch.Tensor, output_tensor: torch.Tensor):
        """
        Updates module's utility score based on KL Divergence.
        High KL -> High Info Gain -> High Utility.
        Low KL (Identity) -> Low Utility.
        """
        with torch.no_grad():
            # If shapes differ, there IS information processing (dim change).
            if input_tensor.shape != output_tensor.shape:
                kl_val = 1.0 # Arbitrary high score for structural changes
            else:
                # Approximate KL Div for tensors as distributions
                # We normalize to probability distributions (Softmax) then compare.
                # This is a heuristic for "signal change".
                
                # Flatten -> Softmax
                p = F.softmax(input_tensor.flatten(1), dim=1) + 1e-8
                q = F.softmax(output_tensor.flatten(1), dim=1) + 1e-8
                
                # KL(P || Q) = sum(p * log(p/q))
                kl_val = torch.sum(p * (torch.log(p) - torch.log(q)), dim=1).mean().item()
            
            # Update Score (Exponential Moving Average)
            alpha = 0.1
            module.utility_score = (1 - alpha) * module.utility_score + alpha * kl_val
            
            # Check Status
            self._update_status(module)

    def _update_status(self, module: BaseEvolvableModule):
        """Transitions module states based on utility."""
        if module.is_immortal:
            return
            
        if module.utility_score < self.utility_threshold:
            module.is_dying = True
        else:
            # Recovery!
            module.is_dying = False
            module.grace_period_steps = 100 # Reset grace

    def apply_selective_decay(self, optimizer: optim.Optimizer):
        """
        Increases weight decay for dying layers.
        """
        if not self.config.enable_negative_cleanup:
            return

        # Map params to modules is hard without a map.
        # We iterate registry.
        
        for uuid, module in self.manager.registry.items():
            if module.is_immortal:
                continue
                
            if module.is_dying:
                if module.grace_period_steps > 0:
                    module.grace_period_steps -= 1
                    continue # Grace period active
                
                # Rot the weights
                decay_rate = self.config.pruning_decay_constant
                
                # Instead of modifying optimizer group (complex), let's manually decay weights
                # w = w - lr * decay * w -> w *= (1 - decay)
                for p in module.parameters():
                    with torch.no_grad():
                        p.mul_(1.0 - decay_rate)
                
                if self.debug:
                    print(f"Decaying layer {uuid[:8]} (Utility: {module.utility_score:.4f})")

    def prune_candidates(self) -> List[str]:
        """IDs of layers ready for deletion."""
        candidates = []
        for uuid, module in self.manager.registry.items():
            if module.is_immortal: continue
            
            # If dying and grace period over and weights almost zero (optional check)
            # Or just if utility stays low for long enough.
            if module.is_dying and module.grace_period_steps <= 0:
                # Add check: are weights basically zero?
                # normalized norm
                norm = sum(p.norm().item() for p in module.parameters())
                if norm < 0.1: # Threshold for "Dead"
                    candidates.append(uuid)
        return candidates

    def garbage_collect(self, network: nn.Module) -> int:
        """
        Darwinian GC.
        Removes dead layers IF graph remains connected/redundant.
        Returns count of removed layers.
        """
        candidates = self.prune_candidates()
        removed_count = 0
        
        for uuid in candidates:
            if self.check_redundancy(uuid):
                # Safe to remove
                self._remove_layer(uuid, network)
                removed_count += 1
                
        return removed_count

    def check_redundancy(self, target_uuid: str) -> bool:
        """
        Can we remove this node without breaking the graph?
        1. If it's a leaf (no children), yes (unless output node).
        2. If it makes a bridge, NO.
        """
        # Simplified connectivity check:
        # If removing it leaves any node unreachable from inputs, it's critical.
        
        # 1. Temporarily simulate removal
        # We need graph traversal logic.
        
        # Get parents and children
        children = self.manager.adjacency_list.get(target_uuid, [])
        parents = self.manager.incoming_edges.get(target_uuid, [])
        
        if not children:
            # It's a dead end. Safe if not a required output.
            if target_uuid in self.manager.output_nodes:
                # Only save it if it's explicitly immortal (protected output)
                module = self.manager.registry.get(target_uuid)
                if module and module.is_immortal:
                    return False
            return True
            
        # Bridge Check:
        # Can parents still reach children (or other paths exist)?
        # If we remove Node X, can we bypass it?
        # Usually, if we prune a hidden layer, we might want to Rewire Parent -> Child directly?
        # Or does pruning mean destroying the path?
        
        # User Logic: "Redundancy Check" -> "Is there a neighbor doing the same job?"
        # If not, we might be breaking a unique feature.
        
        # For Phase 4: We assume pruning = Removing the node AND its connections.
        # If this breaks the path from Input -> Output, we block it.
        
        # Simple Logic: Only prune if it's a dead branch (no children) OR we rewire.
        # Let's implement: "Prune Only Dead Branches" for safety 
        # OR "Bypass" if dimensions match.
        
        # If Bypass is possible (Parent Dim == Child Input Dim), we rewire.
        # But we don't know dims here easily.
        
        # Strict Connectivity for Demo:
        # Only remove if it has NO CHILDREN (Tip of branch).
        if not children:
            return True
            
        return False # Protect bridges for now (safer)

    def _remove_layer(self, uuid: str, network: nn.Module):
        """Physical deletion."""
        print(f"Garbage Collecting {uuid}...")
        
        # 1. Remove from Adjacency (Parents pointing to it)
        parents = self.manager.incoming_edges.get(uuid, [])
        for p in parents:
            if uuid in self.manager.adjacency_list[p]:
                self.manager.adjacency_list[p].remove(uuid)
                
        # 2. Remove from Registry
        del self.manager.registry[uuid]
        del self.manager.adjacency_list[uuid]
        del self.manager.incoming_edges[uuid]
        
        if uuid in self.manager.input_nodes:
            self.manager.input_nodes.remove(uuid)
        if uuid in self.manager.output_nodes:
            self.manager.output_nodes.remove(uuid)
            
        # 3. Remove from Network ModuleDict (Physical VRAM clear)
        if uuid in network.modules_dict:
            del network.modules_dict[uuid]
