    {-# LANGUAGE TypeFamilies #-}

    module Lambdabot.Plugin.Novelty.Numberwang (numberwangPlugin) where

    import Control.Applicative
    import Control.Monad
    import Data.Random
    import Data.Random.Distribution.Poisson
    import Lambdabot.Plugin
    import Numeric

    data NumberwangState = State
        { nextCmd   :: !Int -- number of invocations of @numberwang before the next numberwang
        , nextCon   :: !Int -- number of contextual occurrences of numbers before next numberwang
        }

    cmdDist :: RVar Int
    cmdDist = poisson (3.5 :: Double)

    conDist :: RVar Int
    conDist = poisson (32  :: Double)

    numberwangPlugin :: Module NumberwangState
    numberwangPlugin = newModule
        { moduleDefState = sample (State <$> cmdDist <*> conDist)
        , moduleCmds = return
            [ (command "numberwang")
                { help = say "@numberwang <number>: Determines if it is Numberwang."
                , process = doNumberwang True . length . words
                }
            ]
        , contextual = doNumberwang False . length . (numbers :: String -> [Double])
        }

    numbers :: RealFrac t => String -> [t]
    numbers [] = []
    numbers cs = case readFloat cs of
        (n, rest):_ -> n : numbers rest
        _           -> numbers (tail cs)

    doNumberwang :: (Num a, Ord a, MonadLBState m, LBState m ~ NumberwangState) =>
                    Bool -> a -> Cmd m ()
    doNumberwang cmd n
        | n <= 0    = when cmd $ say "What number?"
        | otherwise = do
            isNumberwang <- checkNumberwang cmd 1
            if isNumberwang
                then say "That's Numberwang!"
                else when cmd $ say "Sorry, that's not Numberwang."

    withState :: (MonadLBState m, LBState m ~ NumberwangState) =>
                 Bool -> (Int -> (Int -> m ()) -> RVar Int -> m a) -> m a
    withState True f = withMS $ \st setST ->
        f (nextCmd st) (\n -> setST st {nextCmd = n}) cmdDist
    withState False f = withMS $ \st setST ->
        f (nextCon st) (\n -> setST st {nextCon = n}) conDist

    checkNumberwang :: (MonadLBState m, LBState m ~ NumberwangState) =>
                       Bool -> Int -> m Bool
    checkNumberwang cmd l = withState cmd $ \ n setN nDist -> do
        if n <= l
            then do
                setN =<< lb (sample nDist)
                return True
            else do
                setN (n - l)
                return False