# Create data
X = np.random.randn(1000, 5)
y = np.sin(X[:, 0] + X[:, 3] + X[:, 4]) * X[:, 2]**2

# Define template: we want f(x1, x4, x5) + g(x3)
# We'll obtain independent expressions for f and g 
# where the respective variables of each are expressed
# as #1 for the first variable, #2 for the second, 
# and so on.
# While the output can be evaluated, it is NOT compatible with
# the exporting features for typesetting and modeling in other libraries.
# 
template_py = TemplateExpressionSpec(
    function_symbols=["f", "g"],
    combine="""((; f, g), (a, b, c, d, e)) -> let
    _f = f(a, d, e) 
    _g = g(c)
    _out = _f + _g
    end""",
    num_features={"f": 2, "g": 1},
)

model = PySRRegressor(
    expression_spec=template_py,
    binary_operators=["+", "*", "-", "/"],
    unary_operators=["sin"],
    model_selection="score",
    maxsize=25,
    tournament_selection_n=15,
    tournament_selection_p=0.9,
    niterations=50,
    populations=100,
    population_size=80,
    should_simplify=True,
    topn=10,
    # annealing=True,
    # alpha=2,
    #full_objective=jl.seval(objective),
)
model.fit(X, y)

template_jl = TemplateExpressionSpec(
    function_symbols=["f", "g"],
    combine="""((; f, g), (a, b, c, d, e)) -> let
    _f = f(a, d, e) 
    _g = g(c)
    _out = _f + _g
    end""",
    num_features=jl.seval("""Base.Dict(Symbol("f") => 3, Symbol("g") => 1)"""),
)