--------------------------------------------
-- Differentiation of univariate expressions
-- Fritz Ruehr, Willamette LLC class
--------------------------------------------

-------- Exprs over literals, binary operators and a single variable

data Expr a b = Lit a | Bop b (Expr a b) (Expr a b) | Var deriving (Show, Eq)

fold f g x (Lit a)     = f a
fold f g x (Bop o l r) = g o (fold f g x l) (fold f g x r)
fold f g x  Var        = x

-------- Arithmetic operators: syntax and semantics

data AOp = Add | Mul | Sub deriving (Show, Eq)

sym Add = "+" ;  sym Mul = "*" ;  sym Sub = "-"
sem Add = (+) ;  sem Mul = (*) ;  sem Sub = (-)

-------- Evaluation and printing for exprs

eval = flip (fold id sem)

prtm :: Show a => Expr a AOp -> String
prtm = fold show (fmt . sym) "x"

-------- Differentiation (for exprs and polys)

diff (Lit _) = Lit 0
diff  Var    = Lit 1
diff (Bop Add a b) = Bop Add (diff a) (diff b)
diff (Bop Sub a b) = Bop Sub (diff a) (diff b)
diff (Bop Mul a b) = Bop Add (Bop Mul a (diff b)) (Bop Mul b (diff a))

-- EXPRESS DIFF AS A FOLD? NO, need access to original args (as well as diff's)

dp cs = zipWith (*) (tail cs) [1..]

simp = fix (push s)
       where s (Bop p (Bop q t (Lit k)) (Lit j)) | p == q = Bop p t (Lit (sem q k j))
             s (Bop op (Lit k) (Lit j)) = Lit (sem op k j)

             s (Bop Add t (Lit 0)) = t
             s (Bop Add (Lit 0) t) = t
             s (Bop Mul t (Lit 0)) = Lit 0
             s (Bop Mul (Lit 0) t) = Lit 0
             s (Bop Mul t (Lit 1)) = t
             s (Bop Mul (Lit 1) t) = t
             s t = t

-- utilities (use "k n" on a term to see its nth derivative)

deriv n = prtm . simp . pow n diff

pow 0 f x = x
pow n f x = pow (n-1) f (f x)

push f (Bop op a b) = f (Bop op (push f a) (push f b))
push f t = t

fix f x = if x == f x then x else fix f (f x)

-------- Evaluation and semantic operations on polynomials
-------- (represented as reversed lists of coefficients)

pval cs x = foldr (\c r -> c + x * r) 0 cs

ps = lzw (+) 0
pd = lzw (-) 0

pm cs ds = foldr (\c r -> ps (map (c *) ds) (0 : r)) [] cs

-------- Conversion from exprs to polynomials
-------- {theorem:  eval  ==  pval . poly }

psem Add = ps ;  psem Mul = pm ;  psem Sub = pd

poly = fold (:[]) psem [0,1]

-------- Convenience functions, sample exprs and testing

(+:) = Bop Add ;  (*:) = Bop Mul ;  (-:) = Bop Sub

t1 = Var -: Lit 2
t2 = ((Var *: Var) +: (Var +: Lit 5)) +: t1
t3 = t2 *: (t2 +: t1)

test t k = all (\x-> eval t x == pval (poly t) x) [1..k]

-------- Utility functions ("long" zipWith, infix format)

lzw f u []     []     = []
lzw f u (x:xs) []     = f x u : lzw f u xs []
lzw f u []     (y:ys) = f u y : lzw f u [] ys
lzw f u (x:xs) (y:ys) = f x y : lzw f u xs ys

fmt o l r = concat ["(",l,o,r,")"]