module Decision where

import Data.Maybe
import Control.Monad
import System.Random
import Data.IORef

lookup' :: (Eq a) => a -> [(a,b)] -> b
lookup' a abs = fromJust (lookup a abs)

type Decision = Integer
type Action = Char
type Variable = String
type Value = Char
type Environment = [(Variable, Value)]
type Probability = Double
type Utility = Double
type PEs = [(Probability, Environment)]
type Influence = Environment -> [Probability]

data Model = Reward Utility
           | Decide Decision [(Action, Model)]
           | Confuse Variable [Value] Influence Model
           | Discover Variable [(Value, Model)]

data Agent = Agent { ask  :: Decision -> [Action] -> IO Action,
                     tell :: Variable -> Value -> IO () }

data Policy = End
            | Do Action Policy
            | Switch Variable [(Value, Policy)]
    deriving (Eq, Show)

pickValue :: (Num prob, Ord prob) => prob -> [prob] -> [value] -> value
pickValue mass (p:ps) (v:vs) | remaining > 0 = pickValue remaining ps vs
                             | otherwise     = v
    where remaining = mass - p

normalize :: (Fractional prob) => [prob] -> [prob]
normalize ps = map (/ sum ps) ps

interactive :: Agent
interactive = Agent ask tell
    where ask decision actions = do
            putStr (show decision ++ " " ++ show actions ++ "? ")
            action <- readLn
            if action `elem` actions then return action else do
              putStrLn "That is not a valid action."
              ask decision actions
          tell variable value =
            putStrLn (show variable ++ " = " ++ show value)

policy :: Policy -> IO Agent
policy p = do r <- newIORef p
              let ask decision actions = do
                    putStr (show decision ++ " " ++ show actions ++ "? ")
                    Do action p <- readIORef r
                    writeIORef r p
                    print action
                    return action
                  tell variable value = do
                    putStrLn (show variable ++ " = " ++ show value)
                    Switch _ branches <- readIORef r
                    writeIORef r (lookup' value branches)
              return (Agent ask tell)

run :: Environment -> Agent -> Model -> IO Utility
run env agt (Reward utility) = return utility
run env agt (Decide decision choices) = do
    value <- ask agt decision (map fst choices)
    run env agt (lookup' value choices)
run env agt (Confuse variable values influence model) = do
    mass <- randomIO
    run ((variable, pickValue mass (influence env) values) : env) agt model
run env agt (Discover variable choices) = do
    let value = lookup' variable env
    tell agt variable value
    run env agt (lookup' value choices)

initialPEs :: PEs
initialPEs = [(1, [])]

confusePEs :: PEs -> Variable -> [Value] -> Influence -> PEs
confusePEs pes variable values influence =
    [ (prob * cond, (variable, value) : env)
    | (prob, env) <- pes, (value, cond) <- zip values (influence env) ]

switchPEs :: PEs -> Variable -> Value -> PEs
switchPEs pes variable value = filter ((value ==) . lookup' variable . snd) pes

utility :: PEs -> Policy -> Model -> Maybe Utility
utility pes End (Reward utility) =
    Just (utility * sum (map fst pes))
utility pes (Do action policy) (Decide _ thens) =
    lookup action thens >>= utility pes policy
utility pes policy (Confuse variable values influence model) =
    utility (confusePEs pes variable values influence) policy model
utility pes (Switch variable branches) (Discover variable' choices)
    | variable == variable' && length branches == length choices = do
    (ns, utilities) <- liftM unzip (zipWithM f branches choices)
    if sum ns == length pes then Just (sum utilities) else Nothing
    where f (value, policy) (value', model) = do
            unless (value == value') Nothing
            let pes' = switchPEs pes variable value
            u <- utility pes' policy model
            Just (length pes', u)
utility _ _ _ = Nothing

data Solution = Solution Policy Utility deriving (Eq, Show)

instance Ord Solution where
    compare (Solution _ u1) (Solution _ u2) = compare u1 u2

solve :: PEs -> Model -> Solution
solve pes (Reward utility) =
    Solution End (utility * sum (map fst pes))
solve pes (Decide _ thens) =
    maximum [ case solve pes model of
                Solution policy utility -> Solution (Do action policy) utility
            | (action, model) <- thens ]
solve pes (Confuse variable values influence model) =
    solve (confusePEs pes variable values influence) model
solve pes (Discover variable choices) =
    let solutions = [ solve (switchPEs pes variable value) model
                    | (value, model) <- choices ]
        collectBranch (value, _) (Solution policy _) = (value, policy)
    in Solution (Switch variable (zipWith collectBranch choices solutions))
                (sum [ utility | Solution _ utility <- solutions ])

observe :: Variable -> [(Value, Probability, Model)] -> Model
observe variable chances =
    Confuse variable
            [ value | (value,_,_) <- chances ]
            (const [ probability | (_,probability,_) <- chances ])
            (Discover variable [ (value,model) | (value,_,model) <- chances ])

among :: a -> [b] -> (a -> [(b, c)] -> w) -> (b -> c) -> w
among a bs maker k = maker a [ (b, k b) | b <- bs ]

dog :: Model
dog = Decide 1 [('f', Reward 0.2),
                ('w', observe "X" [('y', 0.8, Decide 2 [('f', Reward 0.7),
                                                        ('h', Reward 0.25)]),
                                   ('n', 0.2, Decide 3 [('f', Reward 0.1),
                                                        ('h', Reward 0.25)])]),
                ('h', Reward 0.3)]

montyHall :: Model
montyHall =
    Confuse "prize" doors
      (const (normalize [ 1 | prize <- doors ]))
      (among 1 doors Decide (\initial ->
        Confuse "open" doors
          (\env -> let prize = lookup' "prize" env
                   in normalize [ if open `elem` [initial, prize] then 0 else 1
                                | open <- doors ])
          (among "open" doors Discover (\open ->
            (among 2 doors Decide (\final ->
              (among "prize" doors Discover (\prize ->
                (Reward (if final == prize then 1 else 0))))))))))
    where doors = "abc"