2 files changed, 35 insertions, 6 deletions

diff --git a/src/Sequence/Formula.hs b/src/Sequence/Formula.hs
index a3675c6..2ac1210 100644
--- a/src/Sequence/Formula.hs
+++ b/src/Sequence/Formula.hs
@@ -4,6 +4,7 @@ module Sequence.Formula
   ( FormulaM, Formula, quot'
   , (:<:)(..), Context(..), ctx
   , evalFormula, evalFormula'
+  , findDistribution, findDistribution'
   , val
   , d, z
   , Table, table
@@ -109,18 +110,33 @@ askQuestion promptPref input q@(Question{..}) = flip (if keepResult then set $ _
     sep = "»"
 
 evalFormula :: (MonadIO m, sInput :<: lInput) => [String] -> lInput -> FormulaM sInput a -> m (lInput, a)
-evalFormula promptPref input = evalFormula' [] promptPref (input, Nothing)
+evalFormula = primEvalFormula $ liftIO . enact
+
+findDistribution :: (MonadIO m, sInput :<: lInput, Ord a) => [String] -> lInput -> FormulaM sInput a -> m (lInput, (Map a Rational))
+findDistribution = primEvalFormula $ return . fmap Map.fromList . seqEither . outcomes
   where
-    evalFormula' :: (MonadIO m, sInput :<: lInput) => [lInput -> lInput] -> [String] -> (lInput, Maybe (Formula sInput)) -> FormulaM sInput a -> m (lInput, a)
-    evalFormula' finalChanges promptPref (input, fSt) formula = do
-      result <- liftIO . enact . runExceptT . (runReaderT ?? (Context input fSt)) . (evalStateT ?? Set.empty) $ formula
+    seqEither :: [(Either q a, Rational)] -> Either q [(a, Rational)]
+    seqEither = mapM seqEither'
+
+    seqEither' (Left  a, _) = Left a
+    seqEither' (Right b, c) = Right (b, c)
+
+primEvalFormula :: (MonadIO m, sInput :<: lInput) => (EventM (Either (Question sInput) a) -> m (Either (Question sInput) r)) -> [String] -> lInput -> FormulaM sInput a -> m (lInput, r)
+primEvalFormula fromOutcomes promptPref input = evalFormula' fromOutcomes [] promptPref (input, Nothing)
+  where
+    evalFormula' :: (MonadIO m, sInput :<: lInput) => (EventM (Either (Question sInput) a) -> m (Either (Question sInput) r)) -> [lInput -> lInput] -> [String] -> (lInput, Maybe (Formula sInput)) -> FormulaM sInput a -> m (lInput, r)
+    evalFormula' fromOutcomes finalChanges promptPref (input, fSt) formula = do
+      result <- fromOutcomes . runExceptT . (runReaderT ?? (Context input fSt)) . (evalStateT ?? Set.empty) $ formula
       case result of
-        Left q@(Question{..}) -> askQuestion promptPref (input, fSt) q >>= flip (flip evalFormula' promptPref $ set (ctx' . answer) (throwError q) : finalChanges) formula
+        Left q@(Question{..}) -> askQuestion promptPref (input, fSt) q >>= flip (flip (evalFormula' fromOutcomes) promptPref $ set (ctx' . answer) (throwError q) : finalChanges) formula
         Right result -> return (foldr ($) input finalChanges, result)
 
 evalFormula' :: (MonadIO m, sInput :<: lInput, MonadState lInput m) => [String] -> FormulaM sInput a -> m a
 evalFormula' promptPref formula = uncurry (<$) . over _2 put . swap =<< flip (evalFormula promptPref) formula =<< get
 
+findDistribution' :: (MonadIO m, sInput :<: lInput, MonadState lInput m, Ord a) => [String] -> FormulaM sInput a -> m (Map a Rational)
+findDistribution' promptPref formula = uncurry (<$) . over _2 put . swap =<< flip (findDistribution promptPref) formula =<< get
+
 val :: Traversal' input (Formula input) -> [String] -> Bool -> Formula input
 val answer prompt keepResult = do
   gets (Set.member prompt) >>= bool (modify $ Set.insert prompt) (modify (Set.delete prompt) >> throwError Question{..})
diff --git a/src/Sequence/Utils.hs b/src/Sequence/Utils.hs
index 929189e..ba6ce95 100644
--- a/src/Sequence/Utils.hs
+++ b/src/Sequence/Utils.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE ViewPatterns, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, UndecidableInstances, OverloadedStrings, RankNTypes #-}
+{-# LANGUAGE ViewPatterns, FlexibleContexts, MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, UndecidableInstances, OverloadedStrings, OverloadedLists, RankNTypes #-}
 
 module Sequence.Utils
   ( withArg, withFocus, withFocus'
@@ -140,3 +140,16 @@ instance Argument DamageType GameState where
   arg (CI.mk -> word) = return $ Map.lookup word types
     where
       types = Map.fromList [(CI.mk $ show dType, dType) | dType <- [minBound .. maxBound]]
+
+instance Completion (Formula Stats) GameState where
+  completableLabel _ = "<value>"
+  complete _ st (CI.foldCase -> prefix) = return . map CI.original . filter ((prefix `isPrefixOf`) . CI.foldedCase) . Map.keys $ Map.filter (isJust . (\a -> preview (gFocus' . eStats . folding a) st)) statAccessors
+
+instance Argument (Formula Stats) GameState where
+  arg (CI.mk -> name) = runMaybeT $ do
+    accessor <- MaybeT . return $ Map.lookup name statAccessors
+    MaybeT . preuse $ gFocus' . eStats . folding accessor
+
+statAccessors :: Map (CI String) (Stats -> Maybe (Formula Stats))
+statAccessors = [
+                ]