Examples.hs

-- | Facts about the example programs in Imp.Examples.*
module Test.Examples (tests) where

import Data.List (nub, sort)
import qualified Data.Map.Strict as Map
import Test.Tasty
import Test.Tasty.HUnit
import Test.Util (assertApprox, assertBounds, assertDist, assertMap, roundDist)
import Imp.Inference
import Imp.Examples.Basic as Basic
import Imp.Examples.Ellsberg as E
import Imp.Examples.IMDP as IMDP
import Imp.Examples.Iteration as Iter
import Imp.Examples.Knightian as K
import Imp.Examples.MontyHall as MH
import Imp.Examples.Polytope as P
import Imp.Examples.TwoChild as TC

tests :: TestTree
tests = testGroup "Examples"
  [ testGroup "Basic"
    [ testCase "fairCoin" $
        assertDist "fairCoin" [(False, 0.5), (True, 0.5)] (preciseMarginal Basic.fairCoin)

    , testCase "biasedCoin" $
        assertDist "biasedCoin" [(False, 0.3), (True, 0.7)]
          (preciseMarginal Basic.biasedCoin)

    , testCase "twoCoins: uniform over the four pairs" $
        assertDist "twoCoins"
          [ ((False, False), 0.25), ((False, True), 0.25)
          , ((True, False), 0.25), ((True, True), 0.25) ]
          (preciseMarginal Basic.twoCoins)

    , testCase "coinOr: P(True) = 0.75" $
        assertDist "coinOr" [(False, 0.25), (True, 0.75)] (preciseMarginal Basic.coinOr)
    ]

  , testGroup "Knightian names"
    [ testCase "dependent: line segment, 2 vertices" $ do
        let verts = credalVertices K.dependent
        length verts @?= 2
        sort (map roundDist verts) @?= sort
          [ Map.fromList [(K.Blue, 0.0), (K.Green, 0.0), (K.Red, 1.0)]
          , Map.fromList [(K.Blue, 0.5), (K.Green, 0.5), (K.Red, 0.0)]
          ]

    , testCase "independent: quadrilateral, 4 vertices" $ do
        let verts = credalVertices K.independent
        length verts @?= 4
        sort (map roundDist verts) @?= sort
          [ Map.fromList [(K.Blue, 0.0), (K.Green, 0.0), (K.Red, 1.0)]
          , Map.fromList [(K.Blue, 0.5), (K.Green, 0.0), (K.Red, 0.5)]
          , Map.fromList [(K.Blue, 0.0), (K.Green, 0.5), (K.Red, 0.5)]
          , Map.fromList [(K.Blue, 0.5), (K.Green, 0.5), (K.Red, 0.0)]
          ]

    , testCase "independent is strictly larger than dependent" $ do
        let dep = map roundDist (credalVertices K.dependent)
            ind = map roundDist (credalVertices K.independent)
        all (`elem` ind) dep @? "every dependent vertex is an independent one"
        length (nub ind) > length (nub dep) @? "independent has strictly more vertices"

    , testCase "sharing a name does not change the marginals, only the joint" $ do
        let bounds = [ (K.Red, (0.0, 1.0)), (K.Green, (0.0, 0.5)), (K.Blue, (0.0, 0.5)) ]
        assertMap "dependent" bounds (marginal K.dependent)
        assertMap "independent" bounds (marginal K.independent)

    , testCase "dependent: E[score]" $
        assertBounds "E" (0.25, 1.0)
          (intervalExpectation K.dependent
             (\v -> case v of K.Red -> 1.0; K.Green -> 0.5; K.Blue -> 0.0))
    ]

  , testGroup "MontyHall"
    [ testCase "P(switch wins) = 2/3 regardless of host bias" $
        assertMap "montyHall" [(False, (1/3, 1/3)), (True, (2/3, 2/3))]
          (marginal MH.montyHall)
    ]

  , testGroup "TwoChild"
    [ testCase "P(both boys | at least one boy)" $
        assertMap "twoChild" [(False, (0.5, 1.0)), (True, (0.0, 0.5))]
          (marginal TC.twoChild)

    , testCase "credal vertices: P(True) takes values 0 and 1/2" $ do
        let verts = credalVertices TC.twoChild
        sort (map (Map.findWithDefault 0 True) verts) @?= [0.0, 0.5]
    ]

  , testGroup "Polytope"
    [ testCase "polytope: 4 vertices, all summing to 1" $ do
        let verts = credalVertices P.polytope
        length verts @?= 4
        mapM_ (\dist -> assertApprox "sum" 1.0 (sum dist)) verts

    , testCase "polytope: marginal bounds" $
        assertMap "polytope"
          [ (P.Red, (0.2, 0.8)), (P.Green, (0.06, 0.56)), (P.Blue, (0.06, 0.56)) ]
          (marginal P.polytope)

    , testCase "polytope2: 8 vertices, all distinct" $ do
        let verts = credalVertices P.polytope2
        length verts @?= 8
        length (nub (map roundDist verts)) @?= 8

    , testCase "polytope2: Red picks up the fall-through branch" $
        assertMap "polytope2"
          [ (P.Red, (0.4, 0.82)), (P.Green, (0.1, 0.4)), (P.Blue, (0.05, 0.32)) ]
          (marginal P.polytope2)
    ]

  , testGroup "Iteration"
    [ testCase "walk1: 2 valuations, 2 distinct vertices" $ do
        let verts = credalVertices Iter.walk1
        length verts @?= 2
        length (nub (map roundDist verts)) @?= 2
        assertMap "walk1" [(0, (0.3, 0.7)), (1, (0.3, 0.7))] (marginal Iter.walk1)

    , testCase "walk2: 4 valuations, 3 distinct" $ do
        let verts = credalVertices Iter.walk2
        length verts @?= 4
        length (nub (map roundDist verts)) @?= 3
        assertMap "walk2" [(0, (0.09, 0.49)), (1, (0.42, 0.58)), (2, (0.09, 0.49))]
          (marginal Iter.walk2)

    , testCase "walk3: 8 valuations, 4 distinct" $ do
        let verts = credalVertices Iter.walk3
        length verts @?= 8
        length (nub (map roundDist verts)) @?= 4
        assertMap "walk3"
          [ (0, (0.027, 0.343)), (1, (0.189, 0.469))
          , (2, (0.189, 0.469)), (3, (0.027, 0.343)) ]
          (marginal Iter.walk3)

    , testCase "walk3Asym: 8 valuations, 8 distinct" $ do
        let verts = credalVertices Iter.walk3Asym
        length verts @?= 8
        length (nub (map roundDist verts)) @?= 8
        assertMap "walk3Asym"
          [ (0, (0.0135, 0.3465)), (1, (0.1695, 0.4995))
          , (2, (0.1695, 0.4995)), (3, (0.0135, 0.3465)) ]
          (marginal Iter.walk3Asym)
    ]

  , testGroup "Ellsberg paradox"
    [ testCase "credal set: 2 vertices" $
        length (credalVertices E.ellsberg) @?= 2

    , testCase "Red is precise, Black and Yellow are maximally imprecise" $
        assertMap "ellsberg"
          [ (E.Red, (1/3, 1/3)), (E.Black, (0.0, 2/3)), (E.Yellow, (0.0, 2/3)) ]
          (marginal E.ellsberg)

    , testCase "P(Black or Yellow) = [2/3, 2/3]" $
        assertBounds "P"  (2/3, 2/3)
          (intervalProbability E.ellsberg (\b -> b == E.Black || b == E.Yellow))

    , testCase "ambiguity aversion: known gambles are precise" $ do
        assertBounds "I (bet Red)" (1/3, 1/3)
          (intervalExpectation E.ellsberg (\b -> if b == E.Red then 1.0 else 0.0))
        assertBounds "II (bet Black)" (0.0, 2/3)
          (intervalExpectation E.ellsberg (\b -> if b == E.Black then 1.0 else 0.0))
        assertBounds "IV (bet Black or Yellow)" (2/3, 2/3)
          (intervalExpectation E.ellsberg
             (\b -> if b == E.Black || b == E.Yellow then 1.0 else 0.0))
    ]

  , testGroup "Interval MDP"
    [ testCase "simpleRobot: 4 valuations, 3 distinct" $ do
        let verts = credalVertices IMDP.simpleRobot
        length verts @?= 4
        length (nub (map roundDist verts)) @?= 3

    , testCase "simpleRobot: marginal bounds" $
        assertMap "simpleRobot"
          [ (IMDP.P0, (0.01, 0.16)), (IMDP.P1, (0.18, 0.48)), (IMDP.P2, (0.36, 0.81)) ]
          (marginal IMDP.simpleRobot)

    , testCase "simpleRobot3: 8 valuations, 4 distinct" $ do
        let verts = credalVertices IMDP.simpleRobot3
        length verts @?= 8
        length (nub (map roundDist verts)) @?= 4

    , testCase "simpleRobot3: E[position]" $ do
        assertMap "simpleRobot3"
          [ (IMDP.P0, (0.001, 0.064)), (IMDP.P1, (0.027, 0.288))
          , (IMDP.P2, (0.648, 0.972)) ]
          (marginal IMDP.simpleRobot3)
        assertBounds "E[position]" (1.584, 1.971)
          (intervalExpectation IMDP.simpleRobot3
             (\v -> case v of IMDP.P0 -> 0; IMDP.P1 -> 1; IMDP.P2 -> 2))
    ]

  , testGroup "Compositional robot"
    [ testCase "complexRobot: 16 valuations" $ do
        let verts = credalVertices IMDP.complexRobot
        length verts @?= 16
        mapM_ (\dist -> assertApprox "sum" 1.0 (sum dist)) verts

    , testCase "complexRobot: exact marginal bounds" $
        assertMap "complexRobot"
          [ (IMDP.P0, (0.04, 0.3)), (IMDP.P1, (0.288, 0.5)), (IMDP.P2, (0.25, 0.64)) ]
          (marginal IMDP.complexRobot)

    , testCase "complexRobot: P(reach goal) = [0.25, 0.64]" $ do
        assertBounds "enumeration" (0.25, 0.64)
          (intervalProbability IMDP.complexRobot (== IMDP.P2))
        assertBounds "symbolic" (0.25, 0.64)
          (intervalProbabilitySymbolic IMDP.complexRobot (== IMDP.P2))
    ]
  ]