from .factory import EngineFactory
from .engines.base import SymbolicRegressionEngine
from .utils.data_processing import dataset_builder, DataTransformer, partition_by_group
from .utils.utils import build_cv_folds
from .utils.equation_processing import Equation, ClassificationEquation
from .utils.evaluation import evaluate_equations, plot_pareto_front, compute_pareto_front, METRIC_PROPERTIES
from .utils.loss import get_loss, losses_dict
import warnings
from typing import List, Optional, Iterable
import pandas as pd
import numpy as np
import sympy as sp
import copy
from concurrent.futures import ProcessPoolExecutor, as_completed

AVAILABLE_ENGINES = ["rsrm", "pysr", "linear"]

# --------------------------------------------------------------------------------------------------------
# Module-level worker for a single search repeat (must be top-level for pickling in ProcessPoolExecutor)
# --------------------------------------------------------------------------------------------------------
def _worker_search_repeat(
    i: int,
    Xs: pd.DataFrame,
    ys: pd.Series,
    gs: pd.Series,
    engine_name: str,
    engine_args: dict
):
    """
    One iteration of the SR search:
      - Rebuilds its own engine from engine_name + engine_args
      - Optionally samples the data
      - Fits the engine
      - Returns raw expression strings + the fitted engine
    """
    # Reinstantiate engine
    engine = EngineFactory.create(engine_name, **engine_args)

    # Fit
    if gs.nunique() == 1:
        engine.fit(Xs.values, ys.values)
    else:
        X_parts, y_parts = partition_by_group(Xs, ys, gs)
        engine.fit_multi(X_parts, y_parts)

    # Extract raw expressions
    raw_expressions = engine.get_expressions()

    return raw_expressions, engine


class SymbolicRegressor:
    """
    High-level API for the LEAF Symbolic Regression framework. 
    
    This class provides a simplified, function-based interface for conducting
    a full symbolic regression analysis through a user-friendly abstraction layer over 
    different symbolic regression engines.

    Workflow Stages:
    1.  **Initialize**: Create a regressor instance with `initialize_regressor`,
        defining the core engine and loss function.
    2.  **Search**: Discover candidate equations with `search_equations`. This is the
        computationally intensive discovery phase.
    3.  **Train**: Calibrate the parameters of the discovered equations using
        `fit_model_parameters`. This can be done on the search data or new
        data, with or without cross-validation.
    4.  **Evaluate**: Assess the performance of the trained models on any dataset
        (e.g., a hold-out test set) using `evaluate_equations`.
    5.  **Analyze & Manage**: View results (`view_results_table`), plot trade-offs
        (`plot_pareto_front`), inspect individual equations (`show_equation_details`),
        and curate the model set (`exclude_equations`).
    6.  **Persist**: Save or load your progress with `save_analysis` and
        `load_analysis` to resume work later.

    Examples
    --------
    See detailed examples in the examples folder.

    Metrics
    -------
    Allowed metrics:
    ['R2', 'RMSE', 'MAE', 'F1', 'TSS', 'PseudoR2', 'AIC', 'AICc', 'AIC_w', 'AICc_w']
    """

    def __init__(self, engine: str, loss: str = 'MSE', force_consistent_sign: Optional[bool] = False, **engine_args):
        f"""
        Initialize a new symbolic regression analysis, specifying the backend SR engine and its parameters.

        Creates a `SymbolicRegressor` instance that manages the state of the
        symbolic regression pipeline. This is the first step before searching
        for equations.

        Parameters. Some are supported by only some engines and will be ignored.
        ----------
        engine : str
            The name of the symbolic regression engine to use. One of:

        loss : str, optional
            Specifies the loss function to optimize. This argument will be mapped to the engine-specific loss keyword.
            Supported by: "rsrm", "pysr", "linear". 
            Allowed losses:
                PySR: ['MAE', 'MSE', 'BinaryCrossEntropy', 'ExpLoss', 'PoissonDeviance']
                RSRM: ['MAE', 'MSE', 'RMSE', 'R2', 'RobustTanh', 'BinaryCrossEntropy', 'ExpLoss', 'PoissonDeviance']
                Linear: ['MSE', 'PoissonDeviance', 'BinaryCrossEntropy']
            Default is "MSE".

        num_iterations : int, optional
            Specifies the number of training iterations or generations (if applicable to the engine).
            Supported by: "rsrm", "pysr".

        operators : list of str, optional
            Specifies the set of mathematical operators allowed in expressions (e.g., ["+", "-", "*", "/"]).
            Supported by: "rsrm", "pysr".

        max_complexity : int, optional
            Limits the symbolic complexity of generated expressions. Engines may interpret this differently.
            Supported by: "rsrm", "pysr".

        max_params : int, optional
            Limits the maximum number of free parameters (constants) allowed in generated expressions.
            Supported by: "rsrm", "pysr (multiview SR mode only)".

        max_depth : int, optional
            Limits the maximum depth of the expression tree. 
            Supported by: "rsrm", "pysr".

        force_consistent_sign (bool, optional):
            Whether to ensure all parameters maintain the same sign across groups during optimization of parameters.
            This ensure the force effect is consistent across groups (always positive or always negative). 
            This will be enforced both during `search_equations` and `fit_model_parameters`. 
            Supported by: "linear", "manual", "rsrm".
            Defaults to False.
        
        **engine_args : dict
            Additional engine-specific arguments. Unlike other parameters, these are passed directly to the underlying engine
            without translation, allowing for advanced customization.

        Notes
        -----
        - All arguments (including common ones like `loss`, `max_depth`, etc.) are passed through a translator
        specific to the chosen engine. If the argument is not recognized or not supported by the engine,
        it will be ignored with a warning.
        """
        self._engine_name: str = engine
        self.loss: str = loss
        engine_args['loss'] = loss
        self._is_classification_ = loss in ['BinaryCrossEntropy', 'ExpLoss']
        self.force_consistent_sign = False
        if engine in ('rsrm', 'linear', 'manual'):
            engine_args['force_consistent_sign'] = force_consistent_sign
            self.force_consistent_sign = force_consistent_sign
        engine_args = {k:v for k,v in engine_args.items() if v is not None}
        self._engine_args = engine_args

        # Attributes fitted during the `fit` call
        self._engine: Optional[SymbolicRegressionEngine] = None
        self.equations_: List[Equation] = []
        self.results_ = None
        self.transformer_: DataTransformer = None
        # Data needed for scoring after fit
        self._X_train_: Optional[pd.DataFrame] = None
        self._y_train_: Optional[pd.Series] = None
        self._groups_train_: Optional[pd.Series] = None
        self._cv_folds_ = None

    def search_equations(
        self,
        data: pd.DataFrame,
        formula: str,
        normalization: str = "zscore",
        folds: Iterable = None,
        n_threads: int = 1,
    ) -> pd.DataFrame:
        """
        Stage 1: Discover equation skeletons using the SR engine, with optional sampling
        and multithreading over repeats.

        Parameters
        ----------
        data : pandas.DataFrame
            Dataset to search for equations in.
        formula : str
            R-style formula (e.g., "'y ~ f(x1, ..., xn | group_1, ..., group_n)'").
        normalization : {'zscore','minmax','divide_by_gmd','none'}, default='zscore'
            Normalization method for feature columns.
        folds : Iterable of (train_idx, test_idx) tuples, optional
            Each element should be an integer array representing the testing indices for one fold. If provided, 
            the search will be repeated across folds using the specified sampling strategy. If None, the
            full dataset is used without resampling.
        n_threads : int, default=1
            Number of threads to use when running the repeats (1 = serial execution).

        Returns
        -------
        pandas.DataFrame
            Table of unique equations found.
        """
        # If folds is none we run once with the whole dataset
        folds = list(folds) if folds else [np.arange(len(data))]
        n_repeats = len(folds)
        # Validate inputs
        if not isinstance(n_threads, int) or n_threads < 1:
            raise ValueError("n_threads must be an integer ≥ 1")

        if self._engine_name == 'manual':
            self._engine_args['formula'] = formula

        # Data preprocessing
        print("1. Processing data for equation search based on formula...")
        processed = dataset_builder(data, formula, normalization)
        self.transformer_ = processed["transformer"]
        self._X_train_, self._y_train_, self._groups_train_ = (
            processed["X"],
            processed["y"],
            processed["groups"],
        )

        # Launch repeats
        print(
            f"2. Running engine '{self._engine_name}' over {n_repeats} folds "
            f"using up to {n_threads} processes..."
        )
        all_raw: List[str] = []
        last_engine = None

        tasks = []
        for i, train_idx in enumerate(folds):
            # Sampling
            processed = dataset_builder(data.iloc[train_idx], formula, normalization)
            X, y, groups = processed["X"], processed["y"], processed["groups"]
            tasks.append(
                (
                    i,
                    X,
                    y,
                    groups,
                    self._engine_name,
                    self._engine_args
                )
            )

        # Serial execution
        if n_threads == 1 or n_repeats == 1:
            for t in tasks:
                raws, eng = _worker_search_repeat(*t)
                all_raw.extend(raws)
                last_engine = eng
        # Parallel via ProcessPoolExecutor
        else:
            with ProcessPoolExecutor(max_workers=n_threads) as executor:
                futures = [executor.submit(_worker_search_repeat, *t) for t in tasks]
                for fut in as_completed(futures):
                    try:
                        raws, eng = fut.result()
                        all_raw.extend(raws)
                        last_engine = eng
                    except Exception as e:
                        warnings.warn(f"Error in parallel repeat: {e}")

        # Assign the last‐fitted engine back
        if last_engine is not None:
            self._engine = last_engine

        # Handle no results
        if not all_raw:
            warnings.warn("No equations found across any repeats.")
            return pd.DataFrame()

        # Build and dedupe Equation objects
        print(f"3. Found {len(all_raw)} raw skeletons. Deduplicating...")
        engine_map = dict(
            zip(self._engine.engine_variable_names, self.transformer_.placeholder_names)
        )
        unique_eqs: List[Equation] = []
        results: List[dict] = []
        next_temp_id = 0

        for raw in all_raw:
            try:
                skeleton = str(sp.sympify(raw).subs(engine_map))
                if self._is_classification_:
                    eq = ClassificationEquation(skeleton, transformer=self.transformer_, loss=self.loss, force_consistent_sign=self.force_consistent_sign)
                else:
                    eq = Equation(skeleton, transformer=self.transformer_, loss=self.loss, force_consistent_sign=self.force_consistent_sign)
                if eq in unique_eqs:
                    continue
                unique_eqs.append(eq)
                results.append(
                    {"temp_id": next_temp_id, "Equation": eq.readable_skeleton, "Complexity": eq.complexity}
                )
                next_temp_id += 1
            except Exception as e:
                warnings.warn(f"Failed to process Equation from `{raw}`: {e}")

        df = pd.DataFrame(results)
        if df.empty:
            warnings.warn("All discovered equations failed to process.")
            return pd.DataFrame()
        # Sort by Complexity (ascending); add a stable secondary key so ties are deterministic
        df = df.sort_values(by=["Complexity", "Equation"], ascending=[True, True]).reset_index(drop=True)

        df["ID"] = df.index
        id_map = dict(zip(df["temp_id"], df["ID"]))
        for temp_idx, eq in enumerate(unique_eqs):
            new_id = id_map[temp_idx]
            eq.id = int(new_id)
        
        # Reorder the equations by index as well
        ordered_equations = [None] * len(unique_eqs)
        for eq in unique_eqs:
            ordered_equations[eq.id] = eq
        
        # Store and return
        self.equations_ = ordered_equations
        self.results_ = df[["ID", "Equation", "Complexity"]].set_index("ID")
        return self.results_
    
    def fit_model_parameters(self, data: Optional[pd.DataFrame] = None, do_cv: bool = False, k_folds: int = None):
        """
        Stage 2: Fit the parameters of the discovered equations.

        Fits the global and group-specific parameters for each equation found
        during the search phase. By default, it uses the same dataset
        provided to `search_equations`, unless a new dataset is explicitly passed.

        Cross-validation is optional and disabled by default.

        Parameters
        ----------
        data (pd.DataFrame, optional): 
            The data to use for fitting. If None,
            the data from the `search_equations` call is used. Defaults to None.
        do_cv (bool, optional):
            Whether to perform k-fold cross-validation.
            Defaults to False.
        k_folds (int, optional):
            Number of folds used for cross-validation. Defaults to 10 if None.

            When stratification is applied, each fold must contain at least one
            sample from every stratum used for splitting. If the target `y_fit`
            is binary, stratification is performed on the joint (group, class)
            labels, meaning the smallest (group, class) combination must contain
            at least `k_folds` samples. Otherwise, stratification is performed
            only on the group labels, and the smallest group must contain at
            least `k_folds` samples.

            A ValueError is raised if `k_folds` exceeds the number of samples in
            the smallest stratum, since it would be impossible to distribute the
            data across folds while preserving the stratification.
                
            The generated folds are stored in the `_cv_folds_` attribute of this 
            `SymbolicRegressor` instance.

       Returns
        -------
        pandas.DataFrame
            The main results table, updated with:
            - ``'Loss'``
            - ``'Loss (cv)'`` (if `do_cv=True`)
        """
        if not self.equations_:
            raise RuntimeError("You must run equation_search() before fitting parameters.")
        if data is None:
            X_fit, y_fit, groups_fit = self._X_train_, self._y_train_, self._groups_train_
        else:
            processed = self.transformer_.transform(data, fit_mode=True)
            X_fit, y_fit, groups_fit = processed['X'], processed['y'], processed['groups']
            self._X_train_, self._y_train_, self._groups_train_ = X_fit, y_fit, groups_fit

        if do_cv:
            folds = build_cv_folds(X_fit, y_fit, groups_fit, k_folds=k_folds, shuffle=True, random_state=42)
            self._cv_folds_ = list(folds)

        print(f"Fitting parameters for {len(self.equations_)} equations...")
        loss_func = get_loss(self.loss)
        
        for eq in self.equations_:
            try:
                # Always fit on the full dataset first
                main_params = eq.fit(X_fit, y_fit, groups_fit, basinhopping=True)
                preds = eq.predict(X_fit, groups_fit, threshold=False)
                loss_value = loss_func(y_fit.values, preds)
                self.results_.loc[eq.id, 'Loss'] = loss_value
                eq.loss_value = loss_value
                if do_cv:
                    print(f"  - Performing CV for Equation with ID={eq.id}...")
                    out_of_fold_preds = np.zeros_like(y_fit.values, dtype=float)
                    for train_idx, val_idx in self._cv_folds_:
                        fold_params = eq.fit(X_fit.iloc[train_idx], y_fit.iloc[train_idx], groups_fit.iloc[train_idx], initial_params=main_params)
                        eq.store_cv_fold_params(fold_params)    # Store for later use in score
                        val_preds = eq.predict(X_fit.iloc[val_idx], groups_fit.iloc[val_idx], params_override=fold_params, threshold=False)
                        out_of_fold_preds[val_idx] = val_preds

                    eq.fitted_params = main_params  # Restore the main parameters
                    self.results_.loc[eq.id, 'Loss (cv)'] = loss_func(y_fit.values, out_of_fold_preds)
            except Exception as e:
                warnings.warn(f"Could not process equation with ID={eq.id}: {eq.readable_skeleton}. Error: {e}")

        print("Parameter fitting complete.")
        return self.results_.copy()
    
    def evaluate_equations(self, metrics: List[str] = None, metrics_cv: List[str] = None,
              data: Optional[pd.DataFrame] = None, equation_ids: List[int] = None, likelihood_for_AIC: Optional[callable]=None) -> pd.DataFrame:
        """
        Stage 3: Evaluate the fitted equations on a given dataset.

        Uses pre-fitted parameters to calculate performance metrics, typically
        on a hold-out test set or via cross-validation.

        Parameters
        ----------
        metrics : list of str, optional
            List of performance metrics to compute on the data set specified in the argument `test_data` or
            the training data used in `fit_model_parameters` if it is None. 
            To see allowed metrics, see description of SymbolicRegressor class.
        metrics_cv : list of str, optional
            List of performance metrics to compute using cross-validated models.  
            Requires `fit_model_parameters` to have been run with `do_cv=True`. 
            To see allowed metrics, see description of SymbolicRegressor class.
        data : pandas.DataFrame, optional
            An external dataset used to evaluate fitted equations on `metrics`. 
        equation_ids : sequence of int, default=None
            A list of equation IDs to score. If None, all equations are scored.
        likelihood_for_AIC_and_Elbow : callable, optional
            Custom likelihood function for AIC and Elbow calculation. Defaults to the Poisson likelihood if the loss is 
            PoissonDeviance, Bernoulli likelihood if the loss is `BinaryCrossEntropy` or `ExpLoss`, and the Gaussian 
            likelihood otherwise.

        Returns
        -------
        pd.DataFrame
            A table with the requested metrics for the specified equations.
        """
        if self.results_ is None or self.results_.empty or not self.equations_:
            raise RuntimeError("You must run equation_search() before scoring.")
        invalid_ids = [i for i in self.results_.index if not self.equations_[i].fitted_params]
        if equation_ids is None:
            ids = [i for i in self.results_.index if i not in invalid_ids]
        else:
            ids = [i for i in equation_ids if i in self.results_.index and i not in invalid_ids]
        if not ids:
            raise RuntimeError("The equations weren't fitted. You must run fit_model_parameters() before scoring.")
        metrics = metrics or []
        metrics_cv = metrics_cv or []
        eqs_to_score = [self.equations_[i] for i in ids]
        num_params_list = [len(eq.c_params) + len(eq.u_params) for eq in eqs_to_score]

        output_df = pd.DataFrame(index=ids)

        # --- Non-cv Scoring ---
        if metrics:
            if data is None:
                X_score, y_score, groups_score = self._X_train_, self._y_train_.values, self._groups_train_
            else: 
                processed = self.transformer_.transform(data)
                X_score, y_score, groups_score = processed['X'], processed['y'].values, processed['groups']
            
            predictions_list = [eq.predict(X_score, groups_score) for eq in eqs_to_score]
            train_results = evaluate_equations(
                predictions_list=predictions_list, num_params_list=num_params_list, y_true=y_score, groups=groups_score.factorize()[0],
                partial_results_df=self.results_.loc[ids], loss_name=self.loss, metrics_to_calc=metrics,
                likelihood_override=likelihood_for_AIC
            )
            output_df = output_df.join(train_results)
        
        # --- CV Scoring ---
        if metrics_cv:
            if not self.equations_ or self._cv_folds_ is None:
                raise ValueError("Cannot score with CV metrics. Run fit_equations() with do_cv=True first.")
             
            X_score, y_score, groups_score = self._X_train_, self._y_train_.values, self._groups_train_
            
            predictions_list = []
            for eq in eqs_to_score:
                out_of_fold_preds = np.zeros_like(y_score, dtype=float)
                for i, (_, val_idx) in enumerate(self._cv_folds_):
                    fold_params = eq.cv_params_per_fold[i]
                    val_preds = eq.predict(X_score.iloc[val_idx], groups_score.iloc[val_idx], params_override=fold_params)
                    out_of_fold_preds[val_idx] = val_preds
                predictions_list.append(out_of_fold_preds)

            cv_results = evaluate_equations(
                predictions_list=predictions_list, num_params_list=num_params_list, y_true=y_score, groups=groups_score.factorize()[0],
                partial_results_df=self.results_.loc[ids], loss_name=self.loss, metrics_to_calc=metrics_cv,
                likelihood_override=likelihood_for_AIC
            )

            rename_map = {col: f"{col} (cv)" for col in cv_results}
            cv_results.rename(columns=rename_map, inplace=True)
            output_df = output_df.join(cv_results)

        # Update main results table
        for col in output_df.columns: self.results_.loc[ids, col] = output_df[col]
        
        # Prepare and return output
        cols_to_return = ['Equation', 'Complexity'] + list(output_df.columns)
        df = self.results_.loc[ids].reset_index()[['ID'] + cols_to_return].set_index('ID')
        return self.get_pareto_front()
    
    def predict(self, data: pd.DataFrame, equation_id: int, threshold=None):
        """
        Generate predictions using a selected fitted equation.

        Parameters
        ----------
        data : pandas.DataFrame
            Input data on which to make predictions. The columns must match
            the features used during training.
        equation_id : int
            Identifier of the equation to use for prediction.
        threshold : float, optional
            Decision threshold used for binary classification. If provided,
            predicted probabilities greater than or equal to this value are
            mapped to the positive class and those below it to the negative
            class.

            If None, the method returns:
            - raw predictions for regression tasks, or
            - predicted probabilities for classification tasks.


        Returns
        -------
        pandas.Series
            Predictions, indexed to match `data.index`.
        """
        if self.transformer_ is None:
            raise RuntimeError("The model has not been fitted yet. Call `fit` first.")
        if not isinstance(equation_id, int):
            raise TypeError(f"`expression_id` must be an integer. Got {type(equation_id).__name__} instead.")
        if equation_id < 0 or equation_id >= len(self.equations_):
            raise IndexError(f"Invalid expression_id {equation_id}. Must be in range [0, {len(self.equations_)}-1].")
        try:
            response_var = self.transformer_.response_var
            if response_var not in data.columns:
                data = data.copy()
                data[response_var] = 0
            processed_data = self.transformer_.transform(data)
            X_new, _, groups_new = processed_data['X'], processed_data['y'], processed_data['groups']
        except Exception as e:
            raise ValueError(f"Failed to transform input data: {e}")
        try:
            equation = self.equations_[equation_id]
            predictions = equation.predict(X_new, groups_new, threshold=threshold)
            return pd.Series(predictions, index=data.index)
        except Exception as e:
            raise RuntimeError(f"Prediction failed for expression {equation_id}: {e}")
    
    def find_optimal_threshold(self, equation_id: int, data: pd.DataFrame=None, k_folds=5):
        """
        Compute the optimal decision threshold for a classification equation 
        using the True Skill Statistic (TSS) via cross-validation.

        This method evaluates the specified equation on either the provided dataset 
        or the training data used to fit the model. It calls `fit_threshold` 
        internally, which performs k-fold cross-validation to select the threshold 
        that maximizes TSS. For multi-view or grouped models, group labels are 
        preserved during threshold optimization.

        Parameters
        ----------
        data : pandas.DataFrame or None
            Input dataset to use for threshold selection. If None, the method 
            defaults to the training data used when fitting the model.
        equation_id : int
            Identifier of the equation

        Returns
        -------
        float
            Optimal threshold value for the specified classification equation, 
            selected to maximize TSS across cross-validation folds.

        Notes
        -----
        - TSS (True Skill Statistic) is computed as `sensitivity + specificity - 1`.
        - This method ensures the threshold is evaluated consistently even for 
        multi-view or grouped data.
        - The selected threshold can be used to convert probability outputs 
        into binary predictions.
        """
        if not self._is_classification_:
            print("The model isn't a classification model.")
            return None
        if self.transformer_ is None:
            raise RuntimeError("The model has not been fitted yet. Call `fit` first.")
        if not isinstance(equation_id, int):
            raise TypeError(f"`expression_id` must be an integer. Got {type(equation_id).__name__} instead.")
        if equation_id < 0 or equation_id >= len(self.equations_):
            raise IndexError(f"Invalid expression_id {equation_id}. Must be in range [0, {len(self.equations_)}-1].")
        
        if data is None:
            X, y, groups = self._X_train_, self._y_train_, self._groups_train_
        else:
            try:
                processed_data = self.transformer_.transform(data)
                X, y, groups = processed_data['X'], processed_data['y'], processed_data['groups']
            except Exception as e:
                raise ValueError(f"Failed to transform input data: {e}")
        try:
            equation: ClassificationEquation = self.equations_[equation_id]
            threshold = equation.fit_threshold(X, y, groups, k_folds=5)
            return threshold
        except Exception as e:
            raise RuntimeError(f"Couldn't find the threshold failed for expression {equation_id}: {e}")
        
        
    def get_pareto_front(self, metrics: list = [], equation_ids: list = None, pareto_only: bool = True, pareto_metric: str = None) -> pd.DataFrame:
        """
        Returns a filtered view of the symbolic regression results table.

        This function provides a convenient way to inspect the discovered equations,
        by default restricting the output to the Pareto-optimal subset. The Pareto front
        consists of equations that are not outperformed by any simpler alternative on the
        selected performance metric.

        Parameters
        ----------
        metrics : list of str, optional
            Metric columns to include in the results table.
            Each entry must match a column name in `regressor.results_`.
            To see allowed metrics, see description of SymbolicRegressor class.
        equation_ids : list of int, optional
            Specific equation IDs to include. If None, all equations are shown.
        pareto_only : bool, default=True
            If True, return only equations on the Pareto front.
        pareto_metric : str, optional
            Performance metric to use for Pareto front calculation
            (e.g., ``'Loss (cv)'``, ``'R2 (test)'``). If pareto_metric is None, the default is 
            ``'Loss (cv)'`` if that column is available; otherwise, ``'Loss'`` is used.

        Returns
        -------
        pandas.DataFrame
            Copy of the full or Pareto-filtered results table.
        """
        if self.results_ is None:
            raise RuntimeError("No results to view. Run equation_search() first.")
        ids = list(self.results_.index) if equation_ids is None else [i for i in equation_ids if i in self.results_.index]

        invalid_metric_names = []
        if metrics: 
            invalid_metric_names = [metric for metric in metrics if metric not in self.results_]
 
        if invalid_metric_names:
            raise ValueError(
                f"The following metrics aren't in the results: {', '.join(invalid_metric_names)}. "
                f"Available metrics are: {', '.join(map(str, self.results_.columns))}. "
                "Check spelling or run evaluate_equations() with those metric names."
            )
        
        # Start from the base df restricted to the user-requested ids
        base_df = self.results_.loc[self.results_.index.intersection(ids)]
        ids = base_df.index.tolist()

        if metrics:
            cols_to_return = ['Equation', 'Complexity'] + metrics
        else:
            cols_to_return = self.results_.columns.to_list()


        def make_df_for_return(df: pd.DataFrame):
            available_ids = [i for i in ids if i in df.index]
            if not available_ids:
                # Return an empty DataFrame with the expected columns and ID index
                empty = pd.DataFrame(columns=['ID'] + cols_to_return).set_index('ID')
                return empty
            out = df.loc[available_ids].reset_index()[['ID'] + cols_to_return].set_index('ID')
            return out

        if not pareto_only:
            return make_df_for_return(base_df)
        
        if pareto_metric is None:
            # Infer the primary performance metric from column names
            if 'Loss (cv)' in self.results_.columns: pareto_metric = 'Loss (cv)'
            elif 'Loss' in self.results_.columns: pareto_metric = 'Loss'
            else: raise ValueError("Cannot determine primary metric for Pareto front. Please specify using the 'metric' argument.")

        if pareto_metric not in self.results_.columns:
            raise ValueError(f"Metric '{pareto_metric}' not found in results. Please check the spelling or run evaluate_equations() for that metric first.")

        metric_name = pareto_metric
        if any(metric_name.startswith(k) for k in losses_dict) or metric_name.startswith('Loss'):
            y_higher_is_better = False
        else:
            matches = [k for k in METRIC_PROPERTIES if metric_name.startswith(k)]
            if not matches:
                raise ValueError(f"Unknown metric name: {metric_name}.")
            y_higher_is_better = METRIC_PROPERTIES[matches[0]] == 'maximize'
        pareto_df = compute_pareto_front(base_df, 'Complexity', pareto_metric, False, y_higher_is_better)
        return make_df_for_return(pareto_df).sort_values(by=["Complexity", "Equation"], ascending=[True, True])
    
    def print_equation_details(self, equation_id: int):
        """
        Display a detailed, group-by-group view of a fitted equation.

        Shows the full, un-normalized symbolic expression for a specific equation,
        substituting the fitted parameters and normalization constants for each
        group in the data.

        **Coefficient notation**
        - **β** : Global coefficient, fitted once on the entire dataset
        (shared across all groups).
        - **u** : Group-specific coefficient, fitted separately within each group.

        Equations may contain only βs (purely global), only us (purely group-specific),
        or a mix of both (analogous to mixed-effects models).

        Parameters
        ----------
        equation_id : int
            ID of the equation to display.
        """
        if self.transformer_ is None:
            raise RuntimeError("The model has not been fitted yet. Call `fit` first.")
        if equation_id < 0 or equation_id >= len(self.equations_):
            raise IndexError(f"Invalid expression_id {equation_id}. Must be in range [0, {len(self.equations_)-1}].")
            
        equation: Equation = self.equations_[equation_id]
        print(f"Details for Equation ID: {equation_id}\n" + "-"*40)
        print(f"Equation: {equation.readable_skeleton}\n" + "-"*40)
        print("Full Expression per Group:")

        unique_groups = sorted(list(self.transformer_.stats.keys()))
        for group in unique_groups:
            full_expr = self.transformer_.untransform_equation(equation, group)
            print(f"\nGroup '{group}':")
            sp.pprint(full_expr, use_unicode=True, wrap_line=True)

    def plot_pareto_front(self, x_axis: str = 'Complexity', y_axis: str = None, **plot_kwargs):
        """
        Plot the Pareto front of the fitted equations.

        Visualizes the trade-off between two metrics, highlighting the set of
        optimal, non-dominated solutions.

        Parameters
        ----------
        x_axis: str, optional
            The metric for the x-axis. Defaults to 'Complexity'.
        y_axis : str, optional
            The metric for the y-axis. If y_axis is None, the default is 
            ``'Loss (cv)'`` if that column is available; otherwise, ``'Loss'`` is used. Defaults to None.
        **plot_kwargs: Any
            Additional keyword arguments passed to matplotlib.
        """
        if self.results_ is None: raise RuntimeError("No results to plot. Run equation_search() first.")
        
        if y_axis is None: # Infer default y-axis
            if 'Loss (cv)' in self.results_.columns: y_axis = 'Loss (cv)'
            elif 'Loss' in self.results_.columns: y_axis = 'Loss'
            else: raise ValueError("Cannot determine y-axis. Please score a metric or specify 'y_axis'.")
        
        def lookup_metric(name):
            matches = [k for k in METRIC_PROPERTIES if name.startswith(k)]
            if not matches:
                raise ValueError(f"Unknown metric name: {name}.")
            return matches[0]
        y_higher_is_better = METRIC_PROPERTIES[lookup_metric(y_axis)] == 'maximize'
        if x_axis == "Complexity":
            x_direction = "minimize"
        else:
            x_direction = METRIC_PROPERTIES[lookup_metric(x_axis)]
        x_higher_is_better = x_direction == 'maximize'
        plot_pareto_front(
            [(self._engine_name, self.results_)], x_axis, y_axis, x_higher_is_better, y_higher_is_better, **plot_kwargs
            )
    
    def save(self, filepath: str) -> None:
        """
        Saves the entire SymbolicRegressor instance to a file using pickle.

        This method serializes the object, including the SR engine, the data
        transformer, all discovered equations, and the results table, allowing
        an analysis to be paused and resumed.

        Parameters
        ----------
        filepath : str
            Path to the output file (e.g., ``"my_model.pkl"``).
        """
        import pickle
        try:
            with open(filepath, "wb") as f:
                pickle.dump(self, f)
            print(f"SymbolicRegressor instance successfully saved to {filepath}")
        except Exception as e:
            print(f"Error saving model to {filepath}: {e}")

    @classmethod
    def load(cls, filepath: str):
        """
        Loads a SymbolicRegressor instance from a pickle file.

        This class method deserializes a previously saved instance, restoring its
        entire state to continue an analysis.

        Parameters
        ----------
        filepath : str
            Path to the saved model file (e.g., ``"my_model.pkl"``).

        Returns
        -------
        SymbolicRegressor
            The loaded instance of the SymbolicRegressor.
        """
        import pickle
        try:
            with open(filepath, "rb") as f:
                instance = pickle.load(f)
            
            if not isinstance(instance, cls):
                raise TypeError(f"The file at {filepath} does not contain a valid SymbolicRegressor instance.")
            
            print(f"SymbolicRegressor instance loaded successfully from {filepath}")
            return instance
        except FileNotFoundError:
            print(f"Error: No file found at {filepath}")
            return None
        except Exception as e:
            print(f"Error loading model from {filepath}: {e}")
            return None

    def copy(self):
        """
        Creates a deep copy of the SymbolicRegressor instance.

        This is useful for creating a snapshot of the regressor's state after
        the equation search, allowing for different parameter fitting strategies
        to be explored from the same starting point.

        Returns
        -------
        SymbolicRegressor
            A new, independent instance with the same state.
        """
        return copy.deepcopy(self)
    
    @classmethod
    def merge(cls, regressors: List['SymbolicRegressor']):
        """
        Merge multiple `SymbolicRegressor` instances into a single object.

        Combines unique equations from different runs or engines, tracking which
        engine(s) found each equation.

        Notes
        -----
        - All regressors must have been created with a compatible formula and
        normalization scheme.
        - The resulting regressor inherits its transformer (normalization and
        preprocessing) from the first regressor in the list.
        - All fitted parameter values are discarded; you must run
        `fit_model_parameters` again.

        Parameters
        ----------
        regressors : list of SymbolicRegressor
            A list of instances to merge.

        Returns
        -------
        SymbolicRegressor
            A new instance containing the merged results.
        """
        if not regressors:
            raise ValueError("The list of regressors to merge cannot be empty.")

        # Use the first regressor as the template for the new instance
        template = regressors[0]
        merged_regressor = template.copy()
        
        all_equations = []
        all_rows = []

        for reg in regressors:
            if not reg.equations_:
                warnings.warn(f"Skipping a regressor with engine '{reg._engine.name}' as it has no equations.")
                continue

            for eq_id, row in reg.results_.iterrows():
                equation = reg.equations_[eq_id]

                if equation not in all_equations:
                    # This is a new, unique equation
                    new_eq = copy.deepcopy(equation)
                    temp_id = len(all_equations)
                    all_equations.append(new_eq)

                    # Store a row referencing that temp_id
                    all_rows.append({
                        "temp_id": temp_id,
                        "Equation": row["Equation"],
                        "Complexity": row["Complexity"],
                        "Engines": [reg._engine_name],
                    })
                else:
                    # This equation was found before, update the engine list
                    existing_idx = all_equations.index(equation)
                    all_rows[existing_idx]['Engines'].append(reg._engine_name)

        # Re-index the merged results
        merged_results = pd.DataFrame(all_rows)
        if merged_results.empty:
            merged_regressor.equations_ = []
            merged_regressor.results_ = pd.DataFrame()
            return merged_regressor
        
        # Sort by Complexity (ascending). Use Equation as deterministic tiebreaker.
        merged_results = merged_results.sort_values(by=["Complexity", "Equation"], ascending=[True, True]).reset_index(drop=True)

        # Assign final IDs from sorted order and map temp_id to final ID
        merged_results["ID"] = merged_results.index
        id_map = dict(zip(merged_results["temp_id"], merged_results["ID"]))

        # Update Equation objects with final IDs and build an ordered list by ID
        for temp_idx, eq in enumerate(all_equations):
            new_id = int(id_map[temp_idx])
            eq.id = new_id
        
        # Build list ordered by final ID
        ordered_eqs = [None] * len(all_equations)
        for eq in all_equations:
            ordered_eqs[eq.id] = eq

        merged_regressor.equations_ = ordered_eqs
        # Finalize results DataFrame and set ID as index
        final_df = merged_results[["ID", "Equation", "Complexity", "Engines"]].set_index("ID")
        merged_regressor.results_ = final_df

        return merged_regressor
    
    def permanently_exclude_equations(self, equation_ids: List[int]):
        """
        Permanently remove specified equations from the regressor instance.

        Useful for curating the set of models before further analysis, evaluation,
        or exporting.

        Parameters
        ----------
        equation_ids : list of int
            IDs of the equations to remove.
        """
        ids_to_remove = set(equation_ids)
        if not ids_to_remove: 
            print("No equations to remove")
            return
        
        # Filter equations and results
        self.equations_ = [eq for eq in self.equations_ if eq.id not in ids_to_remove]
        self.results_ = self.results_[~self.results_.index.isin(ids_to_remove)].copy()
        
        # Re-index everything to maintain consistency
        self.results_.reset_index(drop=True, inplace=True)
        self.results_['ID'] = self.results_.index
        self.results_.set_index('ID', inplace=True)
        for i, eq in enumerate(self.equations_):
            eq.id = i
        
        print(f"Removed {len(ids_to_remove)} equations. {len(self.equations_)} remain.")

def df_to_latex(
    df,
    filename=None,
    index=False,
    escape=True,
    float_format="%.4f",
    bold_max=False,
    bold_min=False,
    caption="",
    label="tab:",
    position="h!"
) -> str:
    """
    Convert a DataFrame to LaTeX with optional styling.

    Parameters
    ----------
    df : pandas.DataFrame
        Input DataFrame.
    filename : str, optional
        If provided, saves LaTeX output to this file.
    index : bool, default=False
        Whether to include the DataFrame index.
    escape : bool, default=True
        Escape special LaTeX characters.
    float_format : str, default="%.4f"
        Format string for floating-point numbers.
    bold_max : bool, default=False
        Bold the maximum value in each column.
    bold_min : bool, default=False
        Bold the minimum value in each column.
    caption : str, default=""
        Table caption.
    label : str, default="tab:"
        LaTeX label for referencing.
    position : str, default="h!"
        LaTeX table placement specifier.

    Returns
    -------
    str
        LaTeX table code.

    Examples
    --------
    >>> latex_code = df_to_latex(df, bold_max=True, caption="Results")
    >>> print(latex_code)
    """
    # Style the DataFrame (bold max/min if requested)
    styled_df = df.copy()
    if bold_max or bold_min:
        def highlight_extrema(x):
            x = x.astype(float)
            styles = pd.Series("", index=x.index)
            if bold_max:
                styles[x.idxmax()] = "\\textbf{%.4f}" % x.max()
            if bold_min:
                styles[x.idxmin()] = "\\textbf{%.4f}" % x.min()
            return styles
        styled_df = styled_df.style.apply(highlight_extrema)
    
    # Convert to LaTeX
    latex_code = styled_df.to_latex(
        index=index,
        escape=escape,
        float_format=float_format,
        caption=caption,
        label=label,
        position=position
    )
    
    # Save to file if requested
    if filename:
        with open(filename, "w") as f:
            f.write(latex_code)
    
    return latex_code