
import numpy as np
from ezga.variation.variation import Variation_Operator
# from ezga.simple.symbolic_regression import SymbolicMathMutation 
from ezga.simple.semantic import SemanticAnalyzer

class AdaptiveVariation(Variation_Operator):
    """
    Subclass of Variation_Operator that hooks into `adjust_mutation_probabilities`
    to feed population statistics to the SymbolicMathMutation scheduler.
    """
    def adjust_mutation_probabilities(self, dataset, ctx) -> np.ndarray:
        # 1. Gather Stats from Engine Context
        stagnation = getattr(ctx, 'stagnation', 0)
        
        # 2. Compute Bloat (Average Tree Size)
        # Assuming we can inspect structures
        containers = dataset.containers
        sizes = []
        invalids = 0
        
        # Estimate invalidity based on penalized fitness? 
        # Or explicit metadata 'valid'.
        # For now, let's use fitness 1e9 as invalid proxy.
        fitnesses = []
        
        for c in containers:
            # Check Fitness
            f = c.AtomPositionManager.E 
            if f is None or f > 1e8:
                invalids += 1
            if f is not None:
                fitnesses.append(f)
                
        bloat = 0.0 # Placeholder if we can't easily get size here without slow iteration
        invalidity = invalids / max(1, len(containers))
        
        # 3. Compute Semantic Diversity
        # We need signatures.
        # If 'ctx.features' are signatures (from SemanticAnalyzer), we can use them.
        diversity = 1.0
        features = ctx.get_features()
        
        # Heuristic: Check range/std of features to guess if they are signatures
        if features is not None and len(features) > 0:
            # features is N x D
            # Compute mean pairwise distance? Or variance?
            # Simple diversity proxy: average std dev of features across population
            diversity = np.mean(np.std(features, axis=0))
            
        stats = {
            'stagnation': stagnation,
            'diversity': diversity,
            'bloat': bloat, # TODO: pipe tree size
            'invalidity': invalidity
        }
        
        # 4. Pass to Scheduler
        # Find our SymbolicMathMutation instance
        # It's usually in self.mutation_funcs[0]
        for op in self.mutation_funcs:
            # Check if it has a 'mutator' and 'scheduler' attribute using duck typing
            # to avoid importing SymbolicMathMutation (Circular Import)
            if hasattr(op, 'mutator') and hasattr(op.mutator, 'scheduler'):
                 if hasattr(op.mutator.scheduler, 'update_state'):
                    op.mutator.scheduler.update_state(stats)
                    
        # 5. Propagate Stagnation to Problem for Adaptive Parsimony
        for op in self.mutation_funcs:
            if hasattr(op, 'problem'):
                if hasattr(op.problem, 'update_penalty'):
                    op.problem.update_penalty(stagnation)
                    
        return self._mutation_probabilities # Return baseline weights (unused by symbolic ops)
