
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, Cube, 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
import logging
import re

# Configure logger
logger = logging.getLogger(__name__)

# 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.approx_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:
            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 with semantic filtering
        self.call_count += 1
        self.approx_gen = self.call_count // self.pop_size
        progress = min(1.0, self.approx_gen / (self.max_generations + 1e-6))

        # Semantic guidance parameters
        # Sample a subset of X for efficiency
        if hasattr(self.problem, 'X_data') and self.problem.X_data is not None:
            n_samples = min(256, len(self.problem.X_data))
            indices = rng.choice(len(self.problem.X_data), n_samples, replace=False)
            X_semantic = self.problem.X_data[indices]
            
            # Adaptive semantic threshold: start permissive, end strict.
            # Base scale from y variability.
            y_std = np.std(self.problem.y_data) if hasattr(self.problem, 'y_data') and len(self.problem.y_data) > 0 else 1.0
            y_var = y_std**2
            
            # Initial phase (first 10%): Allow anything (Infinite threshold)
            # This is critical to move away from constant-only populations.
            if progress < 0.1:
                max_semantic_delta = np.inf
            else:
                # Relax initial threshold: 10.0 * var at start of phase, 0.5 * var at end.
                max_semantic_delta = y_var * (10.0 * (1.0 - progress) + 0.5 * progress)
        else:
            X_semantic = None
            max_semantic_delta = None



        parent_str = str(tree_obj)

        new_tree = self.mutator.mutate(tree_obj, rng,
                                      generation=self.approx_gen,
                                      total_generations=self.max_generations,
                                      X_semantic=X_semantic,
                                      max_semantic_delta=max_semantic_delta)

        # 3. If mutation produced the exact same structure as the parent (no-op),
        # force a stronger change: random subtree replacement.
        # This acts as a secondary retry if the semantic filter was too strict or the mutator was unlucky.
        for _ in range(3):
            if str(new_tree) != parent_str:
                break
            size = new_tree.size
            site = rng.integers(0, size)
            fresh = generate_random_tree(self.n_features, rng, max_depth=2, # Small replacement
                                         binary_ops=self.binary_ops,
                                         unary_ops=self.unary_ops,
                                         ternary_ops=self.ternary_ops,
                                         literal_types=self.literal_types)
            new_tree = new_tree.set_subtree(site, fresh)

        # 4. Register New Tree with New ID
        new_id = self.problem.get_next_id()
        self.problem.register_tree(new_id, new_tree)

        # 5. Update Container Position ID
        new_pos = np.zeros((1, 3))
        new_pos[0, 0] = int(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 for Symbolic Regression.
    Expected signature: func(c1, c2) -> (new_c1, new_c2)
    """
    def __init__(self, problem):
        self.problem = problem

    def __call__(self, p1, p2):
        import copy
        import numpy as np
        
        # 1. Clone Containers
        c1 = copy.deepcopy(p1)
        c2 = copy.deepcopy(p2)
        
        # 2. Get Trees
        t1_id = float(c1.AtomPositionManager.atomPositions[0, 0])
        t2_id = float(c2.AtomPositionManager.atomPositions[0, 0])
        
        tree1 = self.problem.get_tree(t1_id)
        tree2 = self.problem.get_tree(t2_id)
        
        if tree1 is None or tree2 is None:
            return c1, c2
            
        # Clone actual tree objects to avoid messing with originals
        tree1 = tree1.clone()
        tree2 = tree2.clone()
        
        # 3. Pick random swap points
        rng = np.random.default_rng()
        idx1 = rng.integers(0, tree1.size)
        idx2 = rng.integers(0, tree2.size)
        
        # 4. Perform Swap
        # Note: set_subtree assumes mutable replacement in the clone
        sub1 = tree1.get_subtree(idx1).clone()
        sub2 = tree2.get_subtree(idx2).clone()
        
        tree1 = tree1.set_subtree(idx1, sub2)
        tree2 = tree2.set_subtree(idx2, sub1)
        
        # 5. Register and Update
        new_id1 = int(self.problem.get_next_id())
        new_id2 = int(self.problem.get_next_id())
        
        self.problem.register_tree(new_id1, tree1)
        self.problem.register_tree(new_id2, tree2)
        
        c1.AtomPositionManager.atomPositions[0, 0] = new_id1
        c2.AtomPositionManager.atomPositions[0, 0] = new_id2
        
        # Invalidate Energy/Objectives
        c1.AtomPositionManager.E = None
        c2.AtomPositionManager.E = None
        
        return c1, c2


class SymbolicRegressionProblem(ElementwiseProblem):
    def __init__(self, X, y, n_features, binary_ops, unary_ops, literal_types, 
                 complexity_penalty=0.0, optimizer='BFGS', n_restarts=0,
                 objectives=None, complexity_weights=None, residual_mode="feature_corr",
                 huber_delta=1.0):
        # Default objectives: MAE (robust error), Complexity (weighted), Residual Structure (correlation)
        self.objectives = objectives or ["mae", "complexity", "residual_structure"]
        n_obj = len(self.objectives)
        
        super().__init__(n_var=1, n_obj=n_obj, 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.optimizer = optimizer
        self.n_restarts = n_restarts
        
        # New parameters for multi-objective refinement
        self.complexity_weights = complexity_weights
        self.residual_mode = residual_mode
        self.huber_delta = huber_delta
        
        self.tree_registry = {}
        self.next_id = 10000 
        self.results_cache = {} # (formula, penalty) -> [objs, 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 compute_error(self, y_true, y_pred, kind):
        """ Computes a variety of error metrics for regression. """
        kind = kind.lower().strip()
        diff = y_true - y_pred
        
        if kind == "mse":
            return np.mean(diff**2)
        elif kind == "rmse":
            return np.sqrt(np.mean(diff**2))
        elif kind == "nrmse":
            rmse = np.sqrt(np.mean(diff**2))
            return rmse / (np.std(y_true) + 1e-8)
        elif kind == "mae":
            return np.mean(np.abs(diff))
        elif kind == "huber":
            delta = self.huber_delta
            abs_diff = np.abs(diff)
            return np.mean(np.where(
                abs_diff <= delta,
                0.5 * diff**2,
                delta * (abs_diff - 0.5 * delta)
            ))
        return np.mean(diff**2) # Default to MSE

    def compute_residual_structure(self, y_true, y_pred):
        """ Computes metrics assessing the predictability of residuals. """
        r = y_true - y_pred
        mode = self.residual_mode.lower().strip()
        
        if mode == "feature_corr":
            # Max correlation of residuals with any input feature
            corrs = []
            for j in range(self.X_data.shape[1]):
                xj = self.X_data[:, j]
                if np.std(xj) > 1e-8 and np.std(r) > 1e-8:
                    c = np.corrcoef(r, xj)[0, 1]
                    corrs.append(abs(c))
            return max(corrs) if corrs else 0.0
            
        elif mode == "prediction_corr":
            # Correlation of residuals with predictions (linearity check)
            if np.std(r) > 1e-8 and np.std(y_pred) > 1e-8:
                return abs(np.corrcoef(r, y_pred)[0, 1])
            return 0.0
            
        elif mode == "heteroscedasticity":
            # Correlation of absolute residuals with predictions (scale check)
            abs_r = np.abs(r)
            if np.std(abs_r) > 1e-8 and np.std(y_pred) > 1e-8:
                return abs(np.corrcoef(abs_r, y_pred)[0, 1])
            return 0.0
            
        return 0.0

    def compute_complexity(self, tree):
        """ Computes weighted symbolic complexity of a tree. """
        # Default Weights based on relative 'cost' of discovery/interpretability
        weights = self.complexity_weights or {
            "Add": 1.0, "Sub": 1.0, "Mul": 1.2, "Div": 2.0,
            "Sin": 2.0, "Cos": 2.0, "Exp": 3.0, "Log": 3.0,
            "Sqrt": 3.0, "Abs": 1.5, "Square": 1.5, "Cube": 2.0, "Neg": 1.0,
            "SoftIf": 4.0, "Threshold": 4.0, "Interval": 4.0,
            "Ratio": 3.0, "Gaussian": 3.0, "Variable": 1.0, "Constant": 1.0,
            "OptimizableConstant": 1.0
        }
        
        def traverse(node):
            name = node.__class__.__name__
            w = weights.get(name, 1.0)
            total = w
            # Handle children based on Node type
            for child in node.children:
                total += traverse(child)
            return total
            
        # Add a small depth penalty to encourage flatter trees for same node count
        return traverse(tree) + 0.5 * tree.depth

    def _optimize_constants(self, tree, opt_nodes):
        if not opt_nodes:
            return None, None

        # Determine optimizers list
        optimizers = self.optimizer
        if isinstance(optimizers, str):
            optimizers = [optimizers]
            
        def get_params():
            x = []
            for node, param_name in opt_nodes:
                val = getattr(node, param_name)
                if param_name == 'sharpness':
                    x.append(np.log(max(1e-3, val)))
                else:
                    x.append(val)
            return np.array(x)

        def set_params(params):
            for i, (node, param_name) in enumerate(opt_nodes):
                val = params[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)

        def objective(params):
            set_params(params)
            try:
                yp = tree.evaluate(self.X_data)
                return np.mean((self.y_data - yp)**2)
            except Exception:
                return 1e18

        best_mse = np.inf
        best_params = None
        
        initial_params = get_params()
        
        for attempt in range(self.n_restarts + 1):
            if attempt == 0:
                current_params = initial_params.copy()
            else:
                # Random restart: perturb initial_params
                current_params = initial_params + np.random.normal(0, 1.0, size=len(initial_params))
            
            for opt_name in optimizers:
                opt_name_ln = opt_name.lower().strip()
                try:
                    if opt_name_ln in ['bfgs', 'l-bfgs-b', 'nelder-mead']:
                        res = scipy_minimize(objective, current_params, method=opt_name_ln.upper(), options={'maxiter': 50})
                        current_params = res.x
                    elif opt_name_ln == 'cma-es':
                        try:
                            import cma
                            # Use quiet mode, low iterations for speed
                            res = cma.fmin2(objective, current_params, 0.5, {'maxiter': 50, 'verbose': -9})
                            current_params = res[0]
                        except ImportError:
                            pass # Fallback to next or keep current
                except Exception:
                    continue # Skip failed optimizer steps
                
            current_mse = objective(current_params)
            if current_mse < best_mse:
                best_mse = current_mse
                best_params = current_params.copy()
        
        # Apply the best parameters found
        if best_params is not None:
            set_params(best_params)
        return best_params, best_mse

    def evaluate_tree(self, tree, tree_id=None):
        try:
            # -----------------------------------------------------------
            # Pre-Optimization Simplification (local copy only)
            # -----------------------------------------------------------
            # Work on a simplified copy so the original tree in tree_registry
            # is NOT overwritten — preserving genetic identity / diversity.
            eval_tree = tree.simplify_aggressive()

            # -----------------------------------------------------------
            # Result Caching (by formula of simplified tree)
            # -----------------------------------------------------------
            # Use current complexity penalty (potentially increased by stagnation)
            penalty = getattr(self, "current_complexity_penalty", 0.0)
            
            # -----------------------------------------------------------
            # Result Caching (by formula AND complexity penalty)
            # -----------------------------------------------------------
            # We cache the full result including penalized complexity.
            # If the penalty changes, we need to re-evaluate or at least re-calculate.
            formula_key = (str(eval_tree), float(penalty))
            if formula_key in self.results_cache:
                # Return cached result — but keep the ORIGINAL tree in registry
                return self.results_cache[formula_key]

            # NOTE: We intentionally do NOT persist eval_tree to tree_registry here.
            # The registry should hold the original (unsimplified) genetic tree so
            # that the engine's feature/hash system can distinguish distinct individuals
            # that happen to simplify to the same formula.

            # -----------------------------------------------------------
            # -----------------------------------------------------------
            # Memetic Algorithm Step: Constant Optimization
            # -----------------------------------------------------------
            opt_nodes = []

            def collect_opt_nodes(node):
                # We prioritize OptimizableConstant but also support Threshold/Interval parameters
                if isinstance(node, OptimizableConstant):
                    opt_nodes.append((node, 'value'))
                elif isinstance(node, (Threshold, Interval, Ratio, Gaussian)):
                    for attr in ['val', 'low', 'high', 'threshold', 'sharpness', 'mu', 'sigma']:
                        if hasattr(node, attr):
                            opt_nodes.append((node, attr))

                for child in node.children:
                    collect_opt_nodes(child)

            collect_opt_nodes(eval_tree)

            # -----------------------------------------------------------
            # 2. Run Local Optimizer (BFGS, CMA-ES, etc.)
            # -----------------------------------------------------------
            if opt_nodes:
                # OPTIMIZE in eval_tree
                self._optimize_constants(eval_tree, opt_nodes)

                # --- Lamarckian Back-propagation ---
                # Sync values from optimized eval_tree -> original matching nodes in 'tree'
                orig_opt_nodes = []
                def collect_orig_opt(node):
                    if isinstance(node, OptimizableConstant): orig_opt_nodes.append(node)
                    for child in node.children: collect_orig_opt(child)
                collect_orig_opt(tree)
                
                eval_opt_objs = []
                def collect_eval_opt(node):
                    if isinstance(node, OptimizableConstant): eval_opt_objs.append(node)
                    for child in node.children: collect_eval_opt(child)
                collect_eval_opt(eval_tree)

                # Update the genetic persistent tree if structure matches (most common case)
                if len(orig_opt_nodes) == len(eval_opt_objs):
                    for o, e in zip(orig_opt_nodes, eval_opt_objs):
                        o.value = e.value

            # Final evaluation on the simplified, now optimized tree
            eval_tree = eval_tree.simplify_aggressive()
            y_pred = eval_tree.evaluate(self.X_data)
            mse = np.mean((self.y_data - y_pred)**2)


            # -----------------------------------------------------------
            # Compute Multi-Objective Results
            # -----------------------------------------------------------
            y_pred = eval_tree.evaluate(self.X_data)
            
            objs = []
            for obj_name in self.objectives:
                obj_name_ln = obj_name.lower().strip()
                
                if obj_name_ln in ["mse", "rmse", "nrmse", "mae", "huber"]:
                    objs.append(self.compute_error(self.y_data, y_pred, obj_name_ln))
                elif obj_name_ln == "complexity":
                    objs.append(self.compute_complexity(eval_tree))
                elif obj_name_ln == "residual_structure":
                    objs.append(self.compute_residual_structure(self.y_data, y_pred))
                else:
                    objs.append(1.0) # Placeholder for unknown
                    
            result = objs + [eval_tree]
            self.results_cache[formula_key] = result
            return result
        except Exception as e:
            if getattr(self, "verbose", False):
                logger.error(f"Error evaluating tree: {e}")
            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)
            # Return full set of objectives to the engine
            out["F"] = objs[:len(self.objectives)]
        except Exception as e:
            if getattr(self, "verbose", False):
                logger.error(f"Problem evaluation fail: {e}")
            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, optimizer='BFGS', n_restarts=0,
                 objectives=None, complexity_weights=None, residual_mode="feature_corr",
                 huber_delta=1.0):
        # 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, optimizer, n_restarts,
                         objectives, complexity_weights, residual_mode, huber_delta)
        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:
            if getattr(self, "verbose", False):
                logger.error(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[:len(self.objectives)]
        except Exception as e:
            if getattr(self, "verbose", False):
                logger.error(f"GE evaluation fail: {e}")
            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,
                 optimizer='BFGS',
                 n_restarts=0,
                 seed_fraction=1.0,
                 save_generations=False,
                 objectives=None,
                 complexity_weights=None,
                 residual_mode="feature_corr",
                 huber_delta=1.0):
        
        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
        self.optimizer = optimizer
        self.n_restarts = n_restarts
        self.seed_fraction = seed_fraction
        self.save_generations = save_generations
        self.engine = engine
        
        # Multi-objective configuration
        self.objectives = objectives or ["mae", "complexity", "residual_structure"]
        self.complexity_weights = complexity_weights
        self.residual_mode = residual_mode
        self.huber_delta = huber_delta
        # --- 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', 'square', 'cube'], # Added polynomial power shorthands
            "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', 'sqrt'],
                "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, 'cube': Cube, '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 _get_initial_seeds(self, X, y, rng):
        """Generate heuristic initial trees to seed the population."""
        y_mean = np.mean(y)
        y_std = np.std(y) + 1e-6
        y_max = np.max(y)  # amplitude for Gaussian seeds
        y_argmax_idx = int(np.argmax(y))  # sample index of peak

        feat_stats = []
        for f in range(self.n_features):
            col = X[:, f]
            feat_stats.append({
                'min':    float(np.min(col)),
                'max':    float(np.max(col)),
                'mean':   float(np.mean(col)),
                'std':    float(np.std(col)) + 1e-6,
                'argmax': float(col[y_argmax_idx]),
                'p25':    float(np.percentile(col, 25)),
                'p50':    float(np.percentile(col, 50)),
                'p75':    float(np.percentile(col, 75)),
            })

        seeds = []
        
        # Helper to check if ops are allowed
        has_add = Add in self.binary_ops
        has_sub = Sub in self.binary_ops
        has_mul = Mul in self.binary_ops
        has_div = Div in self.binary_ops
        has_neg = Neg in self.unary_ops
        has_square = Square in self.unary_ops
        has_sin = Sin in self.unary_ops
        has_cos = Cos in self.unary_ops
        has_gauss = Gaussian in self.literal_types
        has_ratio = Ratio in self.literal_types
        has_opt_const = OptimizableConstant in self.literal_types
        has_const = Constant in self.literal_types
        has_var = Variable in self.literal_types

        def wrap_const(val):
            if has_opt_const: return OptimizableConstant(val)
            if has_const: return Constant(val)
            return None

        # 1. Constants (Bias)
        c_mean = wrap_const(y_mean)
        if c_mean: seeds.append(c_mean)
        c_zero = wrap_const(0.0)
        if c_zero: seeds.append(c_zero)

        # 2. Linear Regressions (y ~ a*xi + b)
        # We add one true linear seed per feature with fitted coefficients
        if has_var and has_add and has_mul:
            for f in range(self.n_features):
                xi = X[:, f]
                x_mean = feat_stats[f]['mean']
                x_var = (feat_stats[f]['std'])**2
                cov = np.mean((xi - x_mean) * (y - y_mean))
                
                a = cov / x_var
                b = y_mean - a * x_mean
                
                # Seed with x * a + b
                # This directly provides A*x + B structure for BFGS to refine
                if has_opt_const:
                    # Model: (x * a) + b
                    lin_tree = Add(Mul(Variable(f), OptimizableConstant(a)), OptimizableConstant(b))
                    seeds.append(lin_tree)
                    # Also just the scaled variable: x * a
                    seeds.append(Mul(Variable(f), OptimizableConstant(a)))
                elif has_const:
                    seeds.append(Add(Mul(Variable(f), Constant(a)), Constant(b)))

        # 3. Linear (Multivariate Sum)
        if self.n_features > 1 and has_add and has_mul and has_var:
            poly_sum = wrap_const(y_mean)
            if poly_sum:
                for f in range(min(self.n_features, 10)):
                    c_s = wrap_const(0.1)
                    if c_s:
                        term = Mul(c_s, Variable(f))
                        poly_sum = Add(poly_sum, term)
                seeds.append(poly_sum)

        # 4. Quadratic
        if has_square and has_mul and has_var:
            for f in range(self.n_features):
                q = Square(Variable(f))
                scale = y_std / (feat_stats[f]['std']**2 + 1e-6)
                c_scale = wrap_const(scale)
                if c_scale: seeds.append(Mul(c_scale, q))

        # 5. Bilinear Interactions (x_i * x_j)
        if self.n_features >= 2 and has_mul:
            for i in range(min(self.n_features, 3)):
                for j in range(i + 1, min(self.n_features, 4)):
                    seeds.append(Mul(Variable(i), Variable(j)))

        # 6. Trigonometric
        if (has_sin or has_cos) and has_var:
            for f in range(min(self.n_features, 5)):
                freq = 1.0 / feat_stats[f]['std']
                c_freq = wrap_const(freq)
                if c_freq and has_mul:
                    arg = Mul(c_freq, Variable(f))
                else:
                    arg = Variable(f)
                if has_sin: seeds.append(Sin(arg))
                if has_cos: seeds.append(Cos(arg))

        # 7. Gaussian
        if has_gauss and has_mul and has_var:
            for f in range(min(self.n_features, 5)):
                std = feat_stats[f]['std']
                for mu in [feat_stats[f]['argmax'], feat_stats[f]['p50']]:
                    g = Gaussian(f, mu, std * 0.3)
                    c_amp = wrap_const(y_max)
                    if c_amp:
                        seeds.append(Mul(c_amp, g))

        # 8. Ratio
        if has_ratio and self.n_features >= 2:
            seeds.append(Ratio(0, 1, 1.0))

        if self.verbose:
            logger.info(f"Generated {len(seeds)} heuristic seeds using allowed operators.")
            
        return seeds

    def fit(self, X, y):
        # 0. Data Handling (Auto-detect Pandas)
        if pd is not None:
            if isinstance(X, (pd.DataFrame, pd.Series)):
                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:
            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,
                optimizer=self.optimizer,
                n_restarts=self.n_restarts,
                objectives=self.objectives,
                complexity_weights=self.complexity_weights,
                residual_mode=self.residual_mode,
                huber_delta=self.huber_delta
            )
        else:
            problem = SymbolicRegressionProblem(
                X, y,
                self.n_features,
                self.binary_ops,
                self.unary_ops,
                self.literal_types,
                self.complexity_penalty,
                optimizer=self.optimizer,
                n_restarts=self.n_restarts,
                objectives=self.objectives,
                complexity_weights=self.complexity_weights,
                residual_mode=self.residual_mode,
                huber_delta=self.huber_delta
            )

        # 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:
            # Tree-based: 30% Informed Seeds, 70% Ramped Half-and-Half Random
            init_ids = np.zeros((self.pop_size, 1))
            seeds = self._get_initial_seeds(X, y, rng)
            
            n_seeds = int(self.pop_size * self.seed_fraction)
            
            for i in range(self.pop_size):
                tid = i + 1
                init_ids[i, 0] = tid
                
                if i < n_seeds and len(seeds) > 0:
                    # Pick from seed library (cycling if seeds < n_seeds)
                    tree = copy.deepcopy(seeds[i % len(seeds)])
                else:
                    # Ramped Half-and-Half (depths 2 to 6)
                    # Calculate depth and method based on index remainder
                    pop_idx = i - n_seeds
                    n_random = self.pop_size - n_seeds
                    
                    # Distribute depths 2..6
                    depth = 2 + (pop_idx % 5) 
                    # Alternating method
                    meth = 'grow' if (pop_idx // 5) % 2 == 0 else 'full'
                    
                    tree = generate_random_tree(self.n_features, rng, max_depth=depth,
                                                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",
                eliminate_duplicates=True # Avoid redundant trees
            )
            # Enable elitism by ensuring top-k are kept (GA usually keeps best, but we can be explicit if needed)
            # For NSGA2/3, the best front is always kept.
            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(problem=problem)
        
        # 4. Minimize
        res = minimize(
            problem,
            algorithm,
            termination=('n_gen', self.generations),
            mutation_rate=self.mutation_rate,
            seed=self.random_state,
            verbose=self.verbose,
            save_generations=self.save_generations
        )
        
        # Extract Pareto Front Results
        self.pareto_results_ = []
        if res.pop:
            best_ind = None
            min_f = float('inf')
            
            # Weights for scalarized best extraction (consistent with minimize)
            # Default to MSE + Weighted Complexity for "best" model selection if not specified
            w_scalar = np.ones(len(problem.objectives))
            # If we have complexity, we might want to prioritize it for 'best' selection
            if "complexity" in problem.objectives:
                idx = problem.objectives.index("complexity")
                w_scalar[idx] = problem.current_complexity_penalty
            
            seen_formulas = set()
            
            for ind in res.pop:
                try:
                    meta = getattr(ind.AtomPositionManager, "metadata", {})
                    if "objectives" in meta:
                        objs = np.array(meta["objectives"]).flatten()
                        
                        # Find indices for MSE and Complexity for result mapping
                        # We use MSE (or first error metric) as the primary MSE value
                        error_metrics = ["mse", "rmse", "nrmse", "mae", "huber"]
                        mse_idx = -1
                        for i, name in enumerate(problem.objectives):
                            if name in error_metrics:
                                mse_idx = i
                                break
                        
                        complexity_idx = -1
                        if "complexity" in problem.objectives:
                            complexity_idx = problem.objectives.index("complexity")
                            
                        mse_val = objs[mse_idx] if mse_idx != -1 else objs[0]
                        complexity_val = objs[complexity_idx] if complexity_idx != -1 else 0.0
                        
                        # Scalarize for comparison
                        f_scalar = np.dot(objs, w_scalar[:len(objs)])
                    else:
                        e_val = ind.AtomPositionManager.E
                        if e_val is not None:
                             f_scalar = float(np.atleast_1d(e_val)[0])
                             mse_val = f_scalar
                             complexity_val = 0.0 # Unknown
                        else:
                             continue
                    
                    # Extract Tree
                    tree = None
                    if self.engine == "ge":
                        genes = ind.AtomPositionManager.atomPositions[:, 0].astype(int)
                        tree = problem.mapper.map(genes)
                        if tree:
                            tree = tree.simplify_aggressive()
                    else:
                        pos = ind.AtomPositionManager.atomPositions
                        if len(pos) > 0:
                            tid = pos[0][0]
                            tree = problem.get_tree(tid)
                            if tree:
                                tree = tree.simplify_aggressive()
                    
                    if tree:
                        formula = str(tree)
                        if formula not in seen_formulas:
                            seen_formulas.add(formula)
                            self.pareto_results_.append({
                                'tree': tree,
                                'mse': mse_val,
                                'complexity': complexity_val,
                                'fitness': f_scalar
                            })
                    
                    if f_scalar is not None and f_scalar < min_f:
                        min_f = f_scalar
                        best_ind = ind
                        self.best_tree = tree
                        self.best_fitness = min_f
                except Exception as e:
                    if self.verbose:
                        logger.warning(f"Error extracting individual: {e}")
                    continue

            # Sort Pareto front by complexity for coherent plotting
            self.pareto_results_.sort(key=lambda x: x['complexity'])
                
        return self

    def predict(self, X):
        if self.best_tree:
            X = np.asarray(X)
            if X.ndim == 1:
                X = X.reshape(-1, 1)
            return self.best_tree.evaluate(X)
        X = np.asarray(X)
        return np.zeros(X.shape[0])

    def predict_all(self, X):
        """
        Predicts target values for each model discovered in the Pareto front.
        
        Args:
            X: Input features.
            
        Returns:
            Dictionary mapping model expressions to their predictions.
        """
        if not hasattr(self, 'pareto_results_') or not self.pareto_results_:
            return {}
            
        X = np.asarray(X)
        if X.ndim == 1:
            X = X.reshape(-1, 1)
            
        results = {}
        for res in self.pareto_results_:
            tree = res['tree']
            formula = self._format_expression(tree)
            try:
                results[formula] = tree.evaluate(X)
            except Exception:
                results[formula] = np.zeros(X.shape[0])
                
        return results

    def get_models(self):
        """
        Returns a structured summary of all models discovered in the Pareto front.
        
        Returns:
            List of dictionaries containing 'formula', 'mse', and 'complexity'.
        """
        if not hasattr(self, 'pareto_results_') or not self.pareto_results_:
            return []
            
        models = []
        for res in self.pareto_results_:
            models.append({
                'formula': self._format_expression(res['tree']),
                'mse': float(res['mse']),
                'complexity': float(res['complexity'])
            })
        return models

    def _format_expression(self, tree):
        """Helper to format a tree expression with feature names."""
        if tree is None:
            return ""
        formula = str(tree)
        if self.feature_names:
            # Sort feature names by length descending to avoid partial replacement of x1 with x10
            for i in range(len(self.feature_names) - 1, -1, -1):
                raw_token = f"x{i}"
                pattern = rf"\b{raw_token}\b"
                formula = re.sub(pattern, self.feature_names[i], formula)
        return formula

    def score(self, X, y):
        y_pred = self.predict(X)
        y = np.asarray(y).flatten()
        ss_res = np.sum((y - y_pred)**2)
        ss_tot = np.sum((y - np.mean(y))**2)
        if ss_tot == 0:
            return 1.0 if np.allclose(y, y_pred) else -np.inf
        return 1 - (ss_res / ss_tot)

    @property
    def expression(self):
        return self._format_expression(self.best_tree) or "None"

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

    def plot_pareto_front(self, outfile=None):
        """
        Plots the Complexity vs MSE Pareto front for discovered models.
        """
        if plt is None:
            print("Matplotlib not installed.")
            return
        
        if not hasattr(self, 'pareto_results_') or not self.pareto_results_:
            print("No Pareto results available. Did you run fit()?")
            return
            
        complexities = [x['complexity'] for x in self.pareto_results_]
        mses = [x['mse'] for x in self.pareto_results_]
        
        plt.figure(figsize=(8, 5))
        plt.scatter(complexities, mses, color='blue', alpha=0.6, label='Discovery Front')
        
        # Draw a step-like Pareto front line
        # Since they are already sorted by complexity, we just compute the running minimum of MSE
        sorted_copy = sorted(self.pareto_results_, key=lambda x: x['complexity'])
        px = []
        py = []
        curr_min_mse = float('inf')
        for res in sorted_copy:
            if res['mse'] < curr_min_mse:
                curr_min_mse = res['mse']
                px.append(res['complexity'])
                py.append(res['mse'])
        
        plt.step(px, py, where='post', color='red', linestyle='--', alpha=0.8, label='Pareto Front')
        
        plt.yscale('log')
        plt.xlabel('Complexity (Size + 0.5*Depth)')
        plt.ylabel('MSE (Log Scale)')
        plt.title('Symbolic Regression Pareto Front')
        plt.grid(True, which="both", ls="-", alpha=0.2)
        plt.legend()
        
        if outfile:
            plt.savefig(outfile)
        else:
            plt.show()

    def plot_all_models(self, X, y, outfile=None, max_models=10):
        """
        Plots the top Pareto models overlaid on the data (1D only).
        """
        if plt is None: return
        if self.n_features != 1:
            print("plot_all_models only supports 1D data.")
            return

        if not hasattr(self, 'pareto_results_') or not self.pareto_results_:
            return

        # Prepare Data
        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)
        sort_idx = np.argsort(X[:,0])
        X_sorted = X[sort_idx]
        
        plt.figure(figsize=(10, 6))
        plt.scatter(X, y, color='black', alpha=0.3, label='Data', s=10)

        # Get the real Pareto front (dominant models)
        pareto_models = []
        curr_min_mse = float('inf')
        for res in self.pareto_results_:
            if res['mse'] < curr_min_mse:
                curr_min_mse = res['mse']
                pareto_models.append(res)
        
        # Limit number of models to plot
        if len(pareto_models) > max_models:
            # Pick a subset spaced out by complexity
            indices = np.linspace(0, len(pareto_models) - 1, max_models, dtype=int)
            plot_subset = [pareto_models[i] for i in indices]
        else:
            plot_subset = pareto_models

        # Use a colormap
        cmap = plt.get_cmap('viridis')
        for i, res in enumerate(plot_subset):
            y_pred = res['tree'].evaluate(X)
            label = f"Size: {res['tree'].size:.0f}, MSE: {res['mse']:.3f}"
            
            # Use color based on complexity
            color = cmap(i / max(1, len(plot_subset)-1))
            plt.plot(X_sorted, y_pred[sort_idx], color=color, alpha=0.7, linewidth=1.5, label=label)

        plt.title(f"Discovered Pareto Front Models\nBest: {self.expression}")
        plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize='small')
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.grid(True, alpha=0.2)
        plt.tight_layout()

        if outfile:
            plt.savefig(outfile)
        else:
            plt.show()

class MultiTargetSymbolicRegressor:
    """
    Orchestrator for multi-output symbolic regression (R^n -> R^m).
    Fits an independent SymbolicRegressor for each target dimension.
    """
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self.regressors = []
        self.n_targets = 0

    def fit(self, X, y):
        # 0. Data Handling for Target
        if hasattr(y, 'values'): # Pandas
            y = y.values
            
        if y.ndim == 1:
            y = y.reshape(-1, 1)
        
        self.n_targets = y.shape[1]
        self.regressors = []
        
        for i in range(self.n_targets):
            if self.kwargs.get("verbose"):
                print(f"\n[MultiTarget] Fitting Output Dim {i+1}/{self.n_targets}...")
            
            reg = SymbolicRegressor(**self.kwargs)
            reg.fit(X, y[:, i])
            self.regressors.append(reg)
            
        return self

    def predict(self, X):
        preds = []
        for reg in self.regressors:
            preds.append(reg.predict(X))
        return np.column_stack(preds)

    def score(self, X, y):
        if hasattr(y, 'values'): y = y.values
        if y.ndim == 1: y = y.reshape(-1, 1)
        
        scores = []
        for i, reg in enumerate(self.regressors):
            scores.append(reg.score(X, y[:, i]))
        return np.mean(scores)

    @property
    def expressions(self):
        """ Returns a list of strings representing the expressions for each output dimension. """
        return [reg.expression for reg in self.regressors]
    
    @property
    def best_trees(self):
        """ Returns the list of best Tree objects found for each output. """
        return [reg.best_tree for reg in self.regressors]
