# /// script
# dependencies = [
#     "marimo==0.24.0",
#     "symbolica==3.0.0",
# ]
# requires-python = ">=3.14"
# ///

import marimo

__generated_with = "0.24.0"
app = marimo.App(width="full", app_title="Numerical kernels")


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Creating fast numerical evaluators
    """)
    return


@app.cell
def _():
    import marimo as mo
    from symbolica import S, Expression, FunctionDefinition, Symbol, Float, ComplexFloat

    return ComplexFloat, Expression, Float, FunctionDefinition, S, Symbol, mo


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Construction

    Build an example expression that is nested and contains built-in constants:

    \begin{align*}
    e &= x + \pi + \cos(x) + f(g(x+1),2x)\\
    f(y,z) &=y^2 + z^2 y^2\\
    g(y) &=y + 5
    \end{align*}
    """)
    return


@app.cell
def _(FunctionDefinition, S, Symbol):
    x, y, z, f, g = S('x', 'y', 'z', 'f', 'g')

    expr = x + Symbol.PI + x.cos() + f(g(x+1), x*2)
    f_body = y**2 + z**2*y**2
    g_body = y + 5

    ev = expr.evaluator([x], functions=[
        FunctionDefinition(f, [y, z], f_body),
        FunctionDefinition(g, [y], g_body),
    ])
    return ev, expr, f, f_body, g, g_body, x, y, z


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Instructions

    The evaluator has created a linear set of instructions (inlining `f` and `g`), removed common subexpressions and uses near-minimal temporary storage.
    """)
    return


@app.cell
def _(ev):
    instr = ev.get_instructions()
    for i in instr.instructions:
        print(i)

    print("Number of operations: ", len(instr.instructions))
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ### Fast evaluation
    You can now evaluate the expression for a given value of `x`. Specifying multiple `x` will use a SIMD backend if available. On non-WASM builds, you can generate fast C++ code with inline assembly or JIT-compile the evaluator.
    """)
    return


@app.cell
def _(ev):
    ev.evaluate([1., 2., 3.])
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    Evaluate with double-double precision (faster than arbitrary precision):
    """)
    return


@app.cell
def _(Float, ev):
    ev.evaluate_with_prec([Float(1, precision=32)], 32)
    return


@app.cell
def _(Float, ev):
    ev.evaluate_with_prec([Float(1, decimal_digits=100)], 100)
    return


@app.cell
def _(ev):
    ev.evaluate_complex([1+1j])
    return


@app.cell
def _(ComplexFloat, Float, ev):
    ev.evaluate_complex_with_prec([ComplexFloat(Float(1, decimal_digits=100), Float(1, decimal_digits=100))], 100)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Multiple expression evaluation

    Multiple expression can be jointly optimized. They will share common subexpressions, and therefore can be evaluated faster than when using two separate evaluators.
    """)
    return


@app.cell
def _(Expression, FunctionDefinition, expr, f, f_body, g, g_body, x, y, z):
    ev_m = Expression.evaluator_multiple([expr, expr + 2], [x], functions=[
        FunctionDefinition(f, [y, z], f_body),
        FunctionDefinition(g, [y], g_body),
    ])
    return (ev_m,)


@app.cell
def _(ev_m):
    instr_m = ev_m.get_instructions()
    for i_m in instr_m.instructions:
        print(i_m)

    print("Number of operations: ", len(ev_m.get_instructions().instructions))
    return


@app.cell
def _(ev_m):
    ev_m.evaluate([1.])
    return


if __name__ == "__main__":
    app.run()