{-# LANGUAGE RecordWildCards, RankNTypes, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, ViewPatterns, TypeFamilies, GADTs, TypeOperators, ExistentialQuantification, FlexibleContexts #-}

module Sequence.Formula
  ( FormulaM, Formula, quot'
  , (:<:)(..), Context(..), ctx
  , evalFormula, 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