Util.hs

-- | Assertion and inspection helpers shared by the test modules.
module Test.Util
  ( assertApprox
  , assertBounds
  , assertMap
  , assertDist
  , assertContains
  , assertNoFeasible
  , knightNames
  , roundDist
  ) where

import Control.Exception (ErrorCall, evaluate, try)
import Data.List (isPrefixOf)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Test.Tasty.HUnit

import Imp.DSL (Imp)
import Imp.Inference (optimizeProbability)

-- | The Knightian variable names a program allocates, as the compiler sees them.
knightNames :: Ord a => Imp g a -> [String]
knightNames prog = Map.keys (fst (optimizeProbability prog (const True) 1 0.1))

-- | Round a distribution's probabilities to 3 decimal places for comparison.
roundDist :: Map a Double -> Map a Double
roundDist = fmap (\v -> fromIntegral (round (v * 1000) :: Int) / 1000)

-- | Assert that two Doubles are approximately equal (within 1e-6).
assertApprox :: String -> Double -> Double -> Assertion
assertApprox label expected actual =
  abs (actual - expected) < 1e-6 @?
    (label ++ ": expected " ++ show expected ++ ", got " ++ show actual)

-- | Assert a lower/upper bound pair.
assertBounds :: String -> (Double, Double) -> (Double, Double) -> Assertion
assertBounds label (elo, ehi) (lo, hi) = do
  assertApprox (label ++ " lower") elo lo
  assertApprox (label ++ " upper") ehi hi

-- | Assert the exact key list and pair of a bounds map.
assertMap :: (Ord k, Show k)
          => String -> [(k, (Double, Double))] -> Map k (Double, Double) -> Assertion
assertMap label expected actual = do
  Map.keys actual @?= map fst expected
  mapM_ (\(k, b) -> assertBounds (label ++ " " ++ show k) b (actual Map.! k)) expected

-- | Assert the exact key list and probability of a distribution map.
assertDist :: (Ord k, Show k) => String -> [(k, Double)] -> Map k Double -> Assertion
assertDist label expected actual = do
  Map.keys actual @?= map fst expected
  mapM_ (\(k, p) -> assertApprox (label ++ " " ++ show k) p (actual Map.! k)) expected

-- | Assert that the second pair of bounds contains the first.
assertContains :: String -> (Double, Double) -> (Double, Double) -> Assertion
assertContains label (ilo, ihi) (olo, ohi) = do
  olo <= ilo + 1e-9 @?
    (label ++ ": outer lower " ++ show olo ++ " > inner lower " ++ show ilo)
  ohi >= ihi - 1e-9 @?
    (label ++ ": outer upper " ++ show ohi ++ " < inner upper " ++ show ihi)

-- | Assert that forcing the value throws the empty-credal-set error.
assertNoFeasible :: Show a => String -> a -> Assertion
assertNoFeasible label x = do
  result <- try (evaluate (length (show x)))
  case result of
    Left e  -> "No feasible probabilities" `isPrefixOf` show (e :: ErrorCall) @?
                 (label ++ ": unexpected error: " ++ show e)
    Right _ -> assertFailure (label ++ ": expected an error, got " ++ show x)