Print all applicable values

2 files changed, 5 insertions, 0 deletions

diff --git a/src/Sequence/Formula.hs b/src/Sequence/Formula.hs
index 2ac1210..878ec7f 100644
--- a/src/Sequence/Formula.hs
+++ b/src/Sequence/Formula.hs
@@ -5,6 +5,7 @@ module Sequence.Formula
   , (:<:)(..), Context(..), ctx
   , evalFormula, evalFormula'
   , findDistribution, findDistribution'
+  , findAverage
   , val
   , d, z
   , Table, table
@@ -137,6 +138,9 @@ evalFormula' promptPref formula = uncurry (<$) . over _2 put . swap =<< flip (ev
 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
 
+findAverage :: (MonadIO m, sInput :<: lInput, MonadState lInput m, Real a) => [String] -> FormulaM sInput a -> m Rational
+findAverage promptPref formula = sum . map (\(val, prob) -> toRational val * prob) . Map.toList <$> findDistribution' promptPref formula
+
 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 f51bcd8..9f03ca7 100644
--- a/src/Sequence/Utils.hs
+++ b/src/Sequence/Utils.hs
@@ -7,6 +7,7 @@ module Sequence.Utils
   , toName, toDesc
   , outputLogged
   , scaleTimer
+  , statAccessors
   , Argument(..)
   , Completion(..)
   , module Sequence.Utils.Ask