
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
        
        # Baseline Weights (Mid-phase / Balanced)
        # M1=Literal, M2=Growth, M3=Pruning, M4=Reorder, M5=Subtree, M6=OpSub
        # Exploitation Focus: M1, M3. Exploration Focus: M2, M5, M6.
        self.baseline_weights = np.array([0.35, 0.15, 0.15, 0.10, 0.15, 0.10], dtype=float)
        
        # State tracking (Simple exponential smoothing)
        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

    def update_state(self, stats: dict):
        """
        Updates internal state based on population statistics.
        Calculates target weights and smooths current weights towards them.
        
        stats: {
            'stagnation': int (count),
            'diversity': float (score),
            'bloat': float (avg_size),
            'invalidity': float (fraction)
        }
        """
        if self.strategy == 'fixed': return

        # Extract Targets
        # 1. Stagnation / Diversity -> Exploration
        # If stagnated (>10) or Low Diversity (<0.01), boost Exploration
        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 # Threshold depends on metric scaling
        
        if is_stagnated or is_low_diversity:
             # Exploration Profile: Boost Growth (M2), Subtree (M5), OpSub (M6)
             # Reduce Literal (M1)
             target = np.array([0.10, 0.30, 0.10, 0.10, 0.30, 0.10])
        
        # 2. Bloat -> Pruning
        # If bloat is high (>50 nodes?), boost Pruning (M3)
        bloat = stats.get('bloat', 0)
        if bloat > 50:
            # Pruning Profile: High M3. Reduce Growth M2.
            # Mix with current target
            pruning_target = np.array([0.30, 0.05, 0.40, 0.10, 0.10, 0.05])
            # Blend 50/50 with existing intent
            target = 0.5 * target + 0.5 * pruning_target
            
        # 3. Invalidity -> Safety
        # If invalidity is high (>0.5), maybe reduce destructive mutations?
        # Or Just penalty does the job. 
        # For now, let's assume penalty handles it, but maybe boost M1 (safest)?
        invalid = stats.get('invalidity', 0.0)
        if invalid > 0.5:
            # Safety Profile: Mostly M1 (Tweak constants)
            safety_target = np.array([0.60, 0.10, 0.10, 0.10, 0.05, 0.05])
            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
        
        # Ensure sum 1
        self.current_weights /= self.current_weights.sum()

    def get_weights(self, progress=None):
        """
        Returns the mutation weights.
        If strategy is 'adaptive', returns the learned state.
        If 'linear', uses the deterministic schedule mixed with state if updated.
        """
        if self.strategy == 'fixed':
            return self.baseline_weights
            
        if self.strategy == 'linear':
            # Deterministic override if no stats yet? 
            # Or mix? 
            # Let's use the explicit phase logic as a "base" target, but update_state moves away from it.
            # If update_state is effectively called, we trust current_weights.
            # But update_state might not be called in 'linear' mode if we strictly follow phase 2 logic.
            # Let's keep Phase 2 logic purely deterministic for consistency with previous verification if adaptivity isn't triggered.
            
            # However, user wants Blueprint B "Adaptive Scheduler". 
            # Let's assume strategy='adaptive' will be used for the new logic, 
            # and 'linear' remains backward compatible.
            pass
            
        if self.strategy == 'adaptive':
            return self.current_weights
            
        # Fallback to Phase 2 Linear Logic
        if progress is None: return self.baseline_weights
        
        if progress < 0.3: 
            return [0.20, 0.25, 0.10, 0.10, 0.25, 0.10]
        elif progress > 0.7:
            return [0.50, 0.05, 0.25, 0.10, 0.05, 0.05]
        else:
            return self.baseline_weights
