module UniqueLogic.ST.SystemLog (
   -- * Preparation
   Variable,
   globalVariable,
   -- * Posing statements
   T,
   localVariable,
   constant,
   assignment2,
   assignment3,
   Apply, arg, runApply,
   -- * Solution
   solve,
   query,
   ) where

import qualified Control.Monad.Trans.Writer as MW
import qualified Control.Monad.Trans.Class  as MT
import qualified Data.Foldable as Fold
import Control.Monad.Trans.Writer (WriterT, Writer, )
import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT, mapMaybeT, )
import Control.Monad.ST (ST, )
import Control.Monad.HT ((<=<), )
import Control.Monad (when, liftM2, ap, void, )
import Control.Applicative (Applicative, pure, (<*>), )
import Data.Functor.Compose (Compose(Compose))

import Data.STRef (STRef, newSTRef, modifySTRef, readSTRef, writeSTRef, )
import Data.Monoid (Monoid, )
import Data.Maybe (isNothing, )


data Variable w s a =
   Variable {
      dependsRef :: STRef s [Update w s],
      valueRef :: STRef s (Maybe a)
   }

type Update w s = WriterT w (ST s) ()

newtype T w s a =
   Cons {run :: WriterT [STRef s [Update w s]] (ST s) a}

instance Functor (T w s) where
   fmap f (Cons x) = Cons (fmap f x)

instance Applicative (T w s) where
   pure = Cons . return
   (<*>) = ap

instance Monad (T w s) where
   return = Cons . return
   Cons x >>= k  = Cons $ run . k =<< x


lift :: ST s a -> T w s a
lift = Cons . MT.lift

localVariable :: T w s (Variable w s a)
localVariable = lift globalVariable

globalVariable :: ST s (Variable w s a)
globalVariable = object Nothing

constant :: a -> T w s (Variable w s a)
constant a =
   do v <- lift $ object $ Just a
      Cons $ MW.tell [dependsRef v]
      return v

object :: Maybe a -> ST s (Variable w s a)
object ma =
   liftM2 Variable (newSTRef []) (newSTRef ma)

resolve ::
   Monoid w =>
   STRef s [Update w s] -> Update w s
resolve =
   sequence_ <=< MT.lift . readSTRef

solve ::
   Monoid w =>
   T w s a -> WriterT w (ST s) a
solve (Cons m) = do
   (a,w) <- MT.lift $ MW.runWriterT m
   mapM_ resolve w
   return a

query :: Variable w s a -> ST s (Maybe a)
query = readSTRef . valueRef


mw ::
   (Monoid w, Monad st) =>
   MaybeT st (Writer w a) -> MaybeT (WriterT w st) a
--   st (Maybe (Writer w a)) -> WriterT w st (Maybe a)
mw act = do
   mwa <- mapMaybeT MT.lift act
   case MW.runWriter mwa of
      (a,w) -> MT.lift $ MW.tell w >> return a

updateIfNew ::
   Monoid w =>
   Variable w s a -> MaybeT (ST s) (Writer w a) -> Update w s
updateIfNew (Variable al av) act = do
   as <- MT.lift $ readSTRef av
   when (isNothing as) $ void $ runMaybeT $ do
      MT.lift . MT.lift . writeSTRef av . Just =<< mw act
      MT.lift $ resolve al

readSTRefM :: STRef s (Maybe a) -> MaybeT (ST s) a
readSTRefM = MaybeT . readSTRef


assignment2 ::
   Monoid w =>
   (a -> Writer w b) ->
   Variable w s a -> Variable w s b ->
   T w s ()
assignment2 f (Variable al av) b =
   let update =
          updateIfNew b $ fmap f $ readSTRefM av
   in  lift $
       modifySTRef al (update :)

assignment3 ::
   Monoid w =>
   (a -> b -> Writer w c) ->
   Variable w s a -> Variable w s b -> Variable w s c ->
   T w s ()
assignment3 f (Variable al av) (Variable bl bv) c =
   let update =
          updateIfNew c $
          liftM2 f (readSTRefM av) (readSTRefM bv)
   in  lift $
       modifySTRef al (update :) >>
       modifySTRef bl (update :)


data Apply w s a =
   Apply (Compose (MW.Writer [STRef s [Update w s]]) (MaybeT (ST s)) a)


{- |
This function allows to generalize 'assignment2' and 'assignment3' to more arguments.
You could achieve the same with nested applications of @assignment3 (,)@.
-}
arg :: Monoid w => Variable w s a -> Apply w s a
arg (Variable al av) =
   Apply $ Compose $ MW.writer (readSTRefM av, [al])


instance Monoid w => Functor (Apply w s) where
   fmap f (Apply a) = Apply $ fmap f a

instance Monoid w => Applicative (Apply w s) where
   pure a = Apply $ pure a
   Apply f <*> Apply a = Apply $ f <*> a


runApply ::
   Monoid w =>
   Apply w s (Writer w a) -> Variable w s a -> T w s ()
runApply (Apply (Compose w)) a =
   case MW.runWriter w of
      (f, refs) ->
         lift $ Fold.forM_ refs $ flip modifySTRef (updateIfNew a f :)