
import numpy as np
from .grammar import (
    Node, BinaryOperator, UnaryOperator, 
    AND, OR, NOT, 
    Literal, Threshold, Interval, Ratio, Gaussian,
    Add, Sub, Mul, Div, Min, Max,
    Sin, Cos, Exp, Log, Sqrt, Abs, Square, Cube, Neg,
    Variable, Constant, OptimizableConstant,
    SoftIf, TernaryOperator
)
from .scheduler import MutationScheduler

def generate_random_literal(n_features, rng, literal_types=None):
    """Generates a random literal node."""
    # Default to boolean literals if not specified
    if literal_types is None:
        type_idx = rng.choice(4)
        feat = rng.integers(0, n_features)
        
        if type_idx == 0: # Threshold
            val = rng.uniform(-1.0, 1.0)
            op = rng.choice([0, 1])
            return Threshold(feat, val, op)
            
        elif type_idx == 1: # Interval
            v1 = rng.uniform(-1.0, 1.0)
            v2 = rng.uniform(-1.0, 1.0)
            return Interval(feat, v1, v2)
            
        elif type_idx == 2: # Ratio
            feat2 = rng.integers(0, n_features)
            thresh = rng.uniform(0.1, 5.0)
            return Ratio(feat, feat2, thresh)
            
        elif type_idx == 3: # Gaussian
            mu = rng.uniform(-1.0, 1.0)
            sigma = rng.uniform(0.01, 1.0)
            return Gaussian(feat, mu, sigma)
            
    else:
        # Generic literal generation
        Choice = rng.choice(literal_types)
        
        if Choice == Variable:
            return Variable(rng.integers(0, n_features))
        elif Choice == Constant:
            # Mixture: 50% small (randn), 50% large (uniform -1000 to 1000)
            # AND: 50% chance to be OptimizableConstant (Ephemeral Constant)
            val = rng.standard_normal() if rng.random() < 0.5 else rng.uniform(-1000, 1000)
            
            if rng.random() < 0.5:
                return OptimizableConstant(val)
            else:
                return Constant(val)
        elif Choice in [Threshold, Interval, Ratio, Gaussian]:
             # Fallback to boolean logic above if mixed, but for now simple recursion
             return generate_random_literal(n_features, rng, literal_types=None)
        else:
            # Fallback or error
            return Constant(0.0)

def generate_random_tree(n_features, rng, max_depth=4, current_depth=0, 
                        binary_ops=None, unary_ops=None, ternary_ops=None, literal_types=None,
                        method='grow'):
    """
    Recursively generates a random tree using 'grow' or 'full' method.
    - 'grow': nodes can be terminals or operators until max_depth.
    - 'full': nodes are operators until max_depth-1, then terminals.
    """
    
    # Defaults for backward compatibility (Boolean Logic)
    if binary_ops is None: binary_ops = [AND, OR]
    if unary_ops is None: unary_ops = [NOT]
    
    # Terminal condition
    if current_depth >= max_depth:
        return generate_random_literal(n_features, rng, literal_types)
    
    # In 'grow' method, we can stop early and pick a literal
    if method == 'grow' and current_depth > 0:
        # Probability of stopping increases with depth?
        # Standard: 30% chance or more
        stop_prob = 0.3 + 0.1 * current_depth
        if rng.random() < stop_prob:
            return generate_random_literal(n_features, rng, literal_types)
        
    # Operator
    # Decide between Ternary, Undary, Binary
    # Distribution?
    ops_pool = []
    if binary_ops: ops_pool.append('binary')
    if unary_ops: ops_pool.append('unary')
    if ternary_ops: ops_pool.append('ternary')
    
    if not ops_pool:
         return generate_random_literal(n_features, rng, literal_types)
         
    op_type = rng.choice(ops_pool)

    if op_type == 'unary':
        op_cls = rng.choice(unary_ops)
        child = generate_random_tree(n_features, rng, max_depth, current_depth + 1,
                                    binary_ops, unary_ops, ternary_ops, literal_types)
        return op_cls(child)
    elif op_type == 'binary':
        op_cls = rng.choice(binary_ops)
        left = generate_random_tree(n_features, rng, max_depth, current_depth + 1,
                                   binary_ops, unary_ops, ternary_ops, literal_types)
        right = generate_random_tree(n_features, rng, max_depth, current_depth + 1,
                                    binary_ops, unary_ops, ternary_ops, literal_types)
        return op_cls(left, right)
    elif op_type == 'ternary':
        op_cls = rng.choice(ternary_ops)
        # SoftIf specific: child1 (cond) usually wants [0,1] range.
        # But for general GP, we just generate random trees.
        # Ideally, child1 is a Threshold/Interval/Ratio OR a subtree restricted to that?
        # For simplicity, generate full trees, and rely on SoftIf.evaluate using standard sigmoid logic on whatever input.
        c1 = generate_random_tree(n_features, rng, max_depth, current_depth + 1,
                                 binary_ops, unary_ops, ternary_ops, literal_types, method)
        c2 = generate_random_tree(n_features, rng, max_depth, current_depth + 1,
                                 binary_ops, unary_ops, ternary_ops, literal_types, method)
        c3 = generate_random_tree(n_features, rng, max_depth, current_depth + 1,
                                 binary_ops, unary_ops, ternary_ops, literal_types, method)
        return op_cls(c1, c2, c3)


# Operator Families for Smooth Substitution
BINARY_FAMILIES = {
    Add: [Sub], Sub: [Add],
    Mul: [Div], Div: [Mul],
    Min: [Max], Max: [Min],
    AND: [OR], OR: [AND]
}

UNARY_FAMILIES = {
    Sin: [Cos], Cos: [Sin],
    Log: [Sqrt], Sqrt: [Log],
    Square: [Abs], Abs: [Square]
}

class SymbolicMutator:
    def __init__(self, n_features, mutation_rate=0.1, 
                 binary_ops=None, unary_ops=None, ternary_ops=None, literal_types=None,
                 scheduler=None):
        self.n_features = n_features
        self.mutation_rate = mutation_rate
        
        # Defaults
        self.binary_ops = binary_ops if binary_ops else [Add, Sub, Mul, Div]
        self.unary_ops = unary_ops if unary_ops else [Sin, Cos, Log, Exp, Sqrt]
        self.ternary_ops = ternary_ops # Optional
        self.literal_types = literal_types # None means default Boolean literals
        
        self.scheduler = scheduler if scheduler else MutationScheduler(strategy='linear')
        
    def mutate(self, root: Node, rng, generation=None, total_generations=None, 
               X_semantic=None, max_semantic_delta=None):
        """
        Applies one of the ergodic mutations to the root with optional semantic filtering.
        Returns a NEW root.
        """
        # 1. Probabilistic Gate
        if self.mutation_rate < 1.0 and rng.random() > self.mutation_rate:
            return root.clone()

        parent = root.clone()
        progress = 0.0
        if generation is not None and total_generations is not None and total_generations > 0:
            progress = min(1.0, generation / total_generations)

        child = parent # Fallback
        
        # 2. Semantic Retry Loop
        # We try to find a mutation that respects the semantic threshold.
        for attempt in range(10): # Increased retries
            current_child = self._mutate_once(parent, rng, progress)
            child = current_child # Update fallback to most recent
            
            if X_semantic is None or max_semantic_delta is None:
                return child

            try:
                # Measure semantic distance
                yp = parent.evaluate(X_semantic)
                yc = child.evaluate(X_semantic)
                
                # Filter out NaNs/Infs
                valid = np.isfinite(yp) & np.isfinite(yc)
                if not np.any(valid): 
                    continue
                    
                delta = np.mean((yp[valid] - yc[valid])**2)
                
                # If it's a structural change, we might want to be more permissive?
                # For now, stick to the user's limit but allow the jump if it's "close enough"
                if delta <= max_semantic_delta:
                    return child
            except Exception:
                continue
        
        # If we reached here, it means we couldn't find a child within delta.
        # We return the LAST generated child anyway to ensure the GA progresses.
        # This prevents the population from stagnating on parents that have no "soft" children.
        return child


    def _mutate_once(self, root, rng, progress):
        """Internal single-mutation dispatch."""
        strategies = ['M1', 'M2', 'M3', 'M4', 'M5', 'M6', 'M7', 'M8', 'M9', 'M10']
        weights = self.scheduler.get_weights(progress)
        
        # Ensure weights match strategy count
        if len(weights) < len(strategies):
            # Pad with zeros or use baseline
            temp = np.zeros(len(strategies))
            temp[:len(weights)] = weights
            weights = temp / (temp.sum() + 1e-9)

        choice = rng.choice(strategies, p=weights)
        
        # 0. Cloned root for modification
        m_root = root.clone()

        if choice == 'M1':
            return self.mutate_constant_perturbation(m_root, rng)
        elif choice == 'M2':
            return self.mutate_insert(m_root, rng)
        elif choice == 'M3':
            return self.mutate_delete_shrink(m_root, rng)
        elif choice == 'M4':
            return self.mutate_reorder(m_root, rng)
        elif choice == 'M5':
            return self.mutate_subtree_replacement(m_root, rng)
        elif choice == 'M6':
            return self.mutate_operator_substitution_local(m_root, rng)
        elif choice == 'M7':
            return self.mutate_hoist(m_root, rng)
        elif choice == 'M8':
            return self.mutate_point(m_root, rng)
        elif choice == 'M9':
            return self.mutate_affine_wrap(m_root, rng)
        elif choice == 'M10':
            return self.mutate_residual(m_root, rng)
            
        return m_root

    def mutate_constant_perturbation(self, root, rng, sigma=0.1):
        """M1: Smooth perturbation of constants instead of atomic replacement."""
        indices = []
        for i in range(root.size):
            node = root.get_subtree(i)
            if isinstance(node, (Constant, OptimizableConstant)):
                indices.append(i)
        
        if not indices:
            return root
            
        chosen_idx = rng.choice(indices)
        node = root.get_subtree(chosen_idx)
        
        old_val = float(node.value)
        # Scale perturbation by value magnitude to handle different scales
        noise = rng.standard_normal() * sigma * (1.0 + abs(old_val))
        new_val = old_val + noise
        
        # Always return as OptimizableConstant to allow subsequent fine-tuning
        return root.set_subtree(chosen_idx, OptimizableConstant(new_val))

    def mutate_insert(self, root, rng):
        """M2: Structural Insertion. Depth-limited to control bloat."""
        size = root.size
        # Pick site. Prefer deeper sites for local insertion?
        target_idx = rng.integers(0, size)
        target_node = root.get_subtree(target_idx)
        
        # Total depth check (estimated)
        if root.depth > 10: 
            return root
            
        # For M2, we wrap the target_node in a new operator with a random literal
        new_literal = generate_random_literal(self.n_features, rng, self.literal_types)
        
        ops_pool = []
        if self.binary_ops: ops_pool.append('binary')
        if self.unary_ops: ops_pool.append('unary')
        if self.ternary_ops: ops_pool.append('ternary')
        
        if not ops_pool: return root
        op_type = rng.choice(ops_pool)
        
        if op_type == 'binary':
            op_cls = rng.choice(self.binary_ops)
            if rng.random() < 0.5:
                new_subtree = op_cls(target_node, new_literal)
            else:
                new_subtree = op_cls(new_literal, target_node)
        elif op_type == 'unary':
            op_cls = rng.choice(self.unary_ops)
            new_subtree = op_cls(target_node)
        else: # ternary
            op_cls = rng.choice(self.ternary_ops)
            l1 = generate_random_literal(self.n_features, rng, self.literal_types)
            l2 = generate_random_literal(self.n_features, rng, self.literal_types)
            pos = rng.integers(0, 3)
            # Use a dictionary mapping for cleaner placement
            slots = [l1, l2]
            slots.insert(pos, target_node)
            new_subtree = op_cls(*slots)
            
        return root.set_subtree(target_idx, new_subtree)

    def mutate_delete_shrink(self, root, rng):
        """M3: Structural Deletion (Pruning). Supports Unary, Binary, and Ternary."""
        op_indices = []
        for i in range(root.size):
            node = root.get_subtree(i)
            # Literal cannot be deleted (would leave root empty or invalid)
            if not isinstance(node, Literal):
                op_indices.append(i)
                
        if not op_indices:
            return root 
            
        chosen_idx = rng.choice(op_indices)
        node = root.get_subtree(chosen_idx)
        
        # Pick child to keep
        if isinstance(node, UnaryOperator):
            child_to_keep = node.child
        elif isinstance(node, BinaryOperator):
            child_to_keep = node.left if rng.random() < 0.5 else node.right
        elif isinstance(node, TernaryOperator):
            # Keep one of the non-condition branches (usually child2 or child3)
            child_to_keep = rng.choice([node.child2, node.child3])
        else:
            return root
                
        return root.set_subtree(chosen_idx, child_to_keep)

    def mutate_reorder(self, root, rng):
        """M4: Commutative Reordering. Immutable rebuild."""
        comm_indices = []
        for i in range(root.size):
            node = root.get_subtree(i)
            # Add, Mul, Min, Max, AND, OR are commutative
            if isinstance(node, (Add, Mul, Min, Max, AND, OR, SoftIf)):
                comm_indices.append(i)
                
        if not comm_indices:
            return root
            
        chosen_idx = rng.choice(comm_indices)
        node = root.get_subtree(chosen_idx)
        
        if isinstance(node, SoftIf):
            # Swap true/false branches
            new_node = SoftIf(node.child1, node.child3, node.child2)
        else:
            # Binary swap
            new_node = node.__class__(node.right, node.left)
        
        return root.set_subtree(chosen_idx, new_node)

    def mutate_subtree_replacement(self, root, rng):
        """M5: Subtree Replacement. Restricted depth to stay 'local'."""
        chosen_idx = rng.integers(0, root.size)
        new_subtree = generate_random_tree(
            self.n_features, rng, max_depth=2, # Keep it small
            binary_ops=self.binary_ops, 
            unary_ops=self.unary_ops,
            ternary_ops=self.ternary_ops,
            literal_types=self.literal_types
        )
        return root.set_subtree(chosen_idx, new_subtree)

    def mutate_operator_substitution_local(self, root, rng):
        """M6: Operator Substitution. Focuses on similar operators (families)."""
        op_indices = []
        for i in range(root.size):
            node = root.get_subtree(i)
            if isinstance(node, (BinaryOperator, UnaryOperator)):
                op_indices.append(i)
        
        if not op_indices:
            return root
            
        chosen_idx = rng.choice(op_indices)
        node = root.get_subtree(chosen_idx)
        
        # 1. Try Family Substitution first (smoother)
        if node.__class__ in BINARY_FAMILIES and rng.random() < 0.7:
             new_cls = rng.choice(BINARY_FAMILIES[node.__class__])
             return root.set_subtree(chosen_idx, new_cls(node.left, node.right))
        if node.__class__ in UNARY_FAMILIES and rng.random() < 0.7:
             new_cls = rng.choice(UNARY_FAMILIES[node.__class__])
             return root.set_subtree(chosen_idx, new_cls(node.child))

        # 2. General Substitution (prevent no-op)
        if isinstance(node, BinaryOperator):
            pool = [op for op in self.binary_ops if op is not node.__class__]
            if pool:
                return root.set_subtree(chosen_idx, rng.choice(pool)(node.left, node.right))
        elif isinstance(node, UnaryOperator):
            pool = [op for op in self.unary_ops if op is not node.__class__]
            if pool:
                return root.set_subtree(chosen_idx, rng.choice(pool)(node.child))
             
        return root

    def mutate_hoist(self, root, rng):
        """M7: Hoist Mutation. Replaces a node with one of its subtrees."""
        op_indices = [i for i in range(root.size) if not isinstance(root.get_subtree(i), Literal)]
        if not op_indices: return root
            
        chosen_idx = rng.choice(op_indices)
        op_node = root.get_subtree(chosen_idx)
        if op_node.size <= 1: return root
            
        sub_idx = rng.integers(1, op_node.size)
        hoisted_subtree = op_node.get_subtree(sub_idx).clone()
        return root.set_subtree(chosen_idx, hoisted_subtree)

    def mutate_point(self, root, rng):
        """M8: Point Mutation. Correctly mutates the node at the chosen site."""
        idx = rng.integers(0, root.size)
        node = root.get_subtree(idx)
        
        if isinstance(node, (BinaryOperator, UnaryOperator)):
            # Force a change at exactly 'idx'
            # We use a localized version of operator substitution
            if isinstance(node, BinaryOperator):
                pool = [op for op in self.binary_ops if op is not node.__class__]
                if pool: return root.set_subtree(idx, rng.choice(pool)(node.left, node.right))
            else:
                pool = [op for op in self.unary_ops if op is not node.__class__]
                if pool: return root.set_subtree(idx, rng.choice(pool)(node.child))
        elif isinstance(node, Literal):
            # Atomic replacement of literal
            return root.set_subtree(idx, generate_random_literal(self.n_features, rng, self.literal_types))
            
        return root

    def mutate_affine_wrap(self, root, rng, sigma=0.05):
        """M9: Affine Wrapper (f -> a*f + b). Extremely smooth refinement."""
        a = 1.0 + rng.standard_normal() * sigma
        b = rng.standard_normal() * sigma
        # Use OptimizableConstant so that BFGS can further tune these parameters
        return Add(Mul(OptimizableConstant(a), root), OptimizableConstant(b))

    def mutate_residual(self, root, rng, sigma=0.02):
        """M10: Residual Correction (f -> f + eps * small_h). Local additive update."""
        small_h = generate_random_tree(
            self.n_features, rng, max_depth=2,
            binary_ops=self.binary_ops, unary_ops=self.unary_ops,
            literal_types=self.literal_types
        )
        eps = rng.standard_normal() * sigma
        return Add(root, Mul(OptimizableConstant(eps), small_h))

