app "deriv"
    packages { base: "platform" }
    imports [base.Task]
    provides [ main ] to base

# based on: https://github.com/koka-lang/koka/blob/master/test/bench/haskell/deriv.hs

IO a : Task.Task a []

main : IO {}
main =
    Task.after Task.getInt \n ->
        x : Expr
        x = Var "x"

        f : Expr
        f = pow x x

        nest deriv n f # original koka n = 10
            |> Task.map (\_ -> {})

Expr : [ Val I64, Var Str, Add Expr Expr, Mul Expr Expr, Pow Expr Expr, Ln Expr ]

divmod : I64, I64 -> Result { div: I64, mod: I64 } [ DivByZero ]*
divmod = \l, r ->
    when Pair (l // r) (l % r) is
        Pair (Ok div) (Ok mod) -> Ok { div, mod }
        _ -> Err DivByZero

pown : I64, I64 -> I64
pown = \a, n ->
    when n is
        0 -> 1
        1 -> a
        _ ->
            when divmod n 2 is
                Ok { div, mod } ->
                    b = pown a div
                    b * b * (if mod == 0 then 1 else a)

                Err DivByZero -> -1

add : Expr, Expr -> Expr
add = \a,b ->
    when Pair a b is
        Pair (Val n)   (Val m)         -> Val (n + m)
        Pair (Val 0)   f               -> f
        Pair f         (Val 0)         -> f
        Pair f         (Val n)         -> add (Val n) f
        Pair (Val n)   (Add (Val m) f) -> add (Val (n+m)) f
        Pair f         (Add (Val n) g) -> add (Val n) (add f g)
        Pair (Add f g) h               -> add f (add g h)
        Pair f         g               -> Add f g

mul : Expr, Expr -> Expr
mul = \a,b ->
    when Pair a b is
        Pair (Val n)   (Val m)         -> Val (n*m)
        Pair (Val 0)   _               -> Val 0
        Pair _         (Val 0)         -> Val 0
        Pair (Val 1)   f               -> f
        Pair f         (Val 1)         -> f
        Pair f         (Val n)         -> mul (Val n) f
        Pair (Val n)   (Mul (Val m) f) -> mul (Val (n*m)) f
        Pair f         (Mul (Val n) g) -> mul (Val n) (mul f g)
        Pair (Mul f g) h               -> mul f (mul g h)
        Pair f         g               -> Mul f g

pow : Expr, Expr -> Expr
pow = \a,b ->
    when Pair a b is
        Pair (Val m) (Val n) -> Val (pown m n)
        Pair _       (Val 0) -> Val 1
        Pair f       (Val 1) -> f
        Pair (Val 0) _       -> Val 0
        Pair f       g       -> Pow f g

ln : Expr -> Expr
ln = \f ->
    when f is
        Val 1 -> Val 0
        _ -> Ln f

d : Str, Expr -> Expr
d = \x, expr ->
    when expr is
        (Val _)   -> Val 0
        (Var y)   -> if x == y then Val 1 else Val 0
        (Add f g) -> add (d x f) (d x g)
        (Mul f g) -> add (mul f (d x g)) (mul g (d x f))
        (Pow f g) -> mul (pow f g) (add (mul (mul g (d x f)) (pow f (Val (-1)))) (mul (ln f) (d x g)))
        (Ln f)    -> mul (d x f) (pow f (Val (-1)))

count : Expr -> I64
count = \expr ->
    when expr is
        (Val _) -> 1
        (Var _) -> 1
        (Add f g) -> count f + count g
        (Mul f g) -> count f + count g
        (Pow f g) -> count f + count g
        (Ln f)    -> count f

nest : (I64, Expr -> IO Expr), I64, Expr -> IO Expr
nest = \f, n, e -> nestHelp n f n e

nestHelp : I64, (I64, Expr -> IO Expr), I64, Expr -> IO Expr
nestHelp = \s, f, m, x ->
    when m is
        0 -> Task.succeed x
        _ -> f (s - m) x |> Task.after \w -> nestHelp s f (m-1) w


deriv : I64, Expr -> IO Expr
deriv = \i, f ->
    fprime = d "x" f
    line =
        Str.fromInt (i + 1)
            |> Str.concat " count: "
            |> Str.concat (Str.fromInt (count fprime))

    Task.putLine line
        |> Task.after \_ -> Task.succeed fprime