{-# LANGUAGE RecordWildCards, RankNTypes, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, ViewPatterns, TypeFamilies, GADTs, TypeOperators, ExistentialQuantification, FlexibleContexts #-} module Sequence.Formula ( FormulaM, Formula, quot' , (:<:)(..), Context(..), ctx , evalFormula , 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.Set (Set) import qualified Data.Set as Set import Data.Map (Map) import qualified Data.Map as Map class (:<:) small large where ctx' :: Traversal' large small instance a :<: a where ctx' = simple instance a :<: (a, a) where ctx' = both instance a :<: (a, b) where ctx' = _1 instance a :<: (b, a) where ctx' = _2 instance () :<: a where ctx' = united data Context small = forall large. (small :<: large) => Context large ctx :: Traversal' (Context input) input ctx modifySmall (Context large) = Context <$> ctx' modifySmall large 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 -> Question sInput -> m lInput askQuestion promptPref input q@(Question{..}) = flip (set $ ctx' . answer) input . maybe (throwError q) return <$> askQ (promptPref' ++ prompt) (join . fmap readMaybe) where promptPref' | null promptPref = "" | otherwise = promptPref ++ " ยป " evalFormula :: (MonadIO m, sInput :<: lInput) => String -> lInput -> FormulaM sInput a -> m (lInput, a) evalFormula = evalFormula' [] where evalFormula' finalChanges promptPref input formula = do result <- liftIO . enact . runExceptT . (runReaderT ?? (Context input)) . (evalStateT ?? Set.empty) $ formula case result of Left q@(Question{..}) -> askQuestion promptPref input q >>= flip (flip evalFormula' promptPref $ bool (pure . set (ctx' . answer) $ throwError q) mempty keepResult ++ 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 val :: Integral a => 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 (ctx . answer) >>= maybe (throwError Question{..}) id 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