{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Generics.Annotations.Annotations (Ann, Ann1, mkAnn, AnyAnn, unannotate, children, flatten, filterAnn, debugFlatten) where

import Generics.Annotations.ShowFam
import Generics.Annotations.Any

import Control.Monad.Writer (Writer, execWriter, tell)

import Generics.MultiRec hiding (show)
import Generics.MultiRec.HFix

-- | A fixpoint of a data family @s@ annotated with an @x@ at every recursive position.
type Ann  (x :: *) (s :: * -> *) = HFix (K x :*: PF s)

-- | A fixpoint of a data family @s@ annotated with an @x@ at every recursive position except the root. Can be turned into a fully annotated tree @Ann x s@ using 'mkAny'.
type Ann1 (x :: *) (s :: * -> *) = (PF s) (Ann x s)

mkAnn :: x -> Ann1 x s ix -> Ann x s ix
mkAnn x = HIn . (K x :*:)

-- | A fixpoint of a data family @s@ annotated with an @x@ at every recursive position, with existentially quantified top-level index.
type AnyAnn x s = AnyF s (Ann x s)

-- hmap_ :: HFunctor phi f => (forall ix. phi ix -> r ix -> r' ix) -> phi ix -> f r ix -> f r' ix
-- hmap_ f _ = hmap f
-- 
-- hmapM_ :: (HFunctor phi f, Monad m) => (forall ix. phi ix -> r ix -> m (r' ix)) -> phi ix -> f r ix -> m (f r' ix)
-- hmapM_ f _ = hmapM f

-- | Removes all annotations from a recursively annotated fixpoint.
unannotate :: HFunctor s (PF s) => s ix -> Ann x s ix -> HFix (PF s) ix
unannotate p = HIn . hmap unannotate p . snd' . hout

-- fst' :: (f :*: g) r ix -> f r ix
-- fst' (x :*: _) = x

snd' :: (f :*: g) r ix -> g r ix
snd' (_ :*: y) = y

-- | Collects the direct children of a functor in a list.
children :: HFunctor s (PF s) => s ix -> (PF s) r ix -> [AnyF s r]
children p x = execWriter (hmapM collect p x) where
  collect :: (HFunctor s (PF s)) => s ix -> r ix -> Writer [AnyF s r] (r ix)
  collect w x = tell [AnyF w x] >> return x

-- | Flatten an annotated tree to a list of subtrees coupled with their annotations.
flatten :: forall s x ix. (HFunctor s (PF s), Fam s) => s ix -> Ann x s ix -> [(x, Any s)]
flatten p tree@(HIn (K x :*: y)) = [(x, Any p (hto p (unannotate p tree :: HFix (PF s) ix)))] ++
    concatMap (flatten $?) (children p y)

-- | Yield all subtrees whose annotation matches the predicate.
filterAnn :: (Fam s, HFunctor s (PF s)) => s ix -> (x -> Bool) -> Ann x s ix -> [(x, Any s)]
filterAnn p f = filter (f . fst) . flatten p

-- | Flatten an annotated tree and print all subtrees to stdout.
debugFlatten :: (HFunctor s (PF s), ShowFam s, Show x, Fam s) => s ix -> Ann x s ix -> IO ()
debugFlatten p x = putStr . unlines . map show $ flatten p x