{-# LANGUAGE GADTs #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses, UndecidableInstances #-}

module Events.Types.NDT
  ( NDT
  , foldNDT
  , cons
  , fromFoldable
  ) where

import Data.Monoid
import Data.Foldable (foldr)
import Data.Maybe
import Data.Either

import Control.Applicative (Alternative)
import qualified Control.Applicative as Alt (Alternative(..))
import Control.Monad
import Control.Monad.Identity

import Control.Monad.Trans
import Control.Monad.Reader (MonadReader(..))
import Control.Monad.Trans.Maybe
import Control.Monad.Catch (MonadThrow(..))

import Debug.Trace

data NDT m a where
  NDTBind :: NDT m a -> (a -> NDT m b) -> NDT m b
  NDTCons :: m (Maybe (a, NDT m a)) -> NDT m a

instance Functor m => Functor (NDT m) where
  fmap f (NDTBind a g) = NDTBind a (fmap f . g)
  fmap f (NDTCons x) = NDTCons $ fmap f' x
    where
      f' Nothing = Nothing
      f' (Just (x, xs)) = Just (f x, fmap f xs)

instance Applicative m => Applicative (NDT m) where
  pure x = NDTCons . pure $ Just (x, empty)
  fs <*> xs = fs >>= (\f -> xs >>= pure . (f $))

instance Applicative m => Monad (NDT m) where
  return = pure
  (>>=) = NDTBind

instance Monad m => Monoid (NDT m a) where
  mempty = empty
  mappend (NDTCons x) y'@(NDTCons y) = NDTCons $ maybe y (\(x, xs) -> return $ Just (x, xs <> y')) =<< x
  mappend (NDTBind x f) (NDTBind y g) = NDTBind (fmap Left x <> fmap Right y) (either f g)
  mappend x@(NDTBind _ _) y = x <> NDTBind y return
  mappend x y@(NDTBind _ _) = NDTBind x return <> y

instance MonadTrans NDT where
  lift = NDTCons . fmap Just . fmap (,empty)

instance Monad m => Alternative (NDT m) where
  empty = mempty
  (<|>) = mappend

instance Monad m => MonadPlus (NDT m) where
  mzero = mempty
  mplus = mappend

instance MonadReader r m => MonadReader r (NDT m) where
  reader = lift . reader
  local f (NDTCons x) = NDTCons (local f x)
  local f (NDTBind x g) = NDTBind (local f x) g

instance MonadIO m => MonadIO (NDT m) where
  liftIO = lift . liftIO

instance MonadThrow m => MonadThrow (NDT m) where
  throwM = lift . throwM

empty :: Applicative m => NDT m a
empty = NDTCons $ pure Nothing

cons :: Applicative m => a -> NDT m a -> NDT m a
cons x xs = NDTCons . pure $ Just (x, xs)
  
foldNDT :: (Foldable f, Applicative f, Monoid (f a), Monad m) => (a -> m Bool) -> NDT m a -> m (f a)
-- ^ Evaluate depth-first, pruning leaves under the assumption that the selection predicate is monotonic on siblings and children
foldNDT sel (NDTCons mx) = do
  mx' <- mx
  case mx' of
    Nothing -> return mempty
    Just (x, mxs) -> do
      continue <- sel x
      case continue of
        False -> return mempty
        True  -> (pure x <>) <$> foldNDT sel mxs
foldNDT sel (NDTBind (NDTCons x) f) = do
  x' <- x
  case x' of
    Nothing -> return mempty
    Just (x'', xs) -> do
      x3 <- foldNDT sel $ f x''
      xs' <- if null x3 then return mempty else foldNDT sel (NDTBind xs f)
      return $ x3 <> xs'
foldNDT sel (NDTBind (NDTBind x g) f) = foldNDT sel $ NDTBind x (f <=< g)

fromFoldable :: (Foldable f, Monad m) => f a -> NDT m a
fromFoldable = foldr cons empty