
import numpy as np

class MutationScheduler:
    """
    Manages the probabilities (weights) of different mutation strategies
    based on the evolutionary progress.
    """
    
    def __init__(self, strategy='linear'):
        self.strategy = strategy
        
        # Aggressive Growth Baseline (M2, M5, M10 prioritized to escape constants)
        # M1=Perturbation, M2=Growth, M3=Shrink, M4=Reorder, M5=Subtree, M6=OpSub
        # M7=Hoist, M8=Point, M9=Affine, M10=Residual
        # We increase M2 (Growth), M5 (Subtree) and M10 (Residual) significantly.
        self.baseline_weights = np.array([
            0.10, # M1: Constant Perturbation
            0.20, # M2: Insert Growth (Injects Variables)
            0.05, # M3: Delete/Shrink
            0.05, # M4: Reorder
            0.15, # M5: Subtree Replacement (Injects Structure)
            0.10, # M6: OpSub
            0.05, # M7: Hoist
            0.05, # M8: Point
            0.05, # M9: Affine Wrap
            0.20  # M10: Residual (Injects small functional corrections)
        ], dtype=float)

        
        # State tracking
        self.stagnation_factor = 0.0
        self.diversity_factor = 0.0
        self.bloat_factor = 0.0
        self.invalidity_factor = 0.0
        
        # Adaptive current weights (starts at baseline)
        self.current_weights = self.baseline_weights.copy()
        self.learning_rate = 0.1 # Smoothing factor per update
        self.strategy_count = 10

    def update_state(self, stats: dict):
        if self.strategy == 'fixed': return

        # Extract Targets
        stag = stats.get('stagnation', 0)
        div = stats.get('diversity', 1.0)
        
        # Target Profile
        target = self.baseline_weights.copy()
        
        is_stagnated = stag > 10
        is_low_diversity = div < 0.05
        
        if is_stagnated or is_low_diversity:
             # Exploration Profile: High Growth (M2), Subtree (M5), OpSub (M6)
             target = np.array([0.05, 0.25, 0.05, 0.05, 0.20, 0.15, 0.05, 0.10, 0.05, 0.05])
        
        # 2. Bloat -> Pruning
        bloat = stats.get('bloat', 0)
        if bloat > 50:
            # Pruning Profile: High M3 (Delete) and M7 (Hoist).
            pruning_target = np.array([0.10, 0.05, 0.25, 0.10, 0.05, 0.05, 0.25, 0.05, 0.05, 0.05])
            target = 0.5 * target + 0.5 * pruning_target
            
        # 3. Invalidity -> Safety
        invalid = stats.get('invalidity', 0.0)
        if invalid > 0.5:
            # Safety Profile: Mostly M1 and M8 (Tweak constants/operators without changing structure)
            safety_target = np.array([0.40, 0.05, 0.05, 0.10, 0.05, 0.05, 0.05, 0.20, 0.05, 0.00])
            target = 0.5 * target + 0.5 * safety_target
            
        # Normalize Target
        target /= target.sum()
        
        # Smooth Update
        self.current_weights = (1.0 - self.learning_rate) * self.current_weights + self.learning_rate * target
        self.current_weights /= self.current_weights.sum()

    def get_weights(self, progress=None):
        if self.strategy == 'fixed':
            return self.baseline_weights
            
        if self.strategy == 'adaptive':
            return self.current_weights
            
        # Fallback to Phase 2 Linear Logic (Deterministic Phases)
        if progress is None: return self.baseline_weights
        
        if progress < 0.3: 
            # Early: High Exploration (Growth + Subtree + OpSub)
            return [0.10, 0.20, 0.05, 0.05, 0.20, 0.15, 0.05, 0.05, 0.05, 0.10]
        elif progress > 0.7:
            # Late: High Exploitation (Perturbe + Affine + Residual)
            return [0.30, 0.02, 0.05, 0.10, 0.03, 0.05, 0.05, 0.05, 0.20, 0.15]
        else:
            return self.baseline_weights
