import torch
import torch.nn as nn
import uuid
import math
from typing import Any

class BaseEvolvableModule(nn.Module):
    """
    The fundamental "cell" of the organism.
    Wraps standard PyTorch modules with biological metadata and energy tracking.
    """
    def __init__(self, child_module: nn.Module, creation_step: int = 0):
        super().__init__()
        self.child_module = child_module
        self.dynamic_id = str(uuid.uuid4())
        self.creation_step = creation_step
        
        # Biological Metrics
        self.energy_score = 0.0  # Accumulated signal magnitude
        self.utility_score = 1.0 # Rolling Info Gain (starts high to give chance)
        self.activation_count = 0
        self.last_activation_magnitude = 0.0
        
        # State flags
        self.is_frozen = False
        self.is_prunable = False 
        self.is_dying = False
        self.is_immortal = False # Core protection
        self.grace_period_steps = 100 # Steps to recover if low utility

    def forward(self, x: torch.Tensor, **kwargs) -> Any:
        # We need to handle tuple returns (e.g. output, cache)
        # and pass kwargs (e.g. past_key_values)
        if self.is_frozen:
            # If frozen, we might still want to track usage, or maybe skip gradient?
            # For now, just forward.
            with torch.no_grad():
                result = self.child_module(x, **kwargs)
        else:
            result = self.child_module(x, **kwargs)
        
        # Determine output tensor for energy tracking
        if isinstance(result, tuple):
            output = result[0]
        else:
            output = result
            
        self._update_energy_score(x, output)
        return result

    def _update_energy_score(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor):
        """
        Calculates how much 'work' the layer did.
        Logic: Work = ||Output - Input|| (if shapes match) OR ||Output||
        """
        with torch.no_grad():
            self.activation_count += 1
            
            # Calculate signal magnitude (energy)
            # We use L2 norm averaged over batch size
            if input_tensor.shape == output_tensor.shape:
                # Work done is the transformation magnitude
                diff = output_tensor - input_tensor
                magnitude = diff.norm(p=2).item() / (diff.numel() + 1e-8)
            else:
                # Dimension change implies work
                magnitude = output_tensor.norm(p=2).item() / (output_tensor.numel() + 1e-8)
            
            self.last_activation_magnitude = magnitude
            self.energy_score += magnitude

    def get_metadata(self):
        return {
            "uuid": self.dynamic_id,
            "type": self.child_module.__class__.__name__,
            "creation_step": self.creation_step,
            "energy_score": self.energy_score,
            "activation_count": self.activation_count,
            "frozen": self.is_frozen
        }

    def reset_stats(self):
        self.energy_score = 0.0
        self.activation_count = 0
