MaxiMin algorithm for one player versus uncertainty

MaxiMin.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveFunctor  #-}
{-# LANGUAGE NamedFieldPuns #-}

{-# OPTIONS_GHC -Wall #-}

module MaxiMin where

import Data.List.NonEmpty (NonEmpty(..))
import Data.MemoTrie (HasTrie(..))

import qualified Control.Monad as Monad
import qualified Data.MemoTrie as MemoTrie
import qualified Data.List     as List
import qualified Data.Ord      as Ord

{-| A single possibility, consisting of an outcome paired with the associated
    weight of that outcome
-}
data Possibility a
    = Possibility { outcome :: !a, weight :: !Int }
    deriving (Functor, Show)

-- | A probability distribution, which is a non-empty list of weighted outcomes
newtype Distribution a
    = Distribution { possibilities :: NonEmpty (Possibility a) }
    deriving (Functor, Show)

instance Applicative Distribution where
    pure x = Distribution (pure (Possibility x 1))

    (<*>) = Monad.ap

instance Monad Distribution where
    m >>= f = Distribution do
        Possibility x weight0 <- possibilities m

        Possibility y weight1 <- possibilities (f x)

        return $! Possibility y (weight0 * weight1)

-- | Compute the expected value for a probability distribution
expectedValue :: Fractional n => Distribution n -> n
expectedValue Distribution{ possibilities } =
    totalTally / fromIntegral totalWeight
  where
    totalTally = sum (fmap tally possibilities)

    totalWeight = sum (fmap weight possibilities)

    tally Possibility{ outcome, weight } = fromIntegral weight * outcome

{-| Play the game optimally to its conclusion, always selecting the move that
    leads to the highest expected value for the given objective function
-}
play
    :: (Fractional n, Ord n, HasTrie state)
    => (state -> n)
    -- ^ Objective function, which returns the value we are trying to maximize
    -> (state -> Bool)
    -- ^ Termination function, which returns True if the game is over
    -> (state -> NonEmpty (Distribution state))
    -- ^ A function which generates the available moves from the current state
    -> state
    -- ^ The starting state
    -> Distribution state
    -- ^ The final probability distribution at the end of the game after optimal
    --   play
play objective done choices = MemoTrie.memoFix memoized
  where
    memoized loop status
        | done status = do
            pure status
        | otherwise = do
            next <- List.maximumBy (Ord.comparing predict) (choices status)

            loop next
      where
        predict choice = expectedValue do
            nextStatus <- choice

            finalStatus <- loop nextStatus

            return (objective finalStatus)