import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Any, Optional
from .base import BaseEvolvableModule

class SeedBlock(BaseEvolvableModule):
    """
    The 'Seed' Cell.
    A specialized block capable of genetic masking (partial freezing).
    Stores mask configuration to be inherited by clones.
    """
    def __init__(self, creation_step: int = 0, input_dim: int = 64, hidden_dim: int = 128):
        # Standard MLP block for demo purposes (Transformer block logic would go here)
        child = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim)
        )
        super().__init__(child, creation_step)
        
        # Genetic Memory
        self.mask_config: Dict[str, bool] = {} 

    def apply_mask(self, config: Dict[str, bool]):
        """
        Applies a mask configuration to freeze specific parts of the block.
        Example config: {'layer_0': True, 'layer_2': False} (True=Freeze, False=Train)
        """
        self.mask_config = config
        
        # Iterate through named parameters or sub-modules based on logic
        # For this simple sequential, we map indices to config keys
        # In a real Transformer, keys would be 'attention', 'mlp', etc.
        
        for name, param in self.child_module.named_parameters():
             # Logic to map param name to config key. 
             # For this demo, let's assume config keys match part of the name or index.
             # E.g. "0.weight" -> key "0"
             
             for key, freeze in config.items():
                 if key in name:
                     param.requires_grad = not freeze
                     
    def get_metadata(self) -> Dict[str, Any]:
        meta = super().get_metadata()
        meta['mask_config'] = self.mask_config
        return meta

class IdentityWrapper(BaseEvolvableModule):
    """
    The 'Stabilizer'.
    Wraps a new layer and slowly introduces its effect to the network.
    Output = x + (Layer(x) * gamma * gate)
    """
    def __init__(self, child_module: nn.Module, creation_step: int = 0, stability_threshold: float = 0.5):
        super().__init__(child_module, creation_step)
        
        # Learnable Gamma (starts at 0.0 for smooth entry)
        self.gamma = nn.Parameter(torch.zeros(1))
        
        self.stability_threshold = stability_threshold
        self.rolling_stability_score = 0.0 
        self.alpha = 0.9 
        
    def forward(self, x: torch.Tensor, **kwargs) -> Any:
        # 1. Compute Inner Layer Output
        if self.is_frozen:
             with torch.no_grad():
                result = self.child_module(x, **kwargs)
        else:
            result = self.child_module(x, **kwargs)
            
        # Handle tuple return (e.g. KV Cache)
        if isinstance(result, tuple):
             inner_out = result[0]
             aux = result[1:] # Keep other returns
        else:
             inner_out = result
             aux = None
            
        self._update_energy_score(x, inner_out)

        # 2. Return Residual with Gamma Gating
        # 2. Return Residual with Gamma Gating
        # Hotfix v0.1.7: Check for shape mismatch (e.g. Tokenizer -> Embedding or Pooler)
        if x.shape == inner_out.shape:
            out_residual = x + (inner_out * self.gamma.to(x.device))
        else:
            # Dimensions differ, cannot add residual. Return inner_out directly.
            # This happens when the layer changes dimensions (Embedding, Flatten, etc.)
            out_residual = inner_out
        
        if aux is not None:
             return (out_residual,) + aux
        return out_residual

    def get_metadata(self):
        meta = super().get_metadata()
        meta['gamma'] = self.gamma.item()
        meta['stability_score'] = self.rolling_stability_score
        return meta
