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

module Annotations (Ann, AnyAnn, WithBounds, AnyWithBounds, unannotate, children, flatten, debugFlatten) where

import BoundsParser
import ShowFam
import Any

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

import Generics.MultiRec
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) s)

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

-- | The fixpoint of a data family @s@ recursively annotated with 'Bounds'.
type WithBounds  s = Ann  Bounds s

-- | The fixpoint of a data family @s@ recursively annotated with 'Bounds', with existentially qualified top-level index.
type AnyWithBounds s = AnyAnn Bounds s

-- | Removes all annotations from a recursively annotated fixpoint.
unannotate :: HFunctor g => HFix ((f :*: g) s) ix -> HFix (g s) ix
unannotate = HIn . hmap (const unannotate) . snd' . hout

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

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

-- | Collects the direct children of a functor in a list.
children :: (HFunctor f, HFunctor (PF s)) => f s r ix -> [AnyF s r]
children x = execWriter (hmapM collect x) where
  collect :: (Ix s ix, HFunctor (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 :: (Ix s ix, HFunctor (PF s)) => Ann x s ix -> [(x, Any s)]
flatten (HIn (K x :*: y)) = [(x, mkAny . hto . HIn . hmap (const unannotate) $ y)] ++
    concatMap childToList (children y) where
  childToList :: HFunctor (PF s) => AnyF s (Ann x s) -> [(x, Any s)]
  childToList (AnyF _ z) = flatten z

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