
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, 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)


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 [AND, OR]
        self.unary_ops = unary_ops if unary_ops else [NOT]
        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):
        """
        Applies one of the ergodic mutations to the root.
        Returns a NEW root (or same if no mutation happened).
        
        Args:
            generation (int): Current generation number (for adaptive control).
            total_generations (int): Total max generations.
        """
        if rng.random() > 1.0: 
            pass
            
        root = root.clone()
        
        # Pick strategy
        strategies = ['M1', 'M2', 'M3', 'M4', 'M5', 'M6']
        
        # Adaptive Weights
        progress = None
        if generation is not None and total_generations is not None and total_generations > 0:
            progress = generation / total_generations
            
        weights = self.scheduler.get_weights(progress)
        
        choice = rng.choice(strategies, p=weights)
        
        if choice == 'M1':
            return self.mutate_literal(root, rng)
        elif choice == 'M2':
            return self.mutate_insert(root, rng)
        elif choice == 'M3':
            return self.mutate_delete(root, rng)
        elif choice == 'M4':
            return self.mutate_reorder(root, rng)
        elif choice == 'M5':
            return self.mutate_subtree_replacement(root, rng)
        elif choice == 'M6':
            return self.mutate_operator_substitution(root, rng)
            
        return root

    def mutate_literal(self, root, rng):
        """M1: Atomic Replacement OR Perturbation of a random Literal."""
        size = root.size
        
        literal_indices = []
        for i in range(root.size):
            node = root.get_subtree(i)
            if isinstance(node, Literal):
                literal_indices.append(i)
        
        if not literal_indices:
            return root
            
        chosen_idx = rng.choice(literal_indices)
        node_to_mod = root.get_subtree(chosen_idx)
        
        # Check if it is a Constant -> Perturb!
        if isinstance(node_to_mod, Constant):
            if rng.random() < 0.7:
                 # "Unfreeze" Constant: Convert back to OptimizableConstant
                 # This allows BFGS to re-optimize it in the next evaluation
                 val = node_to_mod.value
                 new_literal = OptimizableConstant(val)
                 return root.set_subtree(chosen_idx, new_literal)
                 
        # Special Mutation for Threshold/Interval/Ratio: Perturb params
        if isinstance(node_to_mod, (Threshold, Interval, Ratio, Gaussian)):
            # Perturb existing params instead of full replacement
            # 50% chance to perturb, 50% replace
            if rng.random() < 0.5:
                # Need to clone or careful with mutation. 
                # Since grammar is mutable-style but we want functional purity, we construct new.
                # Just replace with new random of SAME type for now, maybe with params close to old?
                # Implementing simple perturbation:
                
                if isinstance(node_to_mod, Threshold):
                    delta = rng.uniform(-0.1, 0.1)
                    new_val = np.clip(node_to_mod.val + delta, 0.0, 1.0)
                    return root.set_subtree(chosen_idx, Threshold(node_to_mod.feature, new_val, node_to_mod.op))
                    
                elif isinstance(node_to_mod, Interval):
                    d1 = rng.uniform(-0.1, 0.1)
                    d2 = rng.uniform(-0.1, 0.1)
                    v1 = np.clip(node_to_mod.low + d1, 0.0, 1.0)
                    v2 = np.clip(node_to_mod.high + d2, 0.0, 1.0)
                    return root.set_subtree(chosen_idx, Interval(node_to_mod.feature, v1, v2))
                    
        # Default: Full Replacement
        new_lit = generate_random_literal(self.n_features, rng, self.literal_types)
        return root.set_subtree(chosen_idx, new_lit)

    def mutate_insert(self, root, rng):
        """M2: Structural Insertion (Growth)."""
        size = root.size
        target_idx = rng.integers(0, size)
        target_node = root.get_subtree(target_idx)
        
        if target_node.depth > 6: # Depth limit
            return root
            
        new_literal = generate_random_literal(self.n_features, rng, self.literal_types)
        
        # Pick Operator Type
        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)
            # Randomly decide if N is left or right child
            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) # Wrap current node
        elif op_type == 'ternary':
            op_cls = rng.choice(self.ternary_ops)
            # SoftIf(cond, true, false). Wrap target_node in one of them?
            # Or make target_node the condition?
            # Let's put target_node in 'true' branch for continuity, generate others.
            l1 = generate_random_literal(self.n_features, rng, self.literal_types)
            l2 = generate_random_literal(self.n_features, rng, self.literal_types)
            # 3 slots: cond, true, false
            # Randomly place target
            # Ideally for SoftIf: Cond should be specific. But staying generic.
            pos = rng.integers(0, 3)
            if pos == 0:
                new_subtree = op_cls(target_node, l1, l2)
            elif pos == 1:
                new_subtree = op_cls(l1, target_node, l2)
            else:
                new_subtree = op_cls(l1, l2, target_node)
            
        return root.set_subtree(target_idx, new_subtree)

    def mutate_delete(self, root, rng):
        """M3: Structural Deletion (Pruning)."""
        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)
        op_node = root.get_subtree(chosen_idx)
        
        # Pick child to keep
        if isinstance(op_node, UnaryOperator):
            child_to_keep = op_node.child
        else:
            # Binary
            if rng.random() < 0.5:
                child_to_keep = op_node.left
            else:
                child_to_keep = op_node.right
                
        return root.set_subtree(chosen_idx, child_to_keep)

    def mutate_reorder(self, root, rng):
        """M4: Reordering (Commutativity)."""
        # Restrict to Commutative Operators only: Add, Mul, Min, Max, AND, OR
        # This prevents semantic destruction for Sub, Div, etc.
        
        comm_indices = []
        for i in range(root.size):
            node = root.get_subtree(i)
            # Strict check for allowed commutative operators OR SoftIf for branch swap
            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)
        
        # Swap children
        if isinstance(node, SoftIf):
            # SoftIf(cond, true, false). Swap true/false branches.
            temp = node.child2
            node.child2 = node.child3
            node.child3 = temp
        else:
            temp = node.left
            node.left = node.right
            node.right = temp
        
        return root

    def mutate_subtree_replacement(self, root, rng):
        """M5: Subtree Replacement."""
        # Pick any node
        total_size = root.size
        chosen_idx = rng.integers(0, total_size)
        
        # Determine max depth for new branch relative to current depth?
        # For simplicity, generate a small random tree
        # Limit depth to avoid bloat (e.g. depth=2 or 3)
        new_subtree = generate_random_tree(
            self.n_features, rng, max_depth=3, 
            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(self, root, rng):
        """M6: Operator Substitution."""
        # Pick an operator node
        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)
        
        new_op_node = None
        
        if isinstance(node, BinaryOperator):
            if not self.binary_ops: return root
            # Pick different logic if feasible, or just random
            op_cls = rng.choice(self.binary_ops)
            new_op_node = op_cls(node.left, node.right)
            
        elif isinstance(node, UnaryOperator):
            if not self.unary_ops: return root
            op_cls = rng.choice(self.unary_ops)
            new_op_node = op_cls(node.child)
            
        if new_op_node:
             return root.set_subtree(chosen_idx, new_op_node)
             
        return root
