<commit_before>{-# LANGUAGE TypeFamilies, Rank2Types, ImpredicativeTypes, DataKinds #-}

module Tests.Lazy where

import Prelude hiding (Real)

import Language.Hakaru.Lazy
import Language.Hakaru.Any (Any(Any, unAny))
import Language.Hakaru.Syntax (Hakaru(..), Base(..),
                               ununit, max_, liftM, liftM2, bind_,
                               swap_, bern, fst_,
                               Mochastic(..), Lambda(..), Integrate(..),
                               Order_(..))
import Language.Hakaru.PrettyPrint (PrettyPrint, runPrettyPrint, leftMode)
import Language.Hakaru.Simplify (Simplifiable, closeLoop, simplify)
import Language.Hakaru.Expect (Expect(Expect), Expect', normalize)
import Language.Hakaru.Maple (Maple)
import Language.Hakaru.Inference
import Tests.TestTools
import qualified Tests.Models as RT
import qualified Examples.EasierRoadmap as RM    

import Data.Typeable (Typeable)    
import Test.HUnit
import qualified Data.List as L

recover :: (Typeable a) => PrettyPrint a -> IO (Any a)
recover hakaru = closeLoop ("Any (" ++ leftMode (runPrettyPrint hakaru) ++ ")")

simp :: (Simplifiable a) => Any a -> IO (Any a)
simp = simplify . unAny

testS' :: (Simplifiable a) => Any a -> Assertion
testS' t = testS (unAny t)

testL
    :: (Backward a a, Typeable a, Typeable b,
        Simplifiable a, Simplifiable b, Simplifiable env)
    => Cond PrettyPrint env (HMeasure (HPair a b))
    -> [(Maple env, Maple a, Any (HMeasure b))]
    -> Assertion
testL f slices = do
  ds <- mapM recover (runDisintegrate f)
  assertResult ds
  mapM_ testS' ds
  mapM_ (\(env,a,t') ->
         testSS [(app (app (unAny d) env) a) | d <- ds] (unAny t')) slices
         
exists :: PrettyPrint a -> [PrettyPrint a] -> Assertion
exists t ts' = assertBool "no correct disintegration" $
               elem (result t) (map result ts')

runDisintegratePretty
    :: (Backward a a)
    => Cond PrettyPrint env (HMeasure (HPair a b))
    -> IO ()
runDisintegratePretty = print . map runPrettyPrint . runDisintegrate

nonDefault
    :: (Backward a a)
    => String
    -> Cond PrettyPrint env (HMeasure (HPair a b))
    -> Assertion
nonDefault s = assertBool "not calling non-default implementation" . any id .
               map (L.isInfixOf s . leftMode . runPrettyPrint) . runDisintegrate

justRun :: Test -> IO ()
justRun t = runTestTT t >> return ()
                   
main :: IO ()
main = -- justRun nonDefaultTests
       runDisintegratePretty prog1s

-- | TODO: understand why the following fail to use non-default uniform
-- prog1s, prog2s, prog3s, (const culpepper), t3, t4, t9, marsaglia
nonDefaultTests :: Test
nonDefaultTests = test [ nonDefault "uniform" borelishSub
                       , nonDefault "uniform" borelishDiv
                       , nonDefault "uniform" density1
                       , nonDefault "uniform" density2
                       , nonDefault "uniform" t1
                       , nonDefault "uniform" t2
                       , nonDefault "uniform" t5
                       , nonDefault "uniform" t6
                       , nonDefault "uniform" t7
                       , nonDefault "uniform" t8
                       , nonDefault "sqrt_"   fwdSqrt_
                       , nonDefault "sqrt"    fwdSqrt
                       , nonDefault "**"      fwdPow
                       , nonDefault "pow_"    fwdPow_
                       , nonDefault "logBase" fwdLogBase
                       ]

-- 2015-04-09
--------------------------------------------------------------------------------
important :: Test
important = test
    [ "zeroAddInt" ~: testL zeroAddInt [(unit, 0, Any $ dirac 3)]
    , "zeroAddReal" ~: testL zeroAddReal [(unit, 0, Any $ dirac 3)]
    , "easierRoadmapProg1" ~: testL easierRoadmapProg1 []
    ]
--------------------------------------------------------------------------------
            
allTests :: Test
allTests = test
    [ "easierRoadmapProg1" ~: testL easierRoadmapProg1 []
    , "normalFB1" ~: testL normalFB1 []
    , "normalFB2" ~: testL normalFB2 []
    , "zeroDiv" ~: testL zeroDiv [(unit, 0, Any $ dirac 0)]
    , "zeroAddInt" ~: testL zeroAddInt [(unit, 0, Any $ dirac 3)]
    , "zeroPlusSnd" ~: testL zeroPlusSnd []
    , "prog1s" ~: testL prog1s []
    , "prog2s" ~: testL prog2s []
    , "prog3s" ~: testL prog3s []
    , "pair1fst" ~: testL pair1fst []
    , "pair1fstSwap" ~: testL pair1fstSwap []
    , "borelishSub" ~: testL borelishSub [(unit, 0, Any (uniform 0 1))]
    , "borelishDiv" ~: testL borelishDiv
            [(unit, 1, Any
                (superpose [(1/2, liftM fromProb (beta 2 1))]))]
    , "culpepper" ~: testL (const culpepper)
            [(unit, 0, Any
                (superpose [(fromRational (1/8), dirac true)
                ,(fromRational (1/8), dirac false)]))]
    , "beta" ~: testL testBetaConj
            [(unit, true, Any
                (superpose [(fromRational (1/2), beta 2 1)]))]
    , "betaNorm" ~: testSS [testBetaConjNorm] (beta 2 1)
    , "testGibbsPropUnif" ~: testS testGibbsPropUnif
    , "testGibbs0" ~: testSS [testGibbsProp0]
            (lam $ \x ->
            normal (x * fromRational (1/2))
                (sqrt_ 2 * fromRational (1/2)))
    , "testGibbs1" ~: testSS [testGibbsProp1]
            (lam $ \x ->
            normal (fst_ x) 1 `bind` \y ->
            dirac (pair (fst_ x) y))
    -- , "testGibbs2" ~: testSS [testGibbsProp2]
    --                    (lam $ \x ->
    --                     normal ((snd_ x) * fromRational (1/2))
    --                            (sqrt_ 2 * fromRational (1/2))
    --                            `bind` \y ->
    --                     dirac (pair y (snd_ x)))
    , "density1" ~: testL density1 []
    , "density2" ~: testL density2 []
    -- , "density3" ~: testL density3 []
    , "t0"  ~: testL t0
            [(unit, 1, Any
              (superpose 
               [( recip (sqrt_ pi_) *
                  exp_ (1 * 1 * fromRational (-1/4)) *
                  fromRational (1/2)
                , normal
                  (1 * fromRational (1/2))
                  ((sqrt_ 2) * fromRational (1/2)) )]))]
    , "t1"  ~: testL t1 []
    , "t2"  ~: testL t2 []
    , "t3"  ~: testL t3 []
    , "t4"  ~: testL t4 []
    , "t5"  ~: testL t5 []
    , "t6"  ~: testL t6 []
    , "t7"  ~: testL t7 []
    , "t8"  ~: testL t8 []
    , "t9"  ~: testL t9 []
    , "t10" ~: testL t10 []
    , "marsaglia" ~: testL marsaglia [] -- needs heavy disintegration
    ]

normalFB1
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HReal HUnit))
normalFB1 = \u -> ununit u $
            normal 0 1 `bind` \x ->
            normal x 1 `bind` \y ->
            dirac (pair ((y + y) + x) unit)

normalFB2
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HReal HUnit))
normalFB2 = \u -> ununit u $
            normal 0 1 `bind` \x ->
            normal x 1 `bind` \y ->
            dirac (pair (y + x) unit)

easierRoadmapProg1
    :: (Mochastic repr)
    => Cond repr HUnit
        (HMeasure (HPair (HPair HReal HReal) (HPair HProb HProb)))
easierRoadmapProg1 = \u -> ununit u $ RM.easierRoadmapProg1

zeroDiv
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HReal HReal))
zeroDiv = \u -> ununit u $
          normal 0 1 `bind` \x ->
          dirac (pair x (0 / x))

zeroAddInt
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HInt HInt))
zeroAddInt = \u -> ununit u $<commit_msg>Fixing merge conflicts and commenting out some tests for now<commit_after>{-# LANGUAGE TypeFamilies, Rank2Types, ImpredicativeTypes, DataKinds #-}

module Tests.Lazy where

import Prelude hiding (Real)

import Language.Hakaru.Lazy
import Language.Hakaru.Any (Any(Any, unAny))
import Language.Hakaru.Syntax (Hakaru(..), Base(..),
                               ununit, max_, liftM, liftM2, bind_,
                               swap_, bern, fst_, not_, equal_, weight,
                               Mochastic(..), Lambda(..), Integrate(..),
                               Order_(..))
import Language.Hakaru.PrettyPrint (PrettyPrint, runPrettyPrint, leftMode)
import Language.Hakaru.Simplify (Simplifiable, closeLoop, simplify)
import Language.Hakaru.Expect (Expect(Expect), Expect', normalize, total)
import Language.Hakaru.Maple (Maple)
import Language.Hakaru.Inference
import Tests.TestTools
import qualified Tests.Models as RT
import qualified Examples.EasierRoadmap as RM
import Examples.BenchmarkProblems (runExpect)

import Data.Typeable (Typeable)    
import Test.HUnit
import qualified Data.List as L

import Text.PrettyPrint    

recover :: (Typeable a) => PrettyPrint a -> IO (Any a)
recover hakaru = closeLoop ("Any (" ++ leftMode (runPrettyPrint hakaru) ++ ")")

simp :: (Simplifiable a) => Any a -> IO (Any a)
simp = simplify . unAny

testS' :: (Simplifiable a) => Any a -> Assertion
testS' t = testS (unAny t)

testL
    :: (Backward a a, Typeable a, Typeable b,
        Simplifiable a, Simplifiable b, Simplifiable env)
    => Cond PrettyPrint env (HMeasure (HPair a b))
    -> [(Maple env, Maple a, Any (HMeasure b))]
    -> Assertion
testL f slices = do
  ds <- mapM recover (runDisintegrate f)
  assertResult ds
  mapM_ testS' ds
  mapM_ (\(env,a,t') ->
         testSS [(app (app (unAny d) env) a) | d <- ds] (unAny t')) slices
         
exists :: PrettyPrint a -> [PrettyPrint a] -> Assertion
exists t ts' = assertBool "no correct disintegration" $
               elem (result t) (map result ts')

runDisintegratePretty
    :: (Backward a a)
    => Cond PrettyPrint env (HMeasure (HPair a b))
    -> IO ()
runDisintegratePretty = print . map runPrettyPrint . runDisintegrate

nonDefault
    :: (Backward a a)
    => String
    -> Cond PrettyPrint env (HMeasure (HPair a b))
    -> Assertion
nonDefault s = assertBool "not calling non-default implementation" . any id .
               map (L.isInfixOf s . leftMode . runPrettyPrint) . runDisintegrate

justRun :: Test -> IO ()
justRun t = runTestTT t >> return ()
                   
main :: IO ()
main = -- justRun nonDefaultTests
       runDisintegratePretty prog1s

-- | TODO: understand why the following fail to use non-default uniform
-- prog1s, prog2s, prog3s, (const culpepper), t3, t4, t9, marsaglia
nonDefaultTests :: Test
nonDefaultTests = test [ nonDefault "uniform" borelishSub
                       , nonDefault "uniform" borelishDiv
                       , nonDefault "uniform" density1
                       , nonDefault "uniform" density2
                       , nonDefault "uniform" t1
                       , nonDefault "uniform" t2
                       , nonDefault "uniform" t5
                       , nonDefault "uniform" t6
                       , nonDefault "uniform" t7
                       , nonDefault "uniform" t8
                       , nonDefault "sqrt_"   fwdSqrt_
                       , nonDefault "sqrt"    fwdSqrt
                       , nonDefault "**"      fwdPow
                       , nonDefault "pow_"    fwdPow_
                       , nonDefault "logBase" fwdLogBase
                       ]

-- 2015-04-09
--------------------------------------------------------------------------------
important :: Test
important = test
    [ "zeroAddInt" ~: testL zeroAddInt [(unit, 0, Any $ dirac 3)]
    , "zeroAddReal" ~: testL zeroAddReal [(unit, 0, Any $ dirac 3)]
    , "easierRoadmapProg1" ~: testL easierRoadmapProg1 []
    ]
--------------------------------------------------------------------------------
            
allTests :: Test
allTests = test
    [ "easierRoadmapProg1" ~: testL easierRoadmapProg1 []
    , "normalFB1" ~: testL normalFB1 []
    , "normalFB2" ~: testL normalFB2 []
    , "zeroDiv" ~: testL zeroDiv [(unit, 0, Any $ dirac 0)]
    , "zeroAddInt" ~: testL zeroAddInt [(unit, 0, Any $ dirac 3)]
    , "zeroPlusSnd" ~: testL zeroPlusSnd []
    , "prog1s" ~: testL prog1s []
    , "prog2s" ~: testL prog2s []
    , "prog3s" ~: testL prog3s []
    , "pair1fst" ~: testL pair1fst []
    , "pair1fstSwap" ~: testL pair1fstSwap []
    , "borelishSub" ~: testL borelishSub [(unit, 0, Any (uniform 0 1))]
    , "borelishDiv" ~: testL borelishDiv
            [(unit, 1, Any
                (superpose [(1/2, liftM fromProb (beta 2 1))]))]
    , "culpepper" ~: testL (const culpepper)
            [(unit, 0, Any
                (superpose [(fromRational (1/8), dirac true)
                ,(fromRational (1/8), dirac false)]))]
    , "beta" ~: testL testBetaConj
            [(unit, true, Any
                (superpose [(fromRational (1/2), beta 2 1)]))]
    , "betaNorm" ~: testSS [testBetaConjNorm] (beta 2 1)
    , "testGibbsPropUnif" ~: testS testGibbsPropUnif
    , "testGibbs0" ~: testSS [testGibbsProp0]
            (lam $ \x ->
            normal (x * fromRational (1/2))
                (sqrt_ 2 * fromRational (1/2)))
    , "testGibbs1" ~: testSS [testGibbsProp1]
            (lam $ \x ->
            normal (fst_ x) 1 `bind` \y ->
            dirac (pair (fst_ x) y))
    -- , "testGibbs2" ~: testSS [testGibbsProp2]
    --                    (lam $ \x ->
    --                     normal ((snd_ x) * fromRational (1/2))
    --                            (sqrt_ 2 * fromRational (1/2))
    --                            `bind` \y ->
    --                     dirac (pair y (snd_ x)))
    , "density1" ~: testL density1 []
    , "density2" ~: testL density2 []
    -- , "density3" ~: testL density3 []
    , "t0"  ~: testL t0
            [(unit, 1, Any
              (superpose 
               [( recip (sqrt_ pi_) *
                  exp_ (1 * 1 * fromRational (-1/4)) *
                  fromRational (1/2)
                , normal
                  (1 * fromRational (1/2))
                  ((sqrt_ 2) * fromRational (1/2)) )]))]
    , "t1"  ~: testL t1 []
    , "t2"  ~: testL t2 []
    , "t3"  ~: testL t3 []
    , "t4"  ~: testL t4 []
    , "t5"  ~: testL t5 []
    , "t6"  ~: testL t6 []
    , "t7"  ~: testL t7 []
    , "t8"  ~: testL t8 []
    , "t9"  ~: testL t9 []
    , "t10" ~: testL t10 []
    , "marsaglia" ~: testL marsaglia [] -- needs heavy disintegration
    ]

normalFB1
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HReal HUnit))
normalFB1 = \u -> ununit u $
            normal 0 1 `bind` \x ->
            normal x 1 `bind` \y ->
            dirac (pair ((y + y) + x) unit)

normalFB2
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HReal HUnit))
normalFB2 = \u -> ununit u $
            normal 0 1 `bind` \x ->
            normal x 1 `bind` \y ->
            dirac (pair (y + x) unit)

easierRoadmapProg1
    :: (Mochastic repr)
    => Cond repr HUnit
        (HMeasure (HPair (HPair HReal HReal) (HPair HProb HProb)))
easierRoadmapProg1 = \u -> ununit u $ RM.easierRoadmapProg1

zeroDiv
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HReal HReal))
zeroDiv = \u -> ununit u $
          normal 0 1 `bind` \x ->
          dirac (pair x (0 / x))
            
-- cushing :: (Mochastic repr) => Cond repr () (Measure (Real,Real))
-- cushing = \u -> ununit u $
--           uniform 0 1 `bind` \x ->
--           if_ (not_ (equal_ x (1/2)))
--               (uniform 0 x)
--               (uniform x 1) `bind` \y ->
--           dirac (pair y x)                

-- cushingObs n = simplify (d `app` unit `app` n)
--     where d:_ = runDisintegrate cushing

-- cushingSimpl = mapM simplify d
--     where d = runDisintegrate cushing

-- cobb :: (Mochastic repr) => Cond repr () (Measure (Real,Real))
-- cobb = \u -> ununit u $
--        uniform 0 1 `bind` \x ->
--        uniform (-x) x `bind` \y ->
--        dirac (pair y x)

-- cobbSimpl = mapM simplify d
--     where d = runDisintegrate cobb
             
-- cobbObs n = simplify (d `app` unit `app` n)
--     where d:_ = runDisintegrate cobb

-- cobb0 :: (Mochastic repr, Lambda repr, Integrate repr)
--          => Expect repr (Measure Prob)
-- cobb0 = uniform 0 1 `bind` \x0 ->
--         weight (recip (unsafeProb x0) * (1/2)) $
--         dirac (unsafeProb x0)

-- expectCobb0 :: Doc
-- expectCobb0 = runPrettyPrint (runExpect cobb0)

-- totalCobb0 :: Doc
-- totalCobb0 = runPrettyPrint (total cobb0)

-- thenSimpl1 = simplify $ dirac (total cobb0)
-- thenSimpl2 = simplify $ dirac (runExpect cobb0)
-- thenSimpl3 = simplify $ normalize cobb0
                          
zeroAddInt
    :: (Mochastic repr)
    => Cond repr HUnit (HMeasure (HPair HInt HInt))
zeroAddInt = \u -> ununit u $
{'fraction_non_alphanumeric': 0.14862474842959078, 'fraction_numerical': 0.01957675184484967, 'mean_word_length': 1.8277289187791, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 93, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7005347', 'n_tokens_mistral': 5697, 'n_tokens_neox': 5299, 'n_words': 1533}