from .base import SymbolicRegressionEngine
import numpy as np
import pandas as pd
from typing import List, Sequence
import sympy as sp
import re

def extract_vars(expr):
    # grabs content inside f(...)
    inside = re.search(r"f\((.*)\)", expr).group(1)

    # separate by pipe
    parts = [p.strip() for p in inside.split("|")]

    def split_args(s):
        # captures blocks with parenthesis or simple tokens
        tokens = re.findall(r'\([^()]*\)|[^,]+', s)
        return [t.strip() for t in tokens]

    args = split_args(parts[0])

    return args

class ManualEngine(SymbolicRegressionEngine):
    """
    An engine that evaluates a fixed set of user-specified equations rather than
    searching for new symbolic expressions. This is useful when the researcher
    wants to test explicit hypotheses, benchmark known functional forms, or
    compare candidate models derived from theory against data. 
    """
    def __init__(self, backend_kwargs, wrapper_kwargs):
        """
        equation_list : list of str or callable
            A collection of candidate equations to be evaluated. Each entry can be
            provided as a string formula (e.g. "u1 + u2*log(A)*T"). 
            This argument is required for the ManualEngine, since no equations
            are generated automatically.
        """
        self.equation_list = backend_kwargs.get("equation_list", [])
        formula: str = backend_kwargs.get("formula", "")
        variables = extract_vars(formula)
        self.replacement_map = {var: f'x{i+1}' for i, var in enumerate(variables)}
        print(self.replacement_map)
        self.variable_names = [f"x{i}" for i in range(1, len(variables)+1)]
    
    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        return
    
    def fit_multi(self, Xs: List[np.ndarray], ys: List[np.ndarray]) -> None:
        return
    
    def get_expressions(self) -> pd.Series:
        equations = [str(sp.sympify(eq).subs(self.replacement_map)) for eq in self.equation_list]
        print(equations)

        return pd.Series(equations)
    
    @property
    def engine_variable_names(self) -> Sequence[str]: 
        return self.variable_names