# lambdify an expression



## Mapping of SymPy Values into julia values
val_map = Dict(
               "Zero"             => :(0),
               "One"              => :(1),
               "NegativeOne"      => :(-1),
               "Half"             => :(1/2),
               "Pi"               => :pi,
               "Exp1"             => :ℯ,
               "Infinity"         => :Inf,
               "NegativeInfinity" => :(-Inf),
               "ComplexInfinity"  => :Inf, # error?
               "ImaginaryUnit"    => :im,
               "BooleanTrue"      => :true,
               "BooleanFalse"     => :false
               )

## Mapping of Julia function names into julia ones
## most are handled by Symbol(fnname), this catches exceptions
_heaviside(x) = 1//2 * (1 + sign(x))
function _piecewise(args...)
    as = copy([args...])
    val, cond = pop!(as)
    ex = Expr(:call, :ifelse, cond, val, :nothing)
    while length(as) > 0
        val, cond = pop!(as)
        ex = Expr(:call, :ifelse, cond, val, ex)
    end
    ex
end

## Hack to avoid Expr(:call,  :*,2, x)  being  2x and  not  2*x
__prod__(args...) =  prod(args)
export __prod__

fn_map = Dict(
              "Add" => :+,
              "Sub" => :-,
              "Mul" => :__prod__,
              "Div" => :/,
              "Pow" => :^,
              "re"  => :real,
              "im"  => :imag,
              "Abs" => :abs,
              "Min" => :min,
              "Max" => :max,
              "Poly" => :identity,
              "Heaviside" => :(_heaviside),
              "Piecewise" => :(_piecewise),
              "And" => :(&),
              "Or" => :(|),
              "Less" => :(<),
              "LessThan" => :(<=),
              "StrictLessThan" => :(<),
              "Equal" => :(==),
              "Equality" => :(==),
              "Unequality" => :(!==),
              "StrictGreaterThan" => :(>),
              "GreaterThan" => :(>=),
"Greater" => :(>),
"conjugate" => :conj
              )

map_fn(key, fn_map) = haskey(fn_map, key) ? fn_map[key] : Symbol(key)

Base.convert(::Type{Expr}, x::SymbolicObject) = walk_expression(x)

function walk_expression(ex; values=Dict(), fns=Dict())

    fns_map = merge(fn_map, fns)
    vals_map = merge(val_map, values)

    fn = Introspection.funcname(ex)

    # special case `F(t) = ...` output from ODE
    # this may be removed if it proves a bad idea....
    if fn == "Equality" && lhs(ex).is_Function
        return walk_expression(rhs(ex), values=values, fns=fns)
    end

    if fn == "Symbol" || fn == "Dummy"
        str_ex = string(ex)
        return get(vals_map, str_ex, Symbol(str_ex))
    elseif fn in ["Integer" , "Float"]
        return N(ex)
    elseif fn == "Rational"
        return convert(Int, numer(ex))//convert(Int, denom(ex))
        ## piecewise requires special treatment
    elseif fn == "Piecewise"
        return _piecewise([walk_expression(cond, values=values, fns=fns) for cond in Introspection.args(ex)]...)
    elseif fn == "ExprCondPair"
        val, cond = Introspection.args(ex)
        return (val, walk_expression(cond, values=values, fns=fns))
    elseif fn == "Tuple"
        return walk_expression.(Introspection.args(ex), values=values, fns=fns)
    elseif haskey(vals_map, fn)
        return vals_map[fn]
    end

    as = Introspection.args(ex)

    Expr(:call, map_fn(fn, fns_map), [walk_expression(a, values=values, fns=fns) for a in as]...)
end

"""
     lambdify(ex, vars; typ, fns, values, use_julia_code=false, invoke_latest=true)

Take a symbolic expression and returns an anonymous `Julia` function.

Converts from a SymPy object to a `Julia` expression by walking the SymPy expression tree and converting each step, (e.g. essentially calls `convert(Expr,  ex)`).

Then creates a function. The function arguments are based on `free_symbols` and its ordering, unless `vars` is specified directly.

* `use_julia_code=true` will use SymPy's conversion to an expression, the default is `false`

* `invoke_latest=true` calls `Base.invokelatest` to work around world age issues. This is the safe default, but setting to `false` will result in faster-executing functions.

Example:

```jldoctest
julia> using SymPy

julia> @vars x y z
(x, y, z)

julia> ex = x^2 * sin(x)
 2       
x ⋅sin(x)

julia> fn = lambdify(ex)
#87 (generic function with 1 method)

julia> fn(pi)
1.2086779438644711e-15

julia> ex = x + 2y + 3z
x + 2⋅y + 3⋅z

julia> fn = lambdify(ex)
#87 (generic function with 1 method)

julia> fn(1,2,3) # order is by free_symbols
14

julia> ex(x=>1, y=>2, z=>3)
14

julia> fn = lambdify(ex, (y,x,z))
#87 (generic function with 1 method)

julia> fn(1,2,3)
13
```


"""
function lambdify(ex::Sym, vars=free_symbols(ex);
                  typ=Any, fns=Dict(), values=Dict(),
                  use_julia_code=false,
                  invoke_latest=true
                  )
    # if :julia_code printer is there, use it
    if use_julia_code
        body = Meta.parse(sympy.julia_code(ex)) # issue here with 2.*...
    else
        body = walk_expression(ex, fns=fns, values=values)
    end

    try
        syms = typ == Any ? map(Symbol,vars) : map(s->Expr(:(::),s,typ), Symbol.(vars))
        fn = eval(Expr(:function, Expr(:call, gensym(), syms...), body))
        if invoke_latest
            (args...) -> Base.invokelatest(fn, args...) # https://github.com/JuliaLang/julia/pull/19784
        else
            fn
        end
    catch err
        throw(ArgumentError("Expression does not lambdify"))
    end
end


function _lambdify(ex::Sym, vars=free_symbols(ex))
    body = convert(Expr, ex)
    syms = Symbol.(vars)
    fn = eval(Expr(:function, Expr(:call, gensym(), syms...), body))
    fn
end



# from @mistguy cf. https://github.com/JuliaPy/SymPy.jl/issues/218
# T a data type to convert to, when specified
function lambdify(exs::Array{S, N}, vars = union(free_symbols.(exs)...); T::DataType=Nothing, kwargs...) where {S <: Sym, N}
    f = lambdify.(exs, (vars,)) # prevent broadcast in vars
    if T == Nothing
        (args...) -> map.(f, args...)
    else
        (args...) -> convert(Array{T,N}, map.(f, args...))
    end
end

export(lambdify)

Base.convert(::Type{Function}, ex::Sym) = lambdify(ex)
