{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
{-# LANGUAGE StandaloneDeriving, UndecidableInstances #-}

module Morphisms where

data Expr a = Const Int
            | Add a a
            | Sub a a
            | Mul a a
            | Div a a
            deriving (Show, Eq, Functor, Foldable, Traversable)

{-- equivalent to:
instance Functor Expr where
  fmap _ (Const x) = Const x
  fmap f (Add x y) = Add (fmap f x) (fmap f y)
  fmap f (Sub x y) = Sub (fmap f x) (fmap f y)
  fmap f (Mul x y) = Mul (fmap f x) (fmap f y)
  fmap f (Div x y) = Div (fmap f x) (fmap f y)
--}

newtype Term f = In { out :: f (Term f) }

-- this requires undecidableinstances
deriving instance Show (f (Term f)) => Show (Term f)

bottomUp :: Functor a => (Term a -> Term a) -> Term a -> Term a
bottomUp f = f . In . fmap (bottomUp f) . out
-- bottomUp f t = f $ In $ fmap (bottomUp f) $ out t

topDown :: Functor a => (Term a -> Term a) -> Term a -> Term a
topDown f = In . fmap (topDown f) . out . f
-- topDown f t = In $ fmap (topDown f) $ out $ f t

example = In (Mul
               (In (Add
                 (In (Const 7))
                 (In (Const 9))))
               (In (Sub
                 (In (Const 3))
                 (In (Div (In (Const 9)) (In (Const 2)))))))

plusTwo = bottomUp f example where
  f (In (Const x)) = In $ Const $ x + 2
  f x = x

addMul = topDown f example where
  f (In (Add x y)) = In (Mul x y)
  f x = x

-- a generalized fold; like bottomUp but without the In
foldy :: Functor f => (f a -> a) -> Term f -> a
foldy f = f . fmap (foldy f) . out

countNodes :: Expr Int -> Int
countNodes (Const _) = 1
countNodes (Add x y) = 1 + x + y
countNodes (Sub x y) = 1 + x + y
countNodes (Mul x y) = 1 + x + y
countNodes (Div x y) = 1 + x + y

countExample = foldy countNodes example

-- in the jargon, an f a -> a is an Algebra and foldy is a catamorphism
-- cata == downwards/collapse

type Algebra f a = f a -> a

cata :: Functor f => Algebra f a -> Term f -> a
cata f = f . fmap (cata f) . out

pp :: Expr String -> String
pp (Const n) = show n
pp (Add x y) = "(" ++ x ++ " + " ++ y ++ ")"
pp (Sub x y) = "(" ++ x ++ " - " ++ y ++ ")"
pp (Mul x y) = "(" ++ x ++ " x " ++ y ++ ")"
pp (Div x y) = "(" ++ x ++ " / " ++ y ++ ")"

ppExample = cata pp example

-- n.b. bottomUp f == cata (f . In)
bottomUp' :: Functor a => (Term a -> Term a) -> Term a -> Term a
bottomUp' = cata . (. In)

plusTwo' = bottomUp' f example where
  f (In (Const x)) = In $ Const $ x + 2
  f x = x

-- a generalized unfold, like topDown without the out
unfoldy :: Functor f => (a -> f a) -> a -> Term f
unfoldy f = In . fmap (unfoldy f) . f

-- you guessed it: we flipped the arrows, so:

type CoAlgebra f a = a -> f a

-- ana = building
ana :: Functor f => CoAlgebra f a -> a -> Term f
ana f = In . fmap (ana f) . f

-- we need an extra bit of state for "expand me" vs "const me"
makeFactorial :: Int -> Term Expr
makeFactorial n = ana f (n, True) where
  f :: (Int, Bool) -> Expr (Int, Bool)
  f (0, _) = Const 1
  f (n, True) = Mul (n, False) (n - 1, True)
  f (n, False) = Const n