{-# LANGUAGE GADTs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances #-}
module Data.Machine.Type
  (
  
    MachineT(..)
  , Step(..)
  , Machine
  , runT_
  , runT
  , run
  , runMachine
  , encased
  
  , construct
  , repeatedly
  , unfoldPlan
  , before
  , preplan
  
  , deconstruct
  , tagDone
  , finishWith
  
  , fit
  , fitM
  , pass
  , starve
  , stopped
  , stepMachine
  
  , Appliance(..)
  ) where
import Control.Applicative
import Control.Category
import Control.Monad (liftM)
import Data.Foldable
import Data.Functor.Identity
import Data.Machine.Plan
import Data.Monoid hiding ((<>))
import Data.Pointed
import Data.Profunctor.Unsafe ((#.))
import Data.Semigroup
import Prelude hiding ((.),id)
data Step k o r
  = Stop
  | Yield o r
  | forall t. Await (t -> r) (k t) r
instance Functor (Step k o) where
  fmap _ Stop = Stop
  fmap f (Yield o k) = Yield o (f k)
  fmap f (Await g kg fg) = Await (f . g) kg (f fg)
newtype MachineT m k o = MachineT { runMachineT :: m (Step k o (MachineT m k o)) }
type Machine k o = forall m. Monad m => MachineT m k o
runMachine :: MachineT Identity k o -> Step k o (MachineT Identity k o)
runMachine = runIdentity . runMachineT
encased :: Monad m => Step k o (MachineT m k o) -> MachineT m k o
encased = MachineT #. return
stepMachine :: Monad m => MachineT m k o -> (Step k o (MachineT m k o) -> MachineT m k' o') -> MachineT m k' o'
stepMachine m f = MachineT (runMachineT #. f =<< runMachineT m)
instance Monad m => Functor (MachineT m k) where
  fmap f (MachineT m) = MachineT (liftM f' m) where
    f' (Yield o xs)    = Yield (f o) (f <$> xs)
    f' (Await k kir e) = Await (fmap f . k) kir (f <$> e)
    f' Stop            = Stop
instance Monad m => Pointed (MachineT m k) where
  point = repeatedly . yield
instance Monad m => Semigroup (MachineT m k o) where
  a <> b = stepMachine a $ \step -> case step of
    Yield o a'    -> encased (Yield o (mappend a' b))
    Await k kir e -> encased (Await (\x -> k x <> b) kir (e <> b))
    Stop          -> b
instance Monad m => Monoid (MachineT m k o) where
  mempty        = stopped
  mappend       = (<>)
class Appliance k where
  applied :: Monad m => MachineT m k (a -> b) -> MachineT m k a -> MachineT m k b
instance (Monad m, Appliance k) => Applicative (MachineT m k) where
  pure = point
  (<*>) = applied
{-# INLINABLE runT_ #-}
runT_ :: Monad m => MachineT m k b -> m ()
runT_ m = runMachineT m >>= \v -> case v of
  Stop        -> return ()
  Yield _ k   -> runT_ k
  Await _ _ e -> runT_ e
{-# INLINABLE runT #-}
runT :: Monad m => MachineT m k b -> m [b]
runT (MachineT m) = m >>= \v -> case v of
  Stop        -> return []
  Yield o k   -> liftM (o:) (runT k)
  Await _ _ e -> runT e
run :: MachineT Identity k b -> [b]
run = runIdentity . runT
instance (m ~ Identity) => Foldable (MachineT m k) where
  foldMap f (MachineT (Identity m)) = go m where
    go Stop = mempty
    go (Yield o k) = f o `mappend` foldMap f k
    go (Await _ _ fg) = foldMap f fg
fit :: Monad m => (forall a. k a -> k' a) -> MachineT m k o -> MachineT m k' o
fit f (MachineT m) = MachineT (liftM f' m) where
  f' (Yield o k)     = Yield o (fit f k)
  f' Stop            = Stop
  f' (Await g kir h) = Await (fit f . g) (f kir) (fit f h)
{-# INLINE fit #-}
fitM :: (Monad m, Monad m')
     => (forall a. m a -> m' a) -> MachineT m k o -> MachineT m' k o
fitM f (MachineT m) = MachineT $ f (liftM aux m)
  where aux Stop = Stop
        aux (Yield o k) = Yield o (fitM f k)
        aux (Await g kg gg) = Await (fitM f . g) kg (fitM f gg)
{-# INLINE fitM #-}
construct :: Monad m => PlanT k o m a -> MachineT m k o
construct m = MachineT $ runPlanT m
  (const (return Stop))
  (\o k -> return (Yield o (MachineT k)))
  (\f k g -> return (Await (MachineT #. f) k (MachineT g)))
  (return Stop)
{-# INLINE construct #-}
repeatedly :: Monad m => PlanT k o m a -> MachineT m k o
repeatedly m = r where
  r = MachineT $ runPlanT m
    (const (runMachineT r))
    (\o k -> return (Yield o (MachineT k)))
    (\f k g -> return (Await (MachineT #. f) k (MachineT g)))
    (return Stop)
{-# INLINE repeatedly #-}
unfoldPlan :: Monad m => s -> (s -> PlanT k o m s) -> MachineT m k o
unfoldPlan s0 sp = r s0 where
  r s = MachineT $ runPlanT (sp s)
      (\sx -> runMachineT $ r sx)
      (\o k -> return (Yield o (MachineT k)))
      (\f k g -> return (Await (MachineT #. f) k (MachineT g)))
      (return Stop)
{-# INLINE unfoldPlan #-}
before :: Monad m => MachineT m k o -> PlanT k o m a -> MachineT m k o
before (MachineT n) m = MachineT $ runPlanT m
  (const n)
  (\o k -> return (Yield o (MachineT k)))
  (\f k g -> return (Await (MachineT #. f) k (MachineT g)))
  (return Stop)
{-# INLINE before #-}
preplan :: Monad m => PlanT k o m (MachineT m k o) -> MachineT m k o
preplan m = MachineT $ runPlanT m
  runMachineT
  (\o k -> return (Yield o (MachineT k)))
  (\f k g -> return (Await (MachineT #. f) k (MachineT g)))
  (return Stop)
{-# INLINE preplan #-}
pass :: k o -> Machine k o
pass k =
    loop
  where
    loop = encased (Await (\t -> encased (Yield t loop)) k stopped)
{-# INLINE pass #-}
starve :: Monad m => MachineT m k0 b -> MachineT m k b -> MachineT m k b
starve m cont = MachineT $ runMachineT m >>= \v -> case v of
  Stop            -> runMachineT cont 
  Yield o r       -> return $ Yield o (starve r cont)
  Await _ _ r     -> runMachineT (starve r cont)
{-# INLINE starve #-}
stopped :: Machine k b
stopped = encased Stop
{-# INLINE stopped #-}
deconstruct :: Monad m => MachineT m k (Either a o) -> PlanT k o m a
deconstruct m = PlanT $ \r y a f ->
  let aux k = runPlanT (deconstruct k) r y a f
  in runMachineT m >>= \v -> case v of
       Stop -> f
       Yield (Left o) _ -> r o
       Yield (Right o) k -> y o (aux k)
       Await g fk h -> a (aux . g) fk (aux h)
tagDone :: Monad m => (o -> Bool) -> MachineT m k o -> MachineT m k (Either o o)
tagDone f = fmap aux
  where aux x = if f x then Left x else Right x
finishWith :: Monad m
           => (o -> Maybe r) -> MachineT m k o -> MachineT m k (Either r o)
finishWith f = fmap aux
  where aux x = maybe (Right x) Left $ f x