{-# LANGUAGE GADTs, ScopedTypeVariables, FlexibleContexts #-}
module Make.App 
    (
     App,WApp,collect,expr,runApp,apply,readT,shiftR
    )
    where

import Control.Monad
import Control.Applicative
import Data.Traversable
import Data.Function
import Control.Arrow ((***))

data App rule t r x where
    Ap :: App rule t r (a -> b) -> App rule t r a -> App rule t r b
    Pure :: a -> App rule t r a
    Read :: t -> rule t r -> (r -> a) -> App rule t r a
    Bind :: App rule t r a -> (a -> WApp rule t r b) -> App rule t r b

runApp :: App r t t1 t2 -> t2
runApp (Pure x) = x

newtype WApp rule t r x = WApp { unW :: ([(t,rule t r)],App rule t r x)} deriving (Show)

collect :: WApp rule t r x -> [(t,rule t r)]
collect = fst . unW

expr :: WApp rule t r x -> App rule t r x
expr = snd . unW

readT :: t -> rule t r -> WApp rule t r r
readT l r = WApp ([(l,r)], Read l r id)

apply :: (Eq t) => t -> r -> App rule t r a -> ([(t,rule t r)], App rule t r a)
apply t r expr = unW $ subst t r expr


instance Functor (App rule t r) where
    fmap f (Pure x) = Pure (f x)
    fmap f (Read l r k) = Read l r (f . k)
    fmap f (Ap x y) = Ap (fmap (f .) x) y
    fmap f (Bind x k) = Bind x (fmap f . k)

instance Applicative (App rule t r) where
    Pure f <*> x = fmap f x
    f <*> Pure x = fmap ($x) f
    x <*> y = Ap x y
    pure = Pure


instance Functor (WApp rule t r) where
    fmap f = WApp . fmap (fmap f) . unW

instance Applicative (WApp rule t r) where
    pure = WApp . pure . pure 
    (WApp f) <*> WApp x = WApp $ liftA2 (<*>) f x

instance Monad (WApp rule t r) where
    return = pure
    (WApp (l,x)) >>= f = WApp y
        where y = case x of
                    (Pure v) -> let (l',x') = unW (f v) in (l++l',x')
                    _        -> (l,Bind x f)



shiftR :: forall rule t' r' t r a. (t -> t') -> (r' -> r) 
       -> (rule t r -> rule t' r') -> WApp rule t r a -> WApp rule t' r' a
shiftR inF fromG f = goW
    where 
          goW = WApp . (map (inF *** f) *** goA) . unW
          goA :: forall a. App rule t r a -> App rule t' r' a
          goA (Read t r k) = Read (inF t) (f r) (k . fromG)
          goA (Ap mf x) = Ap (goA mf) (goA x)
          goA (Bind m g) = Bind (goA m) (goW . g)
          goA (Pure x) = Pure x


subst :: forall t r rule a. Eq t => t -> r -> (App rule t r a) -> WApp rule t r a
subst l x = subst'
    where 
      subst' :: forall a. App rule t r a -> WApp rule t r a
      subst' (Read l' r k) | l == l' = inW $ Pure (k x)
                           | otherwise = inW $ (Read l' r k)
      subst' (Ap f a) = subst' f <*> subst' a
      subst' (Bind x f) = subst' x >>= (WApp . join' . fmap subst' . unW . f)
      subst' (Pure x) = inW $ Pure x
      inW e = WApp ([],e)

join' (l,WApp (l',x)) =(l++l',x)

instance (Show (rule t r), Show t) => Show (App rule t r x) where
    show (Ap x y) = "Ap ("++show x++") ("++show y++")"
    show (Pure x) = "Pure"
    show (Read l r _) = "Read "++ show l ++ " (" ++ show r ++ ")"
    show (Bind x _) = "Bind (" ++ show x++ ")"