import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from typing import List, Dict, Tuple, Optional, Any
from .organism import OrganismManager
from .base import BaseEvolvableModule
import copy
import uuid
from .blueprints import IdentityWrapper

class GrowthEngine:
    """
    The Evolution & Growth Engine.
    Manages Neurogenesis (Growth), Entropy Monitoring, and Optimizer adjustments.
    """
    def __init__(self, network: Any, config: Any):
        # We accept 'network' which should be DynamicNetwork
        self.network = network
        self.manager = network.manager
        self.config = config
        
        # Use Config Params
        self.entropy_threshold = config.entropy_threshold
        self.growth_cooldown = config.growth_cooldown
        # Phase 5: Dynamic Cooldown
        self.evolution_cooldown_steps = getattr(config, 'evolution_cooldown_steps', config.growth_cooldown)
        self.min_loss_improvement = getattr(config, 'min_loss_improvement', 0.0)
        
        self.last_growth_step = 0
        self.last_mitosis_loss = float('inf')
        self.pending_stabilizations = {} 

    def check_growth_trigger(self, entropy_history: List[float], current_step: int, current_loss: Optional[float] = None) -> bool:
        """
        Determines if growth is needed based on 'Wisdom Threshold' and 'Resource Management'.
        Trigger if: 
        1. Entropy is consistently high (Confusion).
        2. Entropy is NOT decreasing over time (Stuck in thought).
        3. Sufficient time passed (evolution_cooldown_steps).
        4. Loss has stabilized/improved (min_loss_improvement).
        """
        # 1. Cooldown Check
        if current_step - self.last_growth_step < self.evolution_cooldown_steps:
            return False
            
        # 2. Loss Improvement Check
        if current_loss is not None and self.last_mitosis_loss != float('inf'):
             if current_loss > self.last_mitosis_loss - self.min_loss_improvement:
                 # We haven't improved enough since last birth.
                 return False
            
        # Check Hardware Constraints
        if self._check_resource_limits():
            return False
            
        if len(entropy_history) < 3:
            return False
            
        recent = np.array(entropy_history[-3:])
        mean_entropy = np.mean(recent)
        delta_entropy = recent[-1] - recent[0] # Change over window
        
        # If confused (High Entropy) AND not improving (Delta >= 0 or very small drop)
        is_confused = mean_entropy > self.entropy_threshold
        is_stuck = delta_entropy > -0.01 # Not decreasing significantly
        
        return is_confused and is_stuck

    def _check_resource_limits(self) -> bool:
        """Returns True if limits reached (Block growth)."""
        # 1. Param Count Check
        total_params = self.calculate_current_params()
        if total_params >= self.config.max_param_count:
            return True
            
        # 2. VRAM Check (Simplified heuristic or actual)
        if torch.cuda.is_available():
            vram_gb = torch.cuda.memory_allocated() / 1e9
            if vram_gb >= self.config.vram_limit_gb:
                return True
                
        return False
        
    def check_resources(self, current_params: int) -> bool:
        return not self._check_resource_limits()

    def calculate_current_params(self) -> int:
        total = 0
        for mod in self.manager.registry.values():
            for p in mod.parameters():
                total += p.numel()
        return total

    def smart_clone(self, parent_uuid: str, optimizer: optim.Optimizer) -> BaseEvolvableModule:
        """
        Mitosis with Variance-Based Noise.
        Returns: New Module (Wrapped in IdentityWrapper with gamma=0.0)
        """
        parent_module = self.manager.get_module(parent_uuid)
        if parent_module is None:
             raise ValueError(f"Parent {parent_uuid} not found")

        # 1. Deep Copy
        new_module = copy.deepcopy(parent_module)
        new_module.dynamic_id = str(uuid.uuid4()) # New soul
        new_module.creation_step = self.manager.step_counter
        new_module.registration_step = self.manager.step_counter
        new_module.reset_stats()
        
        # 2. Add Smart Noise
        for (name, child_param), (new_name, new_param) in zip(parent_module.named_parameters(), new_module.named_parameters()):
            # Find stats
            state = optimizer.state.get(child_param, {})
            
            # AdamW second moment 'exp_avg_sq' (v_t)
            if 'exp_avg_sq' in state:
                # Variance map
                v_t = state['exp_avg_sq']
                # Noise std dev proportional to sqrt(v_t)
                sigma = torch.sqrt(v_t) * 0.1 
                
                noise = torch.randn_like(new_param) * sigma
                with torch.no_grad():
                    new_param.add_(noise)
            else:
                # Fallback
                noise = torch.randn_like(new_param) * 0.01
                with torch.no_grad():
                    new_param.add_(noise)
                    
        # 3. Apply Gamma Gating (IdentityWrapper)
        # Ensure the new layer enters as an Identity operation (gamma=0)
        if isinstance(new_module, IdentityWrapper):
            # Already wrapped (parent was wrapped), just reset gamma
            with torch.no_grad():
                new_module.gamma.fill_(0.0)
        else:
            # Wrap it
            new_module = IdentityWrapper(new_module, creation_step=self.manager.step_counter)
            # Ensure the wrapper has a fresh ID (IdentityWrapper init does this, but we want consistency)
            # Actually IdentityWrapper(new_module) creates a new wrapper with a new ID.
            # The inner module (new_module) is the child.
            pass

        self.last_growth_step = self.manager.step_counter
        return new_module

    def register_clone_to_optimizer(self, new_module: BaseEvolvableModule, optimizer: optim.Optimizer, base_lr: float):
        """
        Cold Start: Add new params with high LR.
        """
        # Add param group
        high_lr = base_lr * 10.0
        
        optimizer.add_param_group({
            'params': new_module.parameters(),
            'lr': high_lr,
            'initial_lr': high_lr,
            'name': f"new_layer_{new_module.dynamic_id}"
        })
        
        # Track for stabilization
        group_idx = len(optimizer.param_groups) - 1
        self.pending_stabilizations[group_idx] = {
            'start_step': self.manager.step_counter,
            'target_lr': base_lr,
            'initial_lr': high_lr
        }

    def stabilize_optimizer(self, optimizer: optim.Optimizer):
        """
        Callback to run every step.
        Decays LR of new layers.
        """
        current_step = self.manager.step_counter
        
        # List to remove
        done = []
        
        for idx, config in self.pending_stabilizations.items():
            if idx >= len(optimizer.param_groups):
                continue # Safety check
                
            elapsed = current_step - config['start_step']
            
            if elapsed >= self.config.stabilization_steps:
                # Reset to target
                optimizer.param_groups[idx]['lr'] = config['target_lr']
                done.append(idx)
            else:
                # Linear decay
                progress = elapsed / self.config.stabilization_steps
                diff = config['initial_lr'] - config['target_lr']
                new_lr = config['initial_lr'] - (progress * diff)
                optimizer.param_groups[idx]['lr'] = new_lr
                
        for i in done:
            del self.pending_stabilizations[i]

    def mitosis(self, parent_uuid: str, optimizer: optim.Optimizer, base_lr: float = 0.01, current_loss: Optional[float] = None):
        """
        High-level Mitosis: Checks triggers, clones, rewires, and updates optimizer.
        """
        if current_loss is not None:
             self.last_mitosis_loss = current_loss
             
        # 1. Clone
        print(f"Mitosis: Cloning {parent_uuid}...")
        clone = self.smart_clone(parent_uuid, optimizer)
        
        # 1.5 Sync Device (Hotfix v0.1.6)
        try:
            parent_module = self.manager.get_module(parent_uuid)
            # Find device from parameters
            device = next(parent_module.parameters()).device
            clone.to(device)
            # Support IdentityWrapper gammas which are buffers/params
            if hasattr(clone, 'gamma'):
                 clone.gamma.data = clone.gamma.data.to(device)
        except Exception as e:
            print(f"[WARNING] Could not sync device for clone: {e}")
        
        # 2. Rewire (Parallel)
        # Get parents of parent
        parents = self.manager.incoming_edges.get(parent_uuid, [])
        
        # Register clone with same parents (Input connection)
        self.manager.register_module(clone, parents)
        
        # Connect to same children (Output connection)
        children = self.manager.adjacency_list.get(parent_uuid, [])
        for child_uuid in children:
            # We need to manually add connection
            self.manager.adjacency_list[clone.dynamic_id].append(child_uuid)
            self.manager.incoming_edges[child_uuid].append(clone.dynamic_id)
            
        # 3. Sync Network Body
        if hasattr(self.network, 'sync_with_registry'):
            self.network.sync_with_registry()
            
        # 4. Optimizer Update
        self.register_clone_to_optimizer(clone, optimizer, base_lr)
        
        # 5. Log
        self.manager.log_event("MITOSIS", f"Cloned {parent_uuid} -> {clone.dynamic_id}")
        print(f"Mitosis: Born {clone.dynamic_id} (Parallel)")
