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

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.Trans.Maybe

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 Show (NDT Identity a) where
  show (NDTBind x _) = "Bind (" ++ show x ++ ") _"
  show x = show . (\(NDTCons x) -> x) $ fmap (const ()) x

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) = trace "(cons <> cons)" . NDTCons . trace "[cons <> cons]" $ maybe y (\(x, xs) -> return $ Just (x, xs <> y')) =<< x
  mappend (NDTBind x f) (NDTBind y g) = trace "(bind <> bind)" $ NDTBind (fmap Left x <> fmap Right y) (either f g)
  mappend x@(NDTBind _ _) y = trace "(bind <>)" $ x <> NDTBind y return
  mappend x y@(NDTBind _ _) = trace "(<> bind)" $ 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 MonadFix m => MonadFix (NDT m) where
--   mfix f = NDTCons . runMaybeT $ do
--     x <- mfix (head . f)
--     return (x, trace "tail" . mfix $ tail . f)
--     where
--       head :: Monad m => NDT m a -> MaybeT m a
--       head (NDTCons x) = MaybeT . trace "head (cons)" $ fmap fst <$> x
--       head (NDTBind (NDTBind x g) f) = head $ NDTBind x (f <=< g)
--       head (NDTBind (NDTCons x) f) = MaybeT x >>= head . f . fst . trace "head (bind cons)"
--       tail :: Monad m => NDT m a -> NDT m a
--       tail (NDTCons x) = NDTBind (lift x) $ maybe empty snd . guardNothing
--       tail (NDTBind (NDTBind x g) f) = tail $ NDTBind x (f <=< g)
--       tail (NDTBind (NDTCons x) f) = tail . NDTCons $ fmap (\(_, xs) -> (undefined, NDTBind xs f)) <$> x
--       guardNothing :: Maybe a -> Maybe a
--       guardNothing x@(Just _) = x
--       guardNothing x@(Nothing) = trace "Nothing" x

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

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 trace ("(cons "++ show continue ++ ")") continue of
        False -> return mempty
        True  -> (pure x <>) <$> foldNDT sel mxs
foldNDT sel (NDTBind (NDTCons x) f) = trace "(bind cons)" $ 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) = trace "(bind bind)" . foldNDT sel $ NDTBind x (f <=< g)

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