{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
module UniqueLogic.ST.TF.System (
   -- * Preparation
   Variable,
   globalVariable,
   -- * Handle duplicates
   C, update,
   simpleUpdate, -- should be private in future
   updateIfNew, -- should be private or with special type
   updateAndCheck,
   Fragile(break),
   Value, ValueConstraint, valueConstraint,
   -- * Posing statements
   T,
   localVariable,
   constant,
   assignment2,
   assignment3,
   Apply, arg, runApply,
   -- * Solution
   solve, solveDepthFirst, solveBreadthFirst,
   query,
   ) where

import qualified Control.Monad.Exception.Synchronous as E
import qualified Control.Monad.Trans.Writer as MW
import qualified Control.Monad.Trans.Class  as MT
import qualified UniqueLogic.ST.TF.MonadTrans as UMT
import qualified Data.Sequence as Seq
import qualified Data.Foldable as Fold
import Control.Monad.Trans.Writer (WriterT, )
import Control.Monad.Trans.Maybe (MaybeT(MaybeT), runMaybeT, mapMaybeT, )
import Control.Monad.Trans.Identity (IdentityT, )
import Control.Monad.ST (ST, )
import Control.Monad (when, liftM2, ap, guard, )
import Control.Applicative (Applicative, pure, (<*>), )
import Data.Sequence (Seq, (|>), ViewL((:<)), )
import Data.Functor.Compose (Compose(Compose))

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

import Prelude hiding (break)


data Variable w s a =
   Variable {
      varUpdate :: MaybeT (ST s) a -> Update w s,
      dependsRef :: STRef s (Updates w s),
      valueRef :: STRef s (Maybe a)
   }

type Update w s = UMT.Wrap w (ST s) (Updates w s)
newtype Updates w s = Updates {unpackUpdates :: Seq (Update w s)}

instance Monoid (Updates w s) where
   mempty = Updates Seq.empty
   mappend (Updates x) (Updates y) = Updates $ mappend x y

addUpdate :: Update w s -> Updates w s -> Updates w s
addUpdate x (Updates xs) = Updates $ xs |> x


type Updater w s a =
        STRef s (Updates w s) -> STRef s (Maybe a) ->
        MaybeT (UMT.Wrap w (ST s)) a -> Update w s

type SimpleUpdater w s a =
        STRef s (Updates w s) -> STRef s (Maybe a) ->
        MaybeT (ST s) a -> Update w s

newtype T w s a =
   Cons {run :: WriterT [STRef s (Updates 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

globalVariable ::
   (UMT.C w, Value w a) =>
   SimpleUpdater w s a -> ST s (Variable w s a)
globalVariable triggerUpdate = object triggerUpdate Nothing

localVariable :: (C w, Value w a) => T w s (Variable w s a)
localVariable = lift $ globalVariable simpleUpdate

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

object ::
   SimpleUpdater w s a ->
   Maybe a -> ST s (Variable w s a)
object updater ma = do
   al <- newSTRef mempty
   av <- newSTRef ma
   return $ Variable (updater al av) al av


solve, solveDepthFirst, solveBreadthFirst ::
   UMT.C w =>
   T w s a -> w (ST s) a
solve = solveDepthFirst

data Order = DepthFirst | BreadthFirst
   deriving (Eq, Enum)

solveDepthFirst   = solveOrder DepthFirst
solveBreadthFirst = solveOrder BreadthFirst

solveOrder ::
   UMT.C w =>
   Order -> T w s a -> w (ST s) a
solveOrder order (Cons m) = UMT.unwrap $ do
   let resolve updates =
          case Seq.viewl updates of
             Seq.EmptyL -> return ()
             currentUpdate :< remUpdates -> do
                Updates newUpdates <- currentUpdate
                resolve $
                   case order of
                      DepthFirst -> mappend newUpdates remUpdates
                      BreadthFirst -> mappend remUpdates newUpdates

   (a, w) <- UMT.lift $ MW.runWriterT m
   resolve . unpackUpdates . mconcat =<< mapM (UMT.lift . readSTRef) w
   return a

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


updateIfNew :: (C w) => Updater w s a
updateIfNew al av act = do
   as <- UMT.lift $ readSTRef av
   fmap Fold.fold $ runMaybeT $ do
      guard $ isNothing as
      MT.lift . UMT.lift . writeSTRef av . Just =<< act
      MT.lift $ UMT.lift $ readSTRef al


class Inconsistency e where
   inconsistency :: e

instance
   Inconsistency e =>
      Fragile (E.ExceptionalT e) where
   break =
      UMT.wrap $ E.throwT inconsistency

class C t => Fragile t where
   break :: Monad m => UMT.Wrap t m a

updateAndCheck ::
   (UMT.C w) =>
   (a -> a -> UMT.Wrap w (ST s) ()) ->
   Updater w s a
updateAndCheck customBreak al av act = do
   maold <- UMT.lift $ readSTRef av
   manew <- runMaybeT act
   case manew of
      Nothing -> return mempty
      Just anew -> do
         UMT.lift . writeSTRef av . Just $ anew
         case maold of
            Just aold -> customBreak aold anew >> return mempty
            Nothing -> UMT.lift $ readSTRef al


class C w => Value w a where
   data ValueConstraint w a :: *
   valueConstraint ::
      STRef s (Updates w s) -> STRef s (Maybe a) -> ValueConstraint w a

class UMT.C w => C w where
   update :: (Value w a) => Updater w s a

instance Value IdentityT a where
   data ValueConstraint IdentityT a = IdentityConstraint
   valueConstraint _ _ = IdentityConstraint

instance C IdentityT where
   update = updateIfNew

instance (Monoid w) => Value (MW.WriterT w) a where
   data ValueConstraint (MW.WriterT w) a = WriterConstraint
   valueConstraint _ _ = WriterConstraint

instance (Monoid w) => C (MW.WriterT w) where
   update = updateIfNew

instance (Inconsistency e, Eq a) => Value (E.ExceptionalT e) a where
   data ValueConstraint (E.ExceptionalT e) a =
           Eq a => ExceptionConstraint
   valueConstraint _ _ = ExceptionConstraint

instance (Inconsistency e) => C (E.ExceptionalT e) where
   update al av act =
      case valueConstraint al av of
         ExceptionConstraint ->
            updateAndCheck (\aold anew -> when (aold /= anew) break) al av act

simpleUpdate :: (C w, Value w a) => SimpleUpdater w s a
simpleUpdate al av = update al av . mapMaybeT UMT.lift


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


assignment2 ::
   UMT.C w =>
   (a -> b) ->
   Variable w s a -> Variable w s b ->
   T w s ()
assignment2 f (Variable _ al av) b =
   let triggerUpdate =
          varUpdate b $ fmap f $ readSTRefM av
   in  lift $
       modifySTRef al (addUpdate triggerUpdate)

assignment3 ::
   UMT.C w =>
   (a -> b -> 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 triggerUpdate =
          varUpdate c $
          liftM2 f (readSTRefM av) (readSTRefM bv)
   in  lift $
       modifySTRef al (addUpdate triggerUpdate) >>
       modifySTRef bl (addUpdate triggerUpdate)


newtype Apply w s a =
   Apply (Compose (MW.Writer [STRef s (Updates 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 :: Variable w s a -> Apply w s a
arg (Variable _update al av) =
   Apply $ Compose $ MW.writer (MaybeT $ readSTRef av, [al])

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

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


runApply ::
   UMT.C w =>
   Apply w s 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 (addUpdate $ varUpdate a f)