{-# 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'
  , 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

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