{-# LANGUAGE RecordWildCards, RankNTypes, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, ViewPatterns, TypeFamilies #-}

module Sequence.Formula
  ( FormulaM, Formula, quot'
  , evalFormula
  , val
  , d, z
  ) where

import Control.Lens
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.Set (Set)
import qualified Data.Set as Set


type FormulaM input a = StateT (Set String) (ReaderT 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 => input -> (Question input) -> m input
askQuestion input q@(Question{..}) = flip (set answer) input . maybe (throwError q) return <$> askQ prompt (join . fmap readMaybe)

evalFormula :: (MonadIO m, Ord a, Show a) => input -> FormulaM input a -> m (input, a)
evalFormula = evalFormula' []
  where
    evalFormula' finalChanges input formula = do
      result <- liftIO . enact . runExceptT . (runReaderT ?? input) . (evalStateT ?? Set.empty) $ formula
      case result of
        Left q@(Question{..}) -> askQuestion input q >>= flip (evalFormula' $ bool (pure . set answer $ throwError q) mempty keepResult ++ finalChanges) formula
        Right result -> return (foldr ($) input finalChanges, result)

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 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