{-# 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.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)

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