content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  module OpenGames.Engine.BayesianGamesNonState
    ( StochasticBayesianOpenGame (..),
      dependentDecision,
      dependentEpsilonDecision,
      fromLens,
      fromFunctions,
      nature,
      liftStochastic,
      uniformDist,
      distFromList,
      pureAction,
      playDeterministically,
    )
  where
  
  import Control.Arrow hiding ((+:+))
  import Data.Foldable
  import Data.Ord (comparing)
  import Numeric.Probability.Distribution hiding (filter, lift, map)
  import OpenGames.Engine.Diagnostics
  import OpenGames.Engine.OpenGames hiding (lift)
  import OpenGames.Engine.OpticClass
  import OpenGames.Engine.TLL
  
  ---------------------------------------------
  -- Reimplements stateful bayesian from before
  
  type StochasticBayesianOpenGame a b x s y r = OpenGame StochasticOptic StochasticContext a b x s y r
  
  type Agent = String
  
  support :: Stochastic x -> [x]
  support = map fst . decons
  
  bayes :: (Eq y) => Stochastic (x, y) -> y -> Stochastic x
  bayes a y = mapMaybe (\(x, y') -> if y' == y then Just x else Nothing) a
  
  deviationsInContext ::
    (Show x, Show y, Ord y, Show theta) =>
    Double ->
    Agent ->
    x ->
    theta ->
    Stochastic y ->
    (y -> Double) ->
    [y] ->
    [DiagnosticInfoBayesian x y]
  deviationsInContext epsilon name x theta strategy u ys =
    [ DiagnosticInfoBayesian
        { equilibrium = strategicPayoff >= optimalPayoff - epsilon,
          player = name,
          payoff = strategicPayoff,
          optimalMove = optimalPlay,
          optimalPayoff = optimalPayoff,
          context = u,
          state = x,
          unobservedState = show theta,
          strategy = strategy
        }
    ]
    where
      strategicPayoff = expected (fmap u strategy)
      (optimalPlay, optimalPayoff) = maximumBy (comparing snd) [(y, u y) | y <- ys]
  
  dependentDecision :: (Eq x, Show x, Ord y, Show y) => String -> (x -> [y]) -> StochasticBayesianOpenGame '[Kleisli Stochastic x y] '[[DiagnosticInfoBayesian x y]] x () y Double
  dependentDecision name ys =
    OpenGame
      { play = \(a :- Nil) ->
          let v x = do y <- runKleisli a x; return ((), y)
              u () _ = return ()
           in StochasticOptic v u,
        evaluate = \(a :- Nil) (StochasticContext h k) ->
          ( concat
              [ let u y =
                      expected
                        ( do
                            t <- (bayes h x)
                            k t y
                        )
                    strategy = runKleisli a x
                 in deviationsInContext 0 name x theta strategy u (ys x)
                | (theta, x) <- support h
              ]
          )
            :- Nil
      }
  
  dependentEpsilonDecision :: (Eq x, Show x, Ord y, Show y) => Double -> String -> (x -> [y]) -> StochasticBayesianOpenGame '[Kleisli Stochastic x y] '[[DiagnosticInfoBayesian x y]] x () y Double
  dependentEpsilonDecision epsilon name ys =
    OpenGame
      { play = \(a :- Nil) ->
          let v x = do y <- runKleisli a x; return ((), y)
              u () _ = return ()
           in StochasticOptic v u,
        evaluate = \(a :- Nil) (StochasticContext h k) ->
          ( concat
              [ let u y =
                      expected
                        ( do
                            t <- (bayes h x)
                            k t y
                        )
                    strategy = runKleisli a x
                 in deviationsInContext epsilon name x theta strategy u (ys x)
                | (theta, x) <- support h
              ]
          )
            :- Nil
      }
  
  -- Support functionality for constructing open games
  fromLens :: (x -> y) -> (x -> r -> s) -> StochasticBayesianOpenGame '[] '[] x s y r
  fromLens v u =
    OpenGame
      { play = \Nil -> StochasticOptic (\x -> return (x, v x)) (\x r -> return (u x r)),
        evaluate = \Nil _ -> Nil
      }
  
  fromFunctions :: (x -> y) -> (r -> s) -> StochasticBayesianOpenGame '[] '[] x s y r
  fromFunctions f g = fromLens f (const g)
  
  nature :: Stochastic x -> StochasticBayesianOpenGame '[] '[] () () x ()
  nature a =
    OpenGame
      { play = \Nil -> StochasticOptic (\() -> do x <- a; return ((), x)) (\() () -> return ()),
        evaluate = \Nil _ -> Nil
      }
  
  liftStochastic :: (x -> Stochastic y) -> StochasticBayesianOpenGame '[] '[] x () y ()
  liftStochastic f =
    OpenGame
      { play = \Nil -> StochasticOptic (\x -> do y <- f x; return ((), y)) (\() () -> return ()),
        evaluate = \_ _ -> Nil
      }
  
  -- Support functionality for stochastic processes (also interface to the probability module in use)
  
  -- uniform distribution
  uniformDist = uniform
  
  -- tailored distribution from a list
  distFromList = fromFreqs
  
  -- pure action (no randomization)
  pureAction x = Kleisli $ const $ certainly x
  
  playDeterministically :: a -> Stochastic a
  playDeterministically = certainly
  