module UniqueLogic.ST.System (
   -- * Preparation
   Variable,
   globalVariable,
   -- * Posing statements
   M,
   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.Maybe (MaybeT(MaybeT), runMaybeT, )
import Control.Monad.ST (ST, )
import Control.Monad.HT ((<=<), )
import Control.Monad (when, liftM2, ap, void, )
import Control.Applicative (Applicative, pure, liftA, liftA2, (<*>), )
import Data.Functor.Compose (Compose(Compose))

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


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

newtype M s a =
   M {runM :: MW.WriterT [STRef s [ST s ()]] (ST s) a}

instance Functor (M s) where
   fmap f (M x) = M (fmap f x)

instance Applicative (M s) where
   pure = M . return
   (<*>) = ap

instance Monad (M s) where
   return = M . return
   M x >>= k  = M $ runM . k =<< x


lift :: ST s a -> M s a
lift = M . MT.lift

localVariable :: M s (Variable s a)
localVariable = lift globalVariable

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

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

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

resolve ::
   STRef s [ST s ()] -> ST s ()
resolve =
   sequence_ <=< readSTRef

solve :: M s a -> ST s a
solve (M m) =
   do (a,w) <- MW.runWriterT m
      mapM_ resolve w
      return a

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



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

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

assignment2, _assignment2 ::
   String ->
   (a -> b) ->
   Variable s a -> Variable s b ->
   M s ()
assignment2 _ f (Variable al av) b =
   let update =
          updateIfNew b $ fmap f $ readSTRefM av
   in  lift $
       modifySTRef al (update :)

assignment3, _assignment3 ::
   String ->
   (a -> b -> c) ->
   Variable s a -> Variable s b -> Variable s c ->
   M 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 :)


newtype Apply s a =
   Apply (Compose (MW.Writer [STRef s [ST s ()]]) (MaybeT (ST s)) a)

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

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


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

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


{-
examples of how to use 'arg' and 'runApply'
-}
_assignment2 msg f x = runApply msg (liftA f $ arg x)
_assignment3 msg f x y = runApply msg (liftA2 f (arg x) (arg y))