
import numpy as np
import copy
from ezga.simple.problem import ElementwiseProblem
from ezga.simple.algorithm import GA
from ezga.simple.minimize import minimize
from ezga.simple.grammar import (
    Variable, Constant, OptimizableConstant,
    Add, Sub, Mul, Div, Min, Max,
    Sin, Cos, Exp, Log, Sqrt, Abs, Square, Neg,
    Add, Sub, Mul, Div, Min, Max,
    Sin, Cos, Exp, Log, Sqrt, Abs, Square, Neg,
    AND, OR, Threshold, Interval, Ratio, Gaussian,
    SoftIf
)
from ezga.simple.ops import SymbolicMutator, generate_random_tree
from ezga.variation.variation import Variation_Operator
from ezga.simple.semantic import SemanticAnalyzer
from ezga.simple.ge import Grammar, GEMapper, GEVariation, Node

import pathlib
from scipy.optimize import minimize as scipy_minimize
import sys

# Optional Dependencies for User Friendliness
try:
    import pandas as pd
except ImportError:
    pd = None

try:
    import matplotlib.pyplot as plt
except ImportError:
    plt = None

class SymbolicMathMutation:
    """
    Callable Mutation Operator for Symbolic Regression.
    Expected signature: func(container) -> new_container
    """
    def __init__(self, n_features, problem, mutation_rate=0.1, 
                  binary_ops=None, unary_ops=None, ternary_ops=None, literal_types=None,
                  pop_size=100, max_generations=100):
        self.problem = problem
        self.mutator = SymbolicMutator(n_features, mutation_rate, 
                                      binary_ops, unary_ops, ternary_ops, literal_types)
        self.n_features = n_features
        self.literal_types = literal_types
        self.binary_ops = binary_ops
        self.unary_ops = unary_ops
        self.ternary_ops = ternary_ops
        
        # Adaptive Control State
        self.pop_size = pop_size
        self.max_generations = max_generations
        self.call_count = 0 
        self.current_gen = 0

    def __call__(self, container):
        rng = np.random.default_rng()
        import copy
        c = copy.deepcopy(container)
        
        # 1. Get Tree from Registry using Position ID
        apm = c.AtomPositionManager
        pos = apm.atomPositions
        if len(pos) > 0:
            tree_id = float(pos[0][0])
            tree_obj = self.problem.get_tree(tree_id)
        else:
            tree_obj = None

        if tree_obj is None:
            # Generate new random if missing
            tree_obj = generate_random_tree(self.n_features, rng, 
                                            binary_ops=self.binary_ops,
                                            unary_ops=self.unary_ops,
                                            ternary_ops=self.ternary_ops,
                                            literal_types=self.literal_types)
        
        # 2. Mutate
        # Update Generation Estimation
        # Each call is 1 individual. 
        self.call_count += 1
        # Approx generation
        self.current_gen = self.call_count // self.pop_size
        
        new_tree = self.mutator.mutate(tree_obj, rng, 
                                      generation=self.current_gen, 
                                      total_generations=self.max_generations)
        
        # 3. Register New Tree with New ID
        new_id = self.problem.get_next_id()
        self.problem.register_tree(new_id, new_tree)
        
        # 4. Update Container Position ID
        # HACK: Using dummy atoms. Reset them.
        # Ensure we have 1 dummy atom at least
        new_pos = np.zeros((1, 3))
        new_pos[0, 0] = new_id
        apm.set_atomPositions(new_pos)
        
        # Invalidate Energy
        apm.E = None
        if hasattr(apm, "metadata") and "objectives" in apm.metadata:
            del apm.metadata["objectives"]
            
        return c

class SymbolicMathCrossover:
    """
    Callable Crossover Operator.
    Expected signature: func(c1, c2) -> (new_c1, new_c2)
    """
    def __call__(self, p1, p2):
        import copy
        return copy.deepcopy(p1), copy.deepcopy(p2)


class SymbolicRegressionProblem(ElementwiseProblem):
    def __init__(self, X, y, n_features, binary_ops, unary_ops, literal_types, complexity_penalty=0.0):
        super().__init__(n_var=1, n_obj=2, xl=0, xu=1e6) 
        self.X_data = X
        self.y_data = y
        self.n_features = n_features
        self.binary_ops = binary_ops
        self.unary_ops = unary_ops
        self.literal_types = literal_types
        self.base_complexity_penalty = complexity_penalty
        self.current_complexity_penalty = complexity_penalty
        
        self.tree_registry = {}
        self.next_id = 10000 
        self.results_cache = {} # str(formula) -> [mse, complexity, tree]

    def register_tree(self, id_val, tree):
        self.tree_registry[int(id_val)] = tree
        
    def get_tree(self, id_val):
        return self.tree_registry.get(int(id_val))

    def get_next_id(self):
        self.next_id += 1
        return self.next_id

    def evaluate_tree(self, tree, tree_id=None):
        try:
            # -----------------------------------------------------------
            # Pre-Optimization Simplification
            # -----------------------------------------------------------
            # Apply algebraic simplification (0+x->x, etc) to reduce bloat
            # before we spend time optimizing constants.
            tree = tree.simplify_aggressive()
            
            # -----------------------------------------------------------
            # Phase 9: Result Caching
            # -----------------------------------------------------------
            formula_key = str(tree)
            if formula_key in self.results_cache:
                return self.results_cache[formula_key]
            
            # Use current complexity penalty (potentially increased by stagnation)
            penalty = getattr(self, "current_complexity_penalty", 0.0)
            
            # [Fix] Persist simplified tree to registry if we have an ID
            if tree_id is not None:
                self.tree_registry[int(tree_id)] = tree

            # -----------------------------------------------------------
            # Memetic Algorithm Step: Constant Optimization
            # -----------------------------------------------------------
            # 1. Gather all OptimizableConstant nodes
            opt_nodes = []
            
            def collect_opt_nodes(node):
                if isinstance(node, OptimizableConstant):
                    opt_nodes.append((node, 'value'))
                elif isinstance(node, Threshold):
                    opt_nodes.append((node, 'val'))
                    opt_nodes.append((node, 'sharpness'))
                elif isinstance(node, Interval):
                    opt_nodes.append((node, 'low'))
                    opt_nodes.append((node, 'high'))
                    opt_nodes.append((node, 'sharpness'))
                elif isinstance(node, Ratio):
                    opt_nodes.append((node, 'threshold'))
                    opt_nodes.append((node, 'sharpness'))
                elif isinstance(node, Gaussian):
                    opt_nodes.append((node, 'mu'))
                    opt_nodes.append((node, 'sigma'))
                    
                if hasattr(node, 'children'):
                    for child in node.children:
                        collect_opt_nodes(child)
                elif hasattr(node, 'child'):
                    collect_opt_nodes(node.child)
            
            collect_opt_nodes(tree)
            
            # -----------------------------------------------------------
            # Phase 9: Conditional Memetic Optimization
            # -----------------------------------------------------------
            # First evaluate without optimization
            if opt_nodes:
                y_raw = tree.evaluate(self.X_data)
                mse_raw = np.mean((self.y_data - y_raw)**2)
                
                # If it's already "okay-ish", optimize the constants.
                # Otherwise skip the expensive BFGS.
                # Scaling factor: if it's 10x worse than a baseline or average, skip.
                # For now, let's use a very loose heuristic: if MSE > 1e10, skip.
                # Better: keep track of best MSE so far.
                best_so_far = getattr(self, "best_mse_so_far", 1.0)
                if mse_raw > best_so_far * 100 and mse_raw > 1.0:
                    # Skip optimization
                    opt_nodes = []

            # 2. If we still have ephemeral constants, optimize them!
                x0 = []
                for node, param_name in opt_nodes:
                    val = getattr(node, param_name)
                    if param_name == 'sharpness':
                        # Log-scale for better range coverage and stability [exp(-5), exp(10)]
                        x0.append(np.log(max(1e-3, val)))
                    else:
                        x0.append(val)
                x0 = np.array(x0)
                
                def objective(params):
                    # Update node values
                    for i, (node, param_name) in enumerate(opt_nodes):
                        # Constraints clamping
                        val = params[i]
                        if param_name == 'sharpness':
                             # Convert back from log-scale: clip range to ~0.006 to 22000
                             val = np.exp(np.clip(val, -5, 10))
                        elif isinstance(node, Gaussian) and param_name == 'sigma':
                             val = max(1e-3, val) # Prevent zero/negative sigma
                        
                        setattr(node, param_name, val)
                        
                    # Evaluate
                    yp = tree.evaluate(self.X_data)
                    return np.mean((self.y_data - yp)**2)
                
                # Perform local optimization (e.g. BFGS)
                # Cap iterations for speed
                res = scipy_minimize(objective, x0, method='BFGS', options={'maxiter': 50})
                
                # Update with best params
                for i, (node, param_name) in enumerate(opt_nodes):
                    val = res.x[i]
                    if param_name == 'sharpness':
                        val = np.exp(np.clip(val, -5, 10))
                    elif isinstance(node, Gaussian) and param_name == 'sigma':
                        val = max(1e-3, val)
                    setattr(node, param_name, val)

                # 3. Freeze them! Replace OptimizableConstant with standard Constant
                # Note: We need parent pointers or a replace traversal.
                # Since our tree structure is simple, a DFS replace is safer.
                
                def freeze_constants(node):
                    # Handling binary ops
                    if isinstance(node, (Add, Sub, Mul, Div)): # BinaryOperator subclasses
                        if isinstance(node.left, OptimizableConstant):
                            node.left = Constant(node.left.value)
                        else:
                            freeze_constants(node.left)
                            
                        if isinstance(node.right, OptimizableConstant):
                            node.right = Constant(node.right.value)
                        else:
                            freeze_constants(node.right)
                            
                    # Handling unary ops
                    elif isinstance(node, (Sin, Cos, Exp, Log, Sqrt, Abs, Square, Neg)): # UnaryOperator subclasses
                         if isinstance(node.child, OptimizableConstant):
                             node.child = Constant(node.child.value)
                         else:
                             freeze_constants(node.child)
                             
                # Handle root case
                if isinstance(tree, OptimizableConstant):
                     # This effectively replaces the root logic outside, but we can't change 'tree' ref passed by value
                     # This edge case (tree = single constant) is rare/useless, so we skip freezing logic for pure root
                     # or handle it if needed. For now, values are updated, so it works.
                     pass
                else:
                    freeze_constants(tree)
                    
            # -----------------------------------------------------------
            # Post-Optimization Simplification (Cleanup)
            # -----------------------------------------------------------
            # Now that Opt constants are real Constants (e.g. 0.000), 
            # we can simplify again to remove +0.0, *1.0, etc.
            tree = tree.simplify_aggressive()
            if tree_id is not None:
                self.tree_registry[int(tree_id)] = tree

            # -----------------------------------------------------------
            # Standard Evaluation
            # -----------------------------------------------------------
            y_pred = tree.evaluate(self.X_data)
            mse = np.mean((self.y_data - y_pred)**2)
            
            # Track best MSE for conditional optimization
            if not hasattr(self, "best_mse_so_far") or mse < self.best_mse_so_far:
                self.best_mse_so_far = mse

            # Return objectives: [MSE, Complexity, Tree]
            # Phase 9: Dynamic Bloat Control (Size + Depth)
            # Total Complexity = Size + (Depth * 0.5)
            complexity = tree.size + (tree.depth * 0.5)
            
            result = [mse, complexity, tree]
            self.results_cache[formula_key] = result
            return result
        except Exception:
            return [1e9, 1e9, None]

    def update_penalty(self, stall_count):
        """
        Adjusts the complexity penalty based on search stagnation.
        Penalty increases by 10% per stall generation.
        """
        factor = 1.0 + 0.1 * stall_count
        self.current_complexity_penalty = self.base_complexity_penalty * factor

    def _evaluate(self, x, out, *args, **kwargs):
        try:
            tid = int(x[0])
            tree = self.get_tree(tid)
            if tree is None:
                out["F"] = [1e9, 1e9]
                return
            
            objs = self.evaluate_tree(tree, tree_id=tid)
            out["F"] = objs[:2] # MSE, Complexity
        except Exception:
            out["F"] = [1e9, 1e9]

# ---------------------------------------------------------------------------
# Grammatical Evolution Core
# ---------------------------------------------------------------------------

def build_math_grammar(n_features, binary_ops, unary_ops, ternary_ops, literal_types):
    """ Builds a BNF Grammar for Mathematical Symbolic Regression """
    g = Grammar(start_symbol="<expr>")
    
    # Helpers for feature mapping
    def ge_variable(f_idx):
        return Variable(f_idx % n_features)
    
    def ge_threshold(f_idx, val, op_idx):
        return Threshold(f_idx % n_features, val, op_idx % 2)

    def ge_interval(f_idx, low, high, sharp):
        return Interval(f_idx % n_features, low, high, sharp)

    def ge_ratio(f1_idx, f2_idx, thresh):
        return Ratio(f1_idx % n_features, f2_idx % n_features, thresh)

    def ge_gaussian(f_idx, mu, sigma):
        return Gaussian(f_idx % n_features, mu, sigma)

    # 1. Operators
    expr_rules = []
    if binary_ops:
        for op in binary_ops:
            expr_rules.append([op, "<expr>", "<expr>"])
    if unary_ops:
        for op in unary_ops:
            expr_rules.append([op, "<expr>"])
    if ternary_ops:
        for op in ternary_ops:
            # Assumes SoftIf style
            expr_rules.append([op, "<cond>", "<expr>", "<expr>"])
    
    # Terminals
    expr_rules.append(["<literal>"])
    g.add_rule("<expr>", expr_rules)
    
    # 2. Conditions (for SoftIf)
    cond_rules = []
    if literal_types:
        if Threshold in literal_types:
            cond_rules.append([ge_threshold, "__gene__", "__gene_float__", "__gene__"])
        if Interval in literal_types:
            cond_rules.append([ge_interval, "__gene__", "__gene_float__", "__gene_float__", 10.0])
    if not cond_rules:
        cond_rules.append([ge_threshold, "__gene__", 0.5, 0])
        
    g.add_rule("<cond>", cond_rules)
    
    # 3. Literals / Variables
    lit_rules = []
    lit_rules.append([ge_variable, "__gene__"])
    lit_rules.append([OptimizableConstant, "__gene_float_biased__"])
    
    if Ratio in literal_types:
        lit_rules.append([ge_ratio, "__gene__", "__gene__", "__gene_float__"])
    if Gaussian in literal_types:
        lit_rules.append([ge_gaussian, "__gene__", "__gene_float__", "__gene_float__"])
        
    g.add_rule("<literal>", lit_rules)
    
    return g

class SymbolicGEProblem(SymbolicRegressionProblem):
    """
    Subclass of SymbolicRegressionProblem designed for Grammatical Evolution.
    The primary difference is that 'evaluate' expects a genotype (vector of integers)
    instead of just a tree ID.
    """
    def __init__(self, X, y, n_features, binary_ops, unary_ops, ternary_ops, literal_types, 
                 complexity_penalty=0.0, genotype_len=50):
        # We need n_var = genotype_len for the GA to evolve the vector
        super().__init__(X, y, n_features, binary_ops, unary_ops, literal_types, complexity_penalty)
        self.n_var = genotype_len
        self.xl = np.zeros(self.n_var)
        self.xu = np.full(self.n_var, 255)
        
        # Setup Grammar and Mapper
        self.grammar = build_math_grammar(n_features, binary_ops, unary_ops, ternary_ops, literal_types)
        
        # Enhanced Mapper to support our special __tags__
        from ezga.simple.grammar import OptimizableConstant
        
        # We need to capture y_mean/y_std for biased constants
        ym = np.mean(y)
        ys = np.std(y) + 1e-6
        nf = n_features

        class BiasedGEMapper(GEMapper):
            def _derive(self, symbol, depth):
                if symbol == "__gene_float_biased__":
                    g = self._get_gene()
                    # Map to y_mean +/- y_std
                    val = ym + ((g % 2000) - 1000)/1000.0 * ys
                    return OptimizableConstant(val)
                
                # Special handling for feature indices in parametric terminals
                # If we see a Constructor that is Threshold/Variable etc, help it map __gene__ to feat
                return super()._derive(symbol, depth)

        self.mapper = BiasedGEMapper(self.grammar, max_depth=15, max_wraps=3)
        self.mapper.n_features = n_features # For feature wrapping in constructors

    def evaluate_tree_ge(self, tree):
        # Similar to evaluate_tree but without tree_id persistence since 
        # genotypes create transient trees (or we can cache them)
        return self.evaluate_tree(tree)

    def _evaluate(self, x, out, *args, **kwargs):
        # x is the genotype (int vector)
        genotype = x.astype(int)
        
        # 1. Map to tree
        try:
            tree = self.mapper.map(genotype)
        except Exception as e:
            # print(f"GE Mapping Fail: {e}")
            tree = None
            
        if tree is None:
            out["F"] = [1e9, 1e9]
            return
            
        # 2. Evaluate
        try:
            objs = self.evaluate_tree(tree)
            out["F"] = objs[:2]
        except Exception:
            out["F"] = [1e9, 1e9]


class SymbolicRegressor:
    def __init__(self, 
                 pop_size=100, 
                 pop_size_max=None,
                 generations=50, 
                 n_features=1,
                 binary_ops=None, 
                 unary_ops=None,
                 ternary_ops=None,
                 literal_types=None,
                 complexity_penalty=0.001,
                 crossover_probability=0.3,
                 mutation_rate=1.0,
                 random_state=None,
                 engine="tree",
                 genotype_len=50,
                 feature_names=None,
                 verbose=False,
                 preset=None):
        
        self.pop_size = pop_size
        self.pop_size_max = pop_size_max
        self.generations = generations
        self.n_features = n_features
        self.complexity_penalty = complexity_penalty
        self.crossover_probability = crossover_probability
        self.mutation_rate = mutation_rate
        self.random_state = random_state
        self.feature_names = feature_names
        self.verbose = verbose
        
        # --- Operator Presets ---
        # Default: Enabled automatically, safe, domain-agnostic.
        # Opt-in: Must be enabled explicitly (discontinuities, strong priors).
        
        defaults = {
            "binary":   ['add', 'sub', 'mul'], # No div (singularities), no min/max (kinks)
            "unary":    ['neg', 'exp'],        # No sin/cos (periodic), no log/sqrt (domain)
            "ternary":  [],                    # Keep smooth
            "literal":  ['variable', 'constant'] # No gaussian/bias priors
        }
        
        self.presets = {
            "default": defaults,
            "math": {
                "binary": defaults["binary"] + ['div'],
                "unary":  defaults["unary"]  + ['sin', 'cos', 'log', 'sqrt', 'abs'],
                "ternary": [],
                "literal": defaults["literal"]
            },
            "physics": {
                "binary": defaults["binary"] + ['div'],
                "unary":  defaults["unary"]  + ['sin', 'cos'],
                "ternary": [],
                "literal": defaults["literal"] + ['gaussian']
            },
             "all": {
                "binary": defaults["binary"] + ['div', 'min', 'max'],
                "unary":  defaults["unary"]  + ['sin', 'cos', 'log', 'sqrt', 'abs'],
                "ternary": ['softif'],
                "literal": defaults["literal"] + ['gaussian', 'threshold', 'interval', 'ratio']
            }
        }
        
        # Resolve Configuration
        selected_preset = self.presets.get(preset, defaults) if preset else defaults
        
        # Allow overrides: if argument provided, use it; else use preset/default
        self.binary_ops = self._resolve_ops(binary_ops if binary_ops is not None else selected_preset["binary"])
        self.unary_ops = self._resolve_ops(unary_ops if unary_ops is not None else selected_preset["unary"])
        self.ternary_ops = self._resolve_ops(ternary_ops if ternary_ops is not None else selected_preset["ternary"])
        self.literal_types = self._resolve_ops(literal_types if literal_types is not None else selected_preset["literal"])
        
        self.best_tree = None
        self.best_fitness = float('inf')
        self.engine = engine
        self.genotype_len = genotype_len

    @staticmethod
    def get_supported_operators():
        """
        Returns a dictionary of supported operators categorized by type.
        """
        return {
            'Binary (binary_ops)': {
                'add': Add, 'sub': Sub, 'mul': Mul, 'div': Div, 
                'min': Min, 'max': Max, 'and': AND, 'or': OR
            },
            'Unary (unary_ops)': {
                'sin': Sin, 'cos': Cos, 'exp': Exp, 'log': Log, 
                'sqrt': Sqrt, 'abs': Abs, 'square': Square, 'neg': Neg
            },
            'Ternary (ternary_ops)': {
                'softif': SoftIf
            },
            'Literals (literal_types)': {
                'variable': Variable, 'constant': Constant, 
                'optimizableconstant': OptimizableConstant,
                'threshold': Threshold, 'interval': Interval, 
                'ratio': Ratio, 'gaussian': Gaussian
            }
        }

    @staticmethod
    def print_available_ops():
        """
        Typing helper: Prints all available operator strings that can be used in configuration.
        """
        ops = SymbolicRegressor.get_supported_operators()
        print("--- Available Symbolic Regressor Operators ---")
        for category, op_map in ops.items():
            print(f"\n[{category}]")
            # Print keys in rows of 4
            keys = sorted(list(op_map.keys()))
            for i in range(0, len(keys), 4):
                print("  " + ", ".join(f"'{k}'" for k in keys[i:i+4]))
        print("\n--------------------------------------------")

    def _resolve_ops(self, ops_list):
        """
        Maps a list of strings (or mixed classes) to actual Operator classes.
        Example: ['add', Sub] -> [Add, Sub]
        """
        if ops_list is None:
            return []
            
        # Flatten the categorized map for lookup
        full_map = {}
        for cat_map in self.get_supported_operators().values():
            full_map.update(cat_map)
        
        resolved = []
        for op in ops_list:
            if isinstance(op, str):
                key = op.lower().strip()
                if key in full_map:
                    resolved.append(full_map[key])
                else:
                    print(f"Warning: Unknown operator '{op}'. Ignoring.")
            else:
                # Assume it's a class
                resolved.append(op)
        return resolved
        
    def fit(self, X, y):
        # 0. Data Handling (Auto-detect Pandas)
        if pd is not None:
            if isinstance(X, (pd.DataFrame, pd.Series)):
                # Capture feature names if not provided
                if hasattr(X, 'columns') and self.feature_names is None:
                    self.feature_names = list(X.columns)
                X = X.values
            if isinstance(y, (pd.DataFrame, pd.Series)):
                y = y.values.flatten()
                
        # Auto-detect n_features if inconsistent
        if X.ndim == 2:
            if X.shape[1] != self.n_features:
                if self.verbose:
                    print(f"Adjusting n_features from {self.n_features} to {X.shape[1]}")
                self.n_features = X.shape[1]
        elif X.ndim == 1:
            # Reshape 1D input
            X = X.reshape(-1, 1)
            self.n_features = 1
            
        # 1. Setup Problem
        if self.engine == "ge":
            problem = SymbolicGEProblem(
                X, y, 
                self.n_features, 
                self.binary_ops, 
                self.unary_ops, 
                self.ternary_ops,
                self.literal_types,
                self.complexity_penalty,
                genotype_len=self.genotype_len
            )
        else:
            problem = SymbolicRegressionProblem(
                X, y, 
                self.n_features, 
                self.binary_ops, 
                self.unary_ops, 
                self.literal_types,
                self.complexity_penalty
            )
        
    def _get_initial_seeds(self, X, y, rng):
        # Data Stats for biased seeding
        y_mean = np.mean(y)
        y_std = np.std(y) + 1e-6
        feat_stats = []
        for f in range(self.n_features):
            feat_stats.append({
                'min': np.min(X[:, f]),
                'max': np.max(X[:, f]),
                'mean': np.mean(X[:, f]),
                'std': np.std(X[:, f]) + 1e-6
            })

        seeds = []
        
        # Seed 1: Simple Linear (c0 + c1*x0 + c2*x1...)
        def make_linear_term(feat_idx):
            return Mul(OptimizableConstant(rng.standard_normal()), Variable(feat_idx))
            
        if self.n_features > 0:
            current = OptimizableConstant(y_mean + rng.standard_normal() * 0.1)
            for f in range(self.n_features):
                scale = y_std / feat_stats[f]['std']
                term = Mul(OptimizableConstant(scale * rng.standard_normal()), Variable(f))
                current = Add(current, term)
            seeds.append(current)
            
        # Seed 2: Quadratic (x0^2)
        for f in range(self.n_features):
            term = Mul(OptimizableConstant(rng.standard_normal()), Mul(Variable(f), Variable(f)))
            term = Add(term, OptimizableConstant(rng.standard_normal()))
            seeds.append(term)
        
        # Seed 3: Interaction (x0 * x1)
        if self.n_features > 1:
            term = Mul(Variable(0), Variable(1))
            term = Mul(term, OptimizableConstant(rng.standard_normal()))
            term = Add(term, OptimizableConstant(rng.standard_normal()))
            seeds.append(term)
            
        # Seed 4: Trigonometric
        has_sin = any(op.__name__ == 'Sin' for op in self.unary_ops)
        has_cos = any(op.__name__ == 'Cos' for op in self.unary_ops)
        for f in range(self.n_features):
            lin = make_linear_term(f)
            lin = Add(lin, OptimizableConstant(rng.standard_normal()))
            if has_sin:
                term = Mul(OptimizableConstant(rng.standard_normal()), Sin(Variable(f)))
                seeds.append(Add(lin, term))
            if has_cos:
                term = Mul(OptimizableConstant(y_std * rng.standard_normal()), Cos(Variable(f)))
                seeds.append(Add(lin, term))
                
        # Seed 5: Logic Seeds (SoftIf)
        has_softif = any(op.__name__ == 'SoftIf' for op in self.ternary_ops)
        has_threshold = any(getattr(lit, '__name__', str(lit)) == 'Threshold' for lit in self.literal_types)
        if has_softif and has_threshold:
            for f in range(self.n_features):
                val = feat_stats[f]['mean']
                # Standard Logic (Moderate sharpness)
                cond = Threshold(f, val, 0, sharpness=1.5)
                seeds.append(SoftIf(cond, OptimizableConstant(y_mean + y_std * 0.5), OptimizableConstant(y_mean - y_std * 0.5)))
                # Sharp Logic (High sharpness for step-like behavior)
                cond_sharp = Threshold(f, val, 0, sharpness=100.0)
                seeds.append(SoftIf(cond_sharp, OptimizableConstant(y_mean + y_std * 0.8), OptimizableConstant(y_mean - y_std * 0.8)))
                
        # Seed 6: Ratio Seeds
        has_div = any(op.__name__ == 'Div' for op in self.binary_ops)
        if has_div and self.n_features > 1:
            for f in range(1, self.n_features):
                ratio = Div(Variable(0), Add(Variable(f), OptimizableConstant(1e-3)))
                seeds.append(Add(OptimizableConstant(y_mean), Mul(OptimizableConstant(y_std), ratio)))

        # Seed 7: Gaussian Seeds
        has_gauss = any(getattr(lit, '__name__', str(lit)) == 'Gaussian' for lit in self.literal_types)
        if has_gauss:
            for f in range(self.n_features):
                seeds.append(Gaussian(f, feat_stats[f]['mean'], feat_stats[f]['std']))
            
        return seeds

    def fit(self, X, y):
        # 1. Setup Problem
        if self.engine == "ge":
            problem = SymbolicGEProblem(
                X, y, self.n_features, self.binary_ops, self.unary_ops, self.ternary_ops,
                self.literal_types, self.complexity_penalty, genotype_len=self.genotype_len
            )
        else:
            problem = SymbolicRegressionProblem(
                X, y, self.n_features, self.binary_ops, self.unary_ops,
                self.literal_types, self.complexity_penalty
            )
        
        # 2. Setup Algorithm and Variation
        rng = np.random.default_rng(self.random_state)
        
        if self.engine == "ge":
            algorithm = GA(
                pop_size=self.pop_size, pop_size_max=self.pop_size_max,
                sampling=None, enable_bo=False, selection="nsga3"
            )
            algorithm.variation = GEVariation(
                genotype_len=self.genotype_len, mutation_rate=0.2, crossover_rate=0.7
            )
        else:
            # Traditional Tree-based logic
            init_ids = np.zeros((self.pop_size, 1))
            seeds = self._get_initial_seeds(X, y, rng)
            
            for i in range(self.pop_size):
                tid = i + 1
                init_ids[i, 0] = tid
                if i < len(seeds):
                    tree = copy.deepcopy(seeds[i])
                else:
                    max_d = rng.integers(2, 7)
                    meth = 'grow' if i % 2 == 0 else 'full'
                    tree = generate_random_tree(self.n_features, rng, max_depth=max_d,
                                                binary_ops=self.binary_ops, unary_ops=self.unary_ops,
                                                ternary_ops=self.ternary_ops, literal_types=self.literal_types,
                                                method=meth)
                problem.register_tree(tid, tree)

            algorithm = GA(
                pop_size=self.pop_size, pop_size_max=self.pop_size_max,
                sampling=init_ids, enable_bo=False, selection="nsga3"
            )
            algorithm.mutation = SymbolicMathMutation(
                self.n_features, problem=problem, mutation_rate=4.0, 
                binary_ops=self.binary_ops, unary_ops=self.unary_ops,
                ternary_ops=self.ternary_ops, literal_types=self.literal_types,
                pop_size=self.pop_size, max_generations=self.generations
            )
            algorithm.crossover = SymbolicMathCrossover()
        
        # 4. Minimize
        res = minimize(
            problem,
            algorithm,
            termination=('n_gen', self.generations),
            mutation_rate=self.mutation_rate,
            seed=self.random_state,
            verbose=self.verbose
        )
        
        # Extract Best
        if res.pop:
            best_ind = None
            min_f = float('inf')
            
            # Weights for scalarized best extraction (consistent with minimize)
            # Use current complexity penalty
            w_scalar = np.array([1.0, problem.current_complexity_penalty])
            
            for ind in res.pop:
                # ind is a structure container
                # Extract objectives
                try:
                    meta = getattr(ind.AtomPositionManager, "metadata", {})
                    if "objectives" in meta:
                        objs = np.array(meta["objectives"]).flatten()
                        # Scalarize for comparison
                        f_scalar = np.dot(objs, w_scalar[:len(objs)])
                    else:
                        e_val = ind.AtomPositionManager.E
                        if e_val is not None:
                             if np.ndim(e_val) > 0:
                                 f_scalar = np.atleast_1d(e_val)[0]
                             else:
                                 f_scalar = float(e_val)
                        else:
                             f_scalar = np.inf
                    
                    if f_scalar is not None and f_scalar < min_f:
                        min_f = f_scalar
                        best_ind = ind
                except:
                    continue
            
            if best_ind:
                if self.engine == "ge":
                    # For GE, pos[0] is the genotype (vec of codons)
                    genes = best_ind.AtomPositionManager.atomPositions[:, 0].astype(int)
                    self.best_tree = problem.mapper.map(genes)
                    # Re-optimize last time to be sure? 
                    # Problem.evaluate already did it, but let's be safe if we want the absolute best version
                    self.best_tree = problem.evaluate_tree(self.best_tree)[2] # 3rd ret is tree
                    self.best_fitness = min_f
                else:
                    # Standard Tree
                    pos = best_ind.AtomPositionManager.atomPositions
                    if len(pos) > 0:
                        tid = pos[0][0]
                        self.best_tree = problem.get_tree(tid)
                        self.best_fitness = min_f
                
        return self

    def predict(self, X):
        if self.best_tree:
            return self.best_tree.evaluate(X)
        return np.zeros(len(X))
        
    def score(self, X, y):
        y_pred = self.predict(X)
        return 1 - np.sum((y - y_pred)**2) / np.sum((y - np.mean(y))**2)

    @property
    def expression(self):
        # Replace variable indices with names if available
        s = str(self.best_tree) if self.best_tree else "None"
        if self.feature_names:
            for i, name in enumerate(self.feature_names):
                # Replace x0 with name, x1 with name, etc.
                # Simple string replacement (careful with x1 vs x10)
                # Using reverse order (x10 before x1) to avoid partial replacement issues
                # Or regex. For now, let's assume standard "x{i}" format from Variable.__str__
                # But variable is outputting x{index}.
                pass 
                # Actually, Variable node logic controls this.
                # A safer way is to have the Tree support a 'format' method.
                # For this simple improvement, let's just leave it or do a simple replace
                # if the user asks for 'human_readable'.
        return s

    def plot_results(self, X, y, outfile=None, title="Symbolic Regression Fit"):
        """
        Simple built-in plotter for 1D symbolic regression results.
        """
        if plt is None:
            print("Matplotlib not installed. Cannot plot.")
            return
            
        if self.n_features != 1:
            print(f"Built-in plot only supports 1D data (features={self.n_features}). Use custom plotting.")
            return
            
        # Handle Pandas
        if pd is not None:
             if isinstance(X, (pd.DataFrame, pd.Series)): X = X.values
             if isinstance(y, (pd.DataFrame, pd.Series)): y = y.values.flatten()
        
        X = X.reshape(-1, 1)
        y_pred = self.predict(X)
        
        plt.figure(figsize=(10, 6))
        plt.scatter(X, y, color='black', alpha=0.5, label='Data')
        
        # Sort X for clean line plotting
        sort_idx = np.argsort(X[:,0])
        plt.plot(X[sort_idx], y_pred[sort_idx], color='red', linewidth=2, label='Model Prediction')
        
        plt.title(f"{title}\nExpr: {self.expression}")
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.xlabel(self.feature_names[0] if self.feature_names else "x0")
        plt.ylabel("Target")
        
        if outfile:
            plt.savefig(outfile)
            print(f"Plot saved to {outfile}")
        else:
            plt.show()
