Combinators.hs

-- | Tests the iteration combinators.  The type signatures below pin
--   'GenNames' and 'ConcatMapTag' at compile time.
{-# LANGUAGE QualifiedDo #-}
module Test.Combinators (tests) where

import Prelude hiding (return, (>>=), (>>), flip)
import qualified Data.Map.Strict as Map
import Test.Tasty
import Test.Tasty.HUnit
import Test.Util (assertApprox, assertBounds, assertMap, knightNames)

import qualified Imp.DSL as Imp
import Imp.DSL (Imp, interval, observe)
import Imp.DSL.Combinators
import Imp.Inference
import Imp.Examples.IMDP as IMDP

-- | Must reproduce the hand-written 'IMDP.complexRobot', grade included.
robotFold2 :: Imp '["move1.b", "move1.f", "move2.b", "move2.f"] Position
robotFold2 = foldN @2 @"move" robotDynamics P0 step3

-- | A third step, which is the only way to reach @step3 P2 _@.
robotFold3 :: Imp '["move1.b", "move1.f", "move2.b", "move2.f"
                   , "move3.b", "move3.f"] Position
robotFold3 = foldN @3 @"move" robotDynamics P0 step3

-- | The trajectory rather than the endpoint.
robotScan2 :: Imp '["move1.b", "move1.f", "move2.b", "move2.f"] [Position]
robotScan2 = scanN @2 @"move" robotDynamics P0 step3

-- | The moves themselves, with no accumulator.
robotTag2 :: Imp '["move1.b", "move1.f", "move2.b", "move2.f"] [Move]
robotTag2 = tagN @2 @"move" robotDynamics

-- | Explicit-tag equivalents of the three above.
robotTagFold, robotTagScan :: Imp '["move1.b", "move1.f", "move2.b", "move2.f"] Position
robotTagFold = tagFold @'["move1", "move2"] robotDynamics P0 step3
robotTagScan = Imp.fmap last (tagScan @'["move1", "move2"] robotDynamics P0 step3)

robotTagMap :: Imp '["move1.b", "move1.f", "move2.b", "move2.f"] [Move]
robotTagMap = tagMap @'["move1", "move2"] robotDynamics

-- | A single imprecise step, reused by the numbered folds below.
oneStep :: Imp '["d"] Bool
oneStep = interval @"d" 0.4 0.6

-- | Count the @True@ steps.
countTrue :: Int -> Bool -> Int
countTrue acc x = if x then acc + 1 else acc

trueFold :: Imp '["t1.d", "t2.d", "t3.d"] Int
trueFold = foldN @3 @"t" oneStep 0 countTrue

trueScan :: Imp '["t1.d", "t2.d", "t3.d"] [Int]
trueScan = scanN @3 @"t" oneStep 0 countTrue

-- | 'foldMN' can condition at each step; here at most two steps may succeed.
cappedFold :: Imp '["t1.d", "t2.d", "t3.d"] Int
cappedFold = foldMN @3 @"t" oneStep 0 $ \acc x -> Imp.do
  let acc' = countTrue acc x
  observe (acc' < 3)
  Imp.return acc'

-- | 'GenNames' pads to the width of the largest index, keeping grades sorted.
names0 :: Imp '[] [Bool]
names0 = knightN @0 @"k"

names9 :: Imp '["k1","k2","k3","k4","k5","k6","k7","k8","k9"] [Bool]
names9 = knightN @9 @"k"

names10 :: Imp '["k01","k02","k03","k04","k05","k06","k07","k08","k09","k10"] [Bool]
names10 = knightN @10 @"k"

tests :: TestTree
tests = testGroup "Combinators"
  [ testGroup "mapName family"
    [ testCase "knightMap: one free choice per name" $
        assertMap "knightMap"
          [ ([False, False], (0, 1)), ([False, True], (0, 1))
          , ([True, False], (0, 1)), ([True, True], (0, 1)) ]
          (marginal (knightMap @'["x", "y"]))

    , testCase "intervalMap: two independent [0.3, 0.7] choices" $
        assertMap "intervalMap"
          [ ([False, False], (0.09, 0.49)), ([False, True], (0.09, 0.49))
          , ([True, False], (0.09, 0.49)), ([True, True], (0.09, 0.49)) ]
          (marginal (intervalMap @'["x", "y"] 0.3 0.7))

    , testCase "intervalN @3: P(all three) = [0.3^3, 0.7^3]" $ do
        assertBounds "all" (0.027, 0.343) (intervalProbability (intervalN @3 @"s" 0.3 0.7) and)
        assertBounds "any" (0.657, 0.973) (intervalProbability (intervalN @3 @"s" 0.3 0.7) or)

    , testCase "knightN @0 is the empty list at unit grade" $
        assertMap "knightN @0" [([], (1, 1))] (marginal names0)

    , testCase "knightN @3: 2^3 valuations, all unconstrained" $ do
        knightNames (knightN @3 @"k" :: Imp '["k1", "k2", "k3"] [Bool])
          @?= ["k1", "k2", "k3"]
        length (credalVertices (knightN @3 @"k" :: Imp '["k1", "k2", "k3"] [Bool])) @?= 8
    ]

  , testGroup "Tagged iteration over robotDynamics"
    [ testCase "foldN @2 reproduces complexRobot" $ do
        marginal robotFold2 @?= marginal IMDP.complexRobot
        length (credalVertices robotFold2) @?= 16
        assertBounds "P(P2)" (0.25, 0.64) (intervalProbability robotFold2 (== P2))

    , testCase "foldN @3: three steps, P2 absorbing" $
        assertMap "foldN @3"
          [ (P0, (0.008, 0.195)), (P1, (0.0832, 0.375)), (P2, (0.475, 0.896)) ]
          (marginal robotFold3)

    , testCase "scanN @2: trajectories, and P0 cannot jump to P2" $ do
        assertMap "scanN @2"
          [ ([P0, P0], (0.04, 0.25)), ([P0, P1], (0.1, 0.4))
          , ([P1, P0], (0.0, 0.08)),  ([P1, P1], (0.08, 0.4))
          , ([P1, P2], (0.25, 0.64)) ]
          (marginal robotScan2)
        assertBounds "last = P2" (0.25, 0.64)
          (intervalProbability robotScan2 ((== P2) . last))

    , testCase "tagN @2: the move pairs" $
        assertMap "tagN @2"
          [ ([Backwards, Backwards],  (0.0, 0.01))
          , ([Backwards, Stationary], (0.0, 0.05))
          , ([Backwards, Forwards],   (0.0, 0.08))
          , ([Stationary, Backwards], (0.0, 0.05))
          , ([Stationary, Stationary],(0.0256, 0.25))
          , ([Stationary, Forwards],  (0.08, 0.4))
          , ([Forwards, Backwards],   (0.0, 0.08))
          , ([Forwards, Stationary],  (0.08, 0.4))
          , ([Forwards, Forwards],    (0.25, 0.64)) ]
          (marginal robotTag2)

    , testCase "tagN @2: the list is in tag order, not reversed" $ do
        -- Identical steps make every bound symmetric under reversal,
        -- so using per-tag weights.
        let fwdThenBack ms = case ms of [Forwards, Backwards] -> 1.0; _ -> 0.0
            (weights, val) = optimizeExpectation robotTag2 fwdThenBack 500 0.5
        weights @?= Map.fromList
          [("move1.b", 0.5), ("move1.f", 1.0), ("move2.b", 1.0), ("move2.f", 0.0)]
        assertApprox "P(forwards then backwards)" 0.08 val

    , testCase "the numbered wrappers agree with the explicit-tag versions" $ do
        marginal robotTagFold @?= marginal robotFold2
        marginal robotTagMap @?= marginal robotTag2
        marginal robotTagScan @?= marginal robotFold2
    ]

  , testGroup "Numbered iteration over a single interval"
    [ testCase "foldN @3: count of successes" $
        assertMap "foldN @3"
          [ (0, (0.064, 0.216)), (1, (0.288, 0.432))
          , (2, (0.288, 0.432)), (3, (0.064, 0.216)) ]
          (marginal trueFold)

    , testCase "scanN @3: every path has the same bounds" $
        assertMap "scanN @3"
          [ ([0,0,0], (0.064, 0.216)), ([0,0,1], (0.064, 0.216))
          , ([0,1,1], (0.064, 0.216)), ([0,1,2], (0.064, 0.216))
          , ([1,1,1], (0.064, 0.216)), ([1,1,2], (0.064, 0.216))
          , ([1,2,2], (0.064, 0.216)), ([1,2,3], (0.064, 0.216)) ]
          (marginal trueScan)

    , testCase "foldMN @3: the per-step observe removes the all-successes path" $ do
        assertMap "foldMN @3"
          [ (0, (4/49, 3/13)), (1, (18/49, 6/13))
          , (2, (4/13, 27/49)), (3, (0, 0)) ]
          (marginal cappedFold)
        assertMap "foldMN @3 symbolic"
          [ (0, (4/49, 3/13)), (1, (18/49, 6/13))
          , (2, (4/13, 27/49)), (3, (0, 0)) ]
          (marginalSymbolic cappedFold)

    , testCase "foldMN @3: tag names join with a dot, one per step" $ do
        let (weights, prob) = optimizeProbability cappedFold (== 2) 200 0.1
        weights @?= Map.fromList [("t1.d", 1.0), ("t2.d", 1.0), ("t3.d", 1.0)]
        assertApprox "ascent reaches the exact upper bound" (27/49) prob
    ]

  , testGroup "GenNames padding"
    [ testCase "one digit up to 9, two digits from 10" $ do
        knightNames names0 @?= []
        knightNames names9 @?= ["k1","k2","k3","k4","k5","k6","k7","k8","k9"]
        knightNames names10
          @?= ["k01","k02","k03","k04","k05","k06","k07","k08","k09","k10"]
    ]
  ]