-- | Recursion schemes, also known as scary named folds... 

{-# LANGUAGE CPP #-}
module Data.Generics.Fixplate.Morphisms where

--------------------------------------------------------------------------------

import Prelude hiding ( mapM )
import Data.Foldable
import Data.Traversable
import Data.Generics.Fixplate.Base

#ifdef WITH_QUICKCHECK
import Data.Char ( ord )
import Data.List ( intercalate )
import Test.QuickCheck
-- import Data.Generics.Fixplate.Traversals
import Data.Generics.Fixplate.Test.Tools
#endif 

--------------------------------------------------------------------------------
-- * Classic ana\/cata\/para\/hylo-morphisms

-- | A /catamorphism/ is the generalization of right fold from lists to trees.
cata :: Functor f => (f a -> a) -> Mu f -> a
cata h = go where
  go = h . fmap go . unFix

-- | A /paramorphism/ is a more general version of the catamorphism.
para :: Functor f => (f (Mu f, a) -> a) -> Mu f -> a
para h = go where
  go (Fix t) = h (fmap go' t)
  go' t = (t, go t) 

-- | Another version of 'para' (a bit less natural in some sense).
para' :: Functor f => (Mu f -> f a -> a) -> Mu f -> a
para' h = go where
  go t = h t (fmap go $ unFix t) 

-- | A list version of 'para' (compare with Uniplate)
paraList :: (Functor f, Foldable f) => (Mu f -> [a] -> a) -> Mu f -> a 
paraList f = go where
  go t = f t (toList $ fmap go $ unFix t)

-- | An /anamorphism/ is simply an unfold. Probably not very useful in this context.
ana :: Functor f => (a -> f a) -> a -> Mu f
ana h = go where
  go = Fix . fmap go . h
  -- go x = Fix (fmap go (h x))

-- | An /apomorphism/ is a generalization of the anamorphism.
apo :: Functor f => (a -> f (Either (Mu f) a)) -> a -> Mu f
apo h = go where
  go = Fix . fmap worker . h 
  worker ei = case ei of 
    Left  t -> t
    Right a -> go a

-- | A /hylomorphism/ is the composition of a catamorphism and an anamorphism.
hylo :: Functor f => (f a -> a) -> (b -> f b) -> (b -> a) 
hylo g h = cata g . ana h

--------------------------------------------------------------------------------
-- * Zygomorphisms

-- | A /zygomorphism/ is a basically a catamorphism inside another catamorphism.
-- It could be implemented (somewhat wastefully) by first annotating each subtree
-- with the corresponding values of the inner catamorphism ('synthCata'), then running 
-- a paramorphims on the annotated tree:
-- 
-- > zygo_ g h == para u . synthCata g 
-- >   where
-- >     u = h . fmap (first attribute) . unAnn
-- >     first f (x,y) = (f x, y)
--
zygo_ :: Functor f => (f b -> b) -> (f (b,a) -> a) -> Mu f -> a
zygo_ g h = snd . zygo g h  

zygo :: Functor f => (f b -> b) -> (f (b,a) -> a) -> Mu f -> (b,a)
zygo g h = go where
  go (Fix t) = (b,a) where
    b  = g (fmap fst ba)  -- :: b
    a  = h ba             -- :: a
    ba = fmap go t        -- :: f (b,a)

--------------------------------------------------------------------------------
-- * Futu- and histomorphisms

{-
newtype Free   f a = Free   { unFree   :: Either a (f (Free f a)) }

-- | @CoFree f a@ is basically an @a@-annotated version of @Mu f@. So it is isomorphic to @Attr f a@.
newtype CoFree f a = CoFree { unCoFree :: (a , f (CoFree f a))    }

-- | Futumorphism. Whatever it does.
futu :: Functor f => (a -> f (Free f a)) -> a -> Mu f
futu h = go where
  -- go :: a -> Mu f
  go = Fix . fmap worker . h 
  -- worker :: Free f a -> Mu f
  worker (Free ei) = case ei of
    Left  x -> go x
    Right t -> Fix (fmap worker t)

-- | Histomorphism. 
histo :: Functor f => (f (CoFree f a) -> a) -> Mu f -> a
histo h = go where
  -- go :: Mu f -> a
  go = h . fmap worker . unFix
  -- worker :: Mu f -> CoFree f
  worker t@(Fix s) = CoFree ( go t , fmap worker s )
-}

-- | Histomorphism. This is a kind of glorified cata/paramorphism, after all:
--
-- > cata f == histo (f . fmap attribute)
-- > para f == histo (f . fmap (\t -> (forget t, attribute t)))
--
histo :: Functor f => (f (Attr f a) -> a) -> Mu f -> a
histo h = go where
  go = h . fmap worker . unFix
  worker t@(Fix s) = Fix (Ann (go t) (fmap worker s))


-- | Futumorphism. This is a more interesting unfold.
futu :: Functor f => (a -> f (CoAttr f a)) -> a -> Mu f
futu h = go where
  go = Fix . fmap worker . h 
  worker (Fix ei) = case ei of
    Pure  x -> go x
    CoAnn t -> Fix (fmap worker t)

--------------------------------------------------------------------------------
-- * Monadic versions

-- | Monadic catamorphism.
cataM :: (Monad m, Traversable f) => (f a -> m a) -> Mu f -> m a
cataM h = go where
  go (Fix t) = mapM go t >>= h

cataM_ :: (Monad m, Traversable f) => (f a -> m a) -> Mu f -> m ()
cataM_ h t = do { _ <- cataM h t ; return () }

-- | Monadic paramorphism.
paraM :: (Monad m, Traversable f) => (f (Mu f, a) -> m a) -> Mu f -> m a
paraM h = go where
  go (Fix t) = mapM go' t >>= h
  go' t = go t >>= \x -> return (t,x)

paraM_ :: (Monad m, Traversable f) => (f (Mu f, a) -> m a) -> Mu f -> m ()
paraM_ h t = do { _ <- paraM h t ; return () }

-- | Another version of 'paraM'.
paraM' :: (Monad m, Traversable f) => (Mu f -> f a -> m a) -> Mu f -> m a
paraM' h = go where
  go t = mapM go (unFix t) >>= h t

{-
paraM_ :: (Monad m, Traversable f) => (Mu f -> f a -> m a) -> Mu f -> m ()
paraM_ h t = do { _ <- paraM h t ; return () }
-}

--------------------------------------------------------------------------------
#ifdef WITH_QUICKCHECK
-- * Tests

runtests_Morphisms :: IO ()
runtests_Morphisms = do
  quickCheck prop_para
  quickCheck prop_paraList
  quickCheck prop_cataHisto
  quickCheck prop_paraHisto
  -- quickCheck prop_zygo      -- moved to Attributes.hs, to avoid circular imports
  -- quickCheck prop_zygo_
 
prop_para :: FixT Label -> Bool
prop_para tree = para f tree == para' f' tree where
  f' :: FixT Label -> TreeF Label Integer -> Integer
  f' t@(Fix (TreeF (Label label) sub)) js = h label (toList sub) (toList js)

  f :: TreeF Label (FixT Label, Integer) -> Integer
  f t@(TreeF (Label label) subjs) = h label sub js where
    (sub,js) = unzip $ toList t

  h :: String -> [FixT Label] -> [Integer] -> Integer
  h label ts js = Prelude.sum $ zipWith (*) [3..] (map (fi.ord) label ++ map g ts ++ js)
  g (Fix (TreeF (Label label) _)) = (Prelude.sum (map (fi.ord) label)) `mod` 59

  fi = fromIntegral :: Int -> Integer

prop_paraList :: FixT Label -> Bool
prop_paraList tree = para' f tree == paraList flist tree where
  f t s = flist t (toList s)
  flist :: FixT Label -> [Integer] -> Integer
  flist t@(Fix (TreeF (Label label) sub)) js = Prelude.sum $ zipWith (*) [4..] (map (fi.ord) label ++ js)

  fi = fromIntegral :: Int -> Integer

prop_cataHisto :: FixT Label -> Bool
prop_cataHisto tree = (cata f tree == histo (f . fmap attribute) tree) where

  f :: TreeF Label String -> String
  f t@(TreeF (Label label) child) = "<" ++ label ++ ">[" ++ intercalate "," child ++ "]"

prop_paraHisto :: FixT Label -> Bool
prop_paraHisto tree = (para f tree == histo (f . fmap (\t -> (forget t, attribute t))) tree) where
   
  f :: TreeF Label (FixT Label, Integer) -> Integer
  f t@(TreeF (Label label) subjs) = h label sub js where
    (sub,js) = unzip $ toList t

  h :: String -> [FixT Label] -> [Integer] -> Integer
  h label ts js = Prelude.sum $ zipWith (*) [3..] (map (fi.ord) label ++ map g ts ++ js)
  g (Fix (TreeF (Label label) _)) = (Prelude.sum (map (fi.ord) label)) `mod` 59

  fi = fromIntegral :: Int -> Integer


#endif
--------------------------------------------------------------------------------