{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE Rank2Types #-}

module Annotations (Ann, Ann1, mkAnn, unannotate, errorCata, flattenF, flattenAnn, filterAnn, SearchHint(..), search, cascade, cata, Except, ErrorAlgebra) where

import Control.Functor.Fix
import Control.Functor.Composition
import Control.Functor.Algebra
-- import Control.Monad.Either ()

import Data.Foldable
import Data.Traversable
import Data.Monoid
import Control.Applicative
import Prelude hiding (sequence, concatMap, mapM)


-- | Annotate a functor.
-- type Ann x f = ((,) x :.: f)

-- | The fixpoint of an annotated functor.
type Ann x f = FixF ((,) x :.: f)

type Ann1 x f = f (FixF ((,) x :.: f))

mkAnn :: x -> Ann1 x f -> Ann x f
mkAnn x = InF . CompF . (,) x

-- -- | Monadic fold over fixpoints.
-- foldM :: (Traversable f, Monad m) => (f a -> m a) -> FixF f -> m a
-- foldM f (InF x) = mapM (foldM f) x >>= f

type ErrorAlgebra f e a = f a -> Either e a

data Except e a = OK a | Failed e
  deriving (Eq, Show)

instance Functor (Except e) where
  fmap f (OK x) = OK (f x)
  fmap _ (Failed err) = Failed err

instance Monoid e => Applicative (Except e) where 
  pure = OK 
  OK f         <*> OK x         = OK (f x) 
  OK _         <*> Failed err   = Failed err 
  Failed err   <*> OK _         = Failed err 
  Failed err1  <*> Failed err2  = Failed (err1 `mappend` err2) 


-- | Error fold over annotated fixpoints.
errorCata :: Traversable f => ErrorAlgebra f e a -> Ann x f -> Except [(e, x)] a
errorCata alg (InF (CompF (x, expr))) = case traverse (errorCata alg) expr of
  Failed xs -> Failed xs
  OK expr' -> case alg expr' of
    Left x'  -> Failed [(x', x)]
    Right v  -> OK v

cata :: Functor f => (f a -> a) -> FixF f -> a
cata f = f . fmap (cata f) . outF

ana :: Functor f => (a -> f a) -> a -> FixF f
ana f = InF . fmap (ana f) . f

cascade :: (Traversable f, Monoid e) => ErrorAlgebra f e a -> Algebra f (Except e a)
cascade alg expr = case sequenceA expr of
  -- At least one child produced an error.
  Failed xs -> Failed xs
  
  -- Children produced valid results.
  OK tree' -> case alg tree' of
    
    -- Algebra at this level produced error. Cascade upwards.
    Left xs -> Failed xs
    
    -- Algebra produced valid result.
    Right res -> OK res

-- | Decompose over fixpoints.
decomposeLF :: Functor g => (forall a. f (g a) -> g a) -> FixF (f :.: g) -> FixF g
decomposeLF unF = InF . fmap (decomposeLF unF) . unF . runCompF . outF

-- | Remove a recursive annotation from a tree.
unannotate :: Functor f => Ann x f -> FixF f
unannotate = decomposeLF snd

-- | Pre-order traversal over a fixpoint.
flattenF :: Foldable f => FixF f -> [FixF f]
flattenF r@(InF f) = r : concatMap flattenF f

-- | Pre-order traversal over an annotated fixpoint.
flattenAnn :: Foldable f => Ann x f -> [Ann x f]
flattenAnn root@(InF (CompF (_, sub))) = root : concatMap flattenAnn (toList sub)

-- | Yield those subtrees whose annotation matches the predicate.
filterAnn :: Foldable f => (x -> Bool) -> Ann x f -> [Ann x f]
filterAnn p = filter (p . fst . runCompF . outF) . flattenAnn

-- | Hints the search algorithm which way to search in the tree.
data SearchHint
  = Here  -- ^ Stop searching: this is the node we are looking for.
  | Down  -- ^ Continue searching, but only in the node's children.
  | Skip  -- ^ Continue searching, but only in the node's right siblings.
  deriving (Show, Eq, Enum, Bounded)

-- | Searches a tree for nodes whose annotation matches the hints.
search :: Foldable f => (x -> SearchHint) -> Ann x f -> [(x, f (Ann x f))]
search hint (InF (CompF (ann, tree))) =
    case hint ann of
      Here -> [(ann, tree)]
      Skip -> []
      Down -> concatMap (search hint) (toList tree)