summaryrefslogtreecommitdiff
path: root/src/Sequence/Formula.hs
blob: 878ec7f83fbeea230186800fd4c3b374da7bc89e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
{-# LANGUAGE RecordWildCards, RankNTypes, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, ViewPatterns, TypeFamilies, GADTs, TypeOperators, ExistentialQuantification, FlexibleContexts, IncoherentInstances #-}

module Sequence.Formula
  ( FormulaM, Formula, quot'
  , (:<:)(..), Context(..), ctx
  , evalFormula, evalFormula'
  , findDistribution, findDistribution'
  , findAverage
  , val
  , d, z
  , Table, table
  ) where

import Control.Lens hiding (Context(..))
import Data.Data.Lens

import Control.Monad
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Base
import Numeric.Probability.Game.Event
import qualified Numeric.Probability.Game.Dice as D

import Sequence.Utils.Ask

import Text.Read (readMaybe)

import Data.Bool
import Data.List
import Data.Maybe
import Data.Either
import Data.Tuple

import Data.Map (Map)
import qualified Data.Map as Map

import Data.Set (Set)
import qualified Data.Set as Set

class (:<:) small large where
  ctx' :: Traversal' large small

instance a :<: a where
  ctx' = simple

instance a :<: (a, a) where
  ctx' = both

instance a :<: b => a :<: (b, c) where
  ctx' = _1 . ctx'

instance a :<: (c, a) where
  ctx' = _2

instance () :<: a where
  ctx' = united

instance a :<: b => a :<: Maybe b where
  ctx' = _Just . ctx'

data Context small = forall large. (small :<: large) => Context large (Maybe (Formula small))

ctx :: Traversal' (Context input) input
ctx modifySmall (Context large fSt) = flip Context fSt <$> ctx' modifySmall large

ctxStore :: Traversal' (Context input) (Formula input)
ctxStore modifyF (Context large fSt) = Context large <$> _Just modifyF fSt

type FormulaM input a = StateT (Set [String]) (ReaderT (Context input) (ExceptT (Question input) EventM)) a

type Formula input = FormulaM input Int

instance MonadBase EventM EventM where
  liftBase = id

data Question input = Question
                      { answer :: Traversal' input (Formula input)
                      , prompt :: [String]
                      , keepResult :: Bool
                      }

instance Eq (Question a) where
  (==) _ _ = True

instance Ord (Question a) where
  (<=) _ _ = True

instance Show (Question a) where
  show Question{..} = show prompt

instance Integral a => Num (FormulaM input a) where
  (+) x y = (+) <$> x <*> y
  (-) x y = (-) <$> x <*> y
  negate = fmap negate
  abs = fmap abs
  signum = fmap signum
  (*) x y = do n <- x
               sum <$> replicateM (fromIntegral n) y
  fromInteger = return . fromInteger

quot' :: Integral a => FormulaM input a -> FormulaM input a -> FormulaM input a
quot' = liftM2 quot

askQuestion :: (MonadIO m, sInput :<: lInput) => [String] -> (lInput, Maybe (Formula sInput)) -> Question sInput -> m (lInput, Maybe (Formula sInput))
askQuestion promptPref input q@(Question{..}) = flip (if keepResult then set $ _1 . ctx' . answer else set _2 . Just) input . maybe (throwError q) return <$> askQ (wPromptPref $ promptPref ++ prompt) (join . fmap readMaybe)
  where
    wPromptPref [] = " " ++ sep
    wPromptPref [x] = x ++ " " ++ sep
    wPromptPref (x:xs) = x ++ " " ++ sep ++ " " ++ wPromptPref xs
    sep = "»"

evalFormula :: (MonadIO m, sInput :<: lInput) => [String] -> lInput -> FormulaM sInput a -> m (lInput, a)
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
    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' 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

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{..})
  preview (if keepResult then ctx . answer else ctxStore) >>= (modify (Set.delete prompt) >>) . fromMaybe (throwError Question{..})

d, z :: Integral a => Int -> FormulaM input a
d n = liftBase . fmap fromIntegral $ D.d n
z n = liftBase . fmap fromIntegral $ D.z n

type Table a = Map a Rational

table :: Ord a => Table a -> FormulaM input a
table = liftBase . makeEventProb . Map.assocs