(module
  (aux hakaru10 random)
  *
  (import
    scheme
    (only (scheme base) exact-integer?)
    (chicken base)
    (chicken fixnum)
    (chicken flonum)
    (chicken bitwise)
    (only (chicken format) fprintf)
    (only (chicken condition)
          make-property-condition
          make-composite-condition
          signal
          condition-predicate
          condition-property-accessor)
    (only (chicken bytevector)
          bytevector?
          bytevector-length
          make-bytevector
          bytevector-copy!)
    (only (chicken random)
          set-pseudo-random-seed!
          pseudo-random-integer
          pseudo-random-real)
    (only srfi-1 every split-at append-map)
    srfi-4
    (aux base))
  (declare (local))
  (define (hakaru10-error kind loc msg . args)
    (signal
      (make-composite-condition
        (make-property-condition
          'exn
          'location
          loc
          'message
          msg
          'arguments
          args)
        (make-property-condition 'hakaru10 'kind kind))))
  (define hakaru10-error? (condition-predicate 'hakaru10))
  (define hakaru10-error-kind (condition-property-accessor 'hakaru10 'kind))
  (define (check-integer-bounds who lo hi)
    (unless
      (and (exact-integer? lo) (exact-integer? hi))
      (hakaru10-error 'parameter who "bounds must be exact integers" lo hi)))
  (define-syntax-rule
    (define-binary-syntax (name op) ...)
    (begin (define-syntax-rule (name a b) (op a b)) ...))
  (cond-expand
    (|64bit| (define-binary-syntax
               (stdgen+ fx+)
               (stdgen- fx-)
               (stdgen* fx*)
               (stdgen-quot fx/)
               (stdgen-mod fxmod)
               (stdgen< fx<)))
    (else (define-binary-syntax
            (stdgen+ +)
            (stdgen- -)
            (stdgen* *)
            (stdgen-quot quotient)
            (stdgen-mod modulo)
            (stdgen< <))))
  (define-record stdgen s1 s2)
  (set-record-printer!
    stdgen
    (λ (g port) (fprintf port "#<stdgen ~a ~a>" (stdgen-s1 g) (stdgen-s2 g))))
  (define (stdgen-copy g) (make-stdgen (stdgen-s1 g) (stdgen-s2 g)))
  (define-syntax-rule
    (define-stdgen/pure (name arg ...) name!)
    (define (name arg ... g)
      (let1 (g* (stdgen-copy g)) (let1 (x (name! arg ... g*)) (values x g*)))))
  (define stdgen-b 2147483561)
  (define (seed->stdgen n)
    (unless
      (exact-integer? n)
      (hakaru10-error
        'parameter
        'seed->stdgen
        "seed must be an exact integer"
        n))
    (let* ((s (bitwise-and n 2147483647))
           (q (quotient s 2147483562))
           (s1 (remainder s 2147483562))
           (s2 (remainder q 2147483398)))
      (make-stdgen (+ s1 1) (+ s2 1))))
  (define-syntax-rule
    (let-stdnext ((z next1 next2) (seed1 seed2)) body ...)
    (let* ((s1 seed1)
           (s2 seed2)
           (k (stdgen-quot s1 53668))
           (s1* (stdgen-
                  (stdgen* 40014 (stdgen- s1 (stdgen* k 53668)))
                  (stdgen* k 12211)))
           (next1 (if (stdgen< s1* 0) (stdgen+ s1* 2147483563) s1*))
           (k* (stdgen-quot s2 52774))
           (s2* (stdgen-
                  (stdgen* 40692 (stdgen- s2 (stdgen* k* 52774)))
                  (stdgen* k* 3791)))
           (next2 (if (stdgen< s2* 0) (stdgen+ s2* 2147483399) s2*))
           (z* (stdgen- next1 next2))
           (z (if (stdgen< z* 1) (stdgen+ z* 2147483562) z*)))
      body
      ...))
  (define-syntax-rule
    (stdgen-fold! g n (z acc) step init)
    (let1 (gen g)
          (let loop ((i n)
                     (acc init)
                     (s1 (stdgen-s1 gen))
                     (s2 (stdgen-s2 gen)))
            (if (eq? i 0)
              (begin (stdgen-s1-set! gen s1) (stdgen-s2-set! gen s2) acc)
              (let-stdnext
                ((z next1 next2) (s1 s2))
                (loop (sub1 i) step next1 next2))))))
  (define-syntax-rule (stdnext! g) (stdgen-fold! g 1 (z _) z #f))
  (define (stdgen-next! g) (stdnext! g))
  (define-stdgen/pure (stdgen-next) stdgen-next!)
  (define (ilog-base b i) (if (< i b) 1 (add1 (ilog-base b (quotient i b)))))
  (define (stdgen-ival-integer! l h g)
    (check-integer-bounds 'stdgen-ival-integer l h)
    (if (> l h)
      (stdgen-ival-integer! h l g)
      (let1 (k (add1 (- h l)))
            (+ l
               (if (< k stdgen-b)
                 (stdgen-mod (stdgen+ (stdnext! g) stdgen-b) k)
                 (modulo
                   (stdgen-fold!
                     g
                     (ilog-base stdgen-b k)
                     (z acc)
                     (+ z (* acc stdgen-b))
                     1)
                   k))))))
  (define-stdgen/pure (stdgen-ival-integer l h) stdgen-ival-integer!)
  (cond-expand
    (|64bit| (define-inline
               (stdgen-acc53 x acc)
               (let* ((hi (fxshr acc 27))
                      (lo (fxand acc 134217727))
                      (hib (fxshl (fxand (fx* hi stdgen-b) 67108863) 27)))
                 (fxand (fx+ x (fx+ hib (fx* lo stdgen-b)))
                        9007199254740991))))
    (else (define-inline
            (stdgen-acc53 x acc)
            (bitwise-and (+ x (* acc stdgen-b)) 9007199254740991))))
  (define (stdgen-double! g)
    (fp/ (exact->inexact (stdgen-fold! g 3 (z acc) (stdgen-acc53 z acc) 1))
         9007199254740992.0))
  (define-stdgen/pure (stdgen-double) stdgen-double!)
  (define (stdgen-double/reference g)
    (let*-values
      (((twoto53) (expt 2 53))
       ((x g*) (stdgen-ival-integer (- (expt 2 63)) (sub1 (expt 2 63)) g)))
      (values
        (/ (exact->inexact (bitwise-and x (sub1 twoto53)))
           (exact->inexact twoto53))
        g*)))
  (define (stdgen-range-double! lo hi g)
    (if (fp> lo hi)
      (stdgen-range-double! hi lo g)
      (let1 (coef (stdgen-double! g))
            (fp* 2.0
                 (fp+ (fp* 0.5 lo)
                      (fp* coef (fp- (fp* 0.5 hi) (fp* 0.5 lo))))))))
  (define-stdgen/pure (stdgen-range-double lo hi) stdgen-range-double!)
  (define-record rng-source name state draw-double draw-integer copier)
  (set-record-printer!
    rng-source
    (λ (src port)
       (fprintf
         port
         "#<rng-source ~a ~s>"
         (rng-source-name src)
         (rng-source-state src))))
  (define (rng-double! src lo hi)
    ((rng-source-draw-double src)
     (rng-source-state src)
     (exact->inexact lo)
     (exact->inexact hi)))
  (define (rng-integer! src lo hi)
    ((rng-source-draw-integer src) (rng-source-state src) lo hi))
  (define (rng-source-copy src)
    (make-rng-source
      (rng-source-name src)
      ((rng-source-copier src) (rng-source-state src))
      (rng-source-draw-double src)
      (rng-source-draw-integer src)
      (rng-source-copier src)))
  (define (make-stdgen-source #!optional (seed 17))
    (make-rng-source
      'stdgen
      (if (stdgen? seed) (stdgen-copy seed) (seed->stdgen seed))
      (λ (g lo hi) (stdgen-range-double! lo hi g))
      (λ (g lo hi) (stdgen-ival-integer! lo hi g))
      stdgen-copy))
  (define (unit->range u lo hi)
    (if (fp> lo hi) (unit->range u hi lo) (fp+ lo (fp* u (fp- hi lo)))))
  (define (unit->integer u lo hi)
    (if (> lo hi)
      (unit->integer u hi lo)
      (min hi
           (+ lo
              (inexact->exact
                (fpfloor (fp* u (exact->inexact (add1 (- hi lo))))))))))
  (define mrg32k3a-m1 4294967087.0)
  (define mrg32k3a-m2 4294944443.0)
  (define mrg32k3a-norm 2.32830654929573e-10)
  (define-inline
    (mrg32k3a-component a x b y m)
    (let* ((p (fp- (fp* a x) (fp* b y)))
           (p (fp- p (fp* (fptruncate (fp/ p m)) m))))
      (if (fp< p 0.0) (fp+ p m) p)))
  (define (mrg32k3a-next! s)
    (let ((p1 (mrg32k3a-component
                1403580.0
                (f64vector-ref s 1)
                810728.0
                (f64vector-ref s 0)
                mrg32k3a-m1))
          (p2 (mrg32k3a-component
                527612.0
                (f64vector-ref s 5)
                1370589.0
                (f64vector-ref s 3)
                mrg32k3a-m2)))
      (f64vector-set! s 0 (f64vector-ref s 1))
      (f64vector-set! s 1 (f64vector-ref s 2))
      (f64vector-set! s 2 p1)
      (f64vector-set! s 3 (f64vector-ref s 4))
      (f64vector-set! s 4 (f64vector-ref s 5))
      (f64vector-set! s 5 p2)
      (fp* (if (fp<= p1 p2) (fp+ (fp- p1 p2) mrg32k3a-m1) (fp- p1 p2))
           mrg32k3a-norm)))
  (define (mrg32k3a-seed->state seed)
    (define moduli
      '(4294967087 4294967087 4294967087 4294944443 4294944443 4294944443))
    (define (all-zero? triple) (every zero? triple))
    (define (component? x m) (and (real? x) (integer? x) (<= 0 x) (< x m)))
    (define (spread seed)
      (let loop ((moduli moduli) (x (modulo seed 4294967296)))
        (if (null? moduli)
          '()
          (let1 (x* (modulo (add1 (* 69069 x)) 4294967296))
                (cons (modulo x* (car moduli)) (loop (cdr moduli) x*))))))
    (define (->state words) (list->f64vector (map exact->inexact words)))
    (define seed-words
      (cond ((list? seed) seed)
            ((vector? seed) (vector->list seed))
            ((f64vector? seed) (f64vector->list seed))
            (else #f)))
    (cond ((exact-integer? seed)
           (let-values
             (((t1 t2) (split-at (spread seed) 3)))
             (->state
               (append-map
                 (λ (t) (if (all-zero? t) '(12345 12345 12345) t))
                 (list t1 t2)))))
          ((and seed-words
                (= (length seed-words) 6)
                (every component? seed-words moduli))
           (let-values
             (((t1 t2) (split-at seed-words 3)))
             (when (or (all-zero? t1) (all-zero? t2))
                   (hakaru10-error
                     'parameter
                     'make-mrg32k3a-source
                     "a seed triple is all zero"
                     seed))
             (->state seed-words)))
          (else
           (hakaru10-error
             'parameter
             'make-mrg32k3a-source
             "seed must be an exact integer or six integers in [0,m1)^3 x [0,m2)^3"
             seed))))
  (define (make-mrg32k3a-source #!optional (seed 17))
    (make-rng-source
      'mrg32k3a
      (mrg32k3a-seed->state seed)
      (λ (s lo hi) (unit->range (mrg32k3a-next! s) lo hi))
      (λ (s lo hi)
         (check-integer-bounds 'rng-integer! lo hi)
         (unit->integer (mrg32k3a-next! s) lo hi))
      (λ (s) (list->f64vector (f64vector->list s)))))
  (define (chicken-random-seed seed)
    (unless
      (bytevector? seed)
      (hakaru10-error
        'parameter
        'make-chicken-random-source
        "seed must be a bytevector"
        seed))
    (when (every zero? (u8vector->list seed))
          (hakaru10-error
            'parameter
            'make-chicken-random-source
            "an all-zero seed would hang (chicken random)"
            seed))
    (let* ((n (bytevector-length seed))
           (padded (make-bytevector (* 8 (max 1 (quotient (+ n 7) 8))) 0)))
      (bytevector-copy! padded 0 seed)
      padded))
  (define (chicken-random-unit)
    (let1 (u (pseudo-random-real))
          (if (and (fp< 0.0 u) (fp< u 1.0)) u (chicken-random-unit))))
  (define (make-chicken-random-source #!optional seed)
    (when seed (set-pseudo-random-seed! (chicken-random-seed seed)))
    (make-rng-source
      'chicken-random
      'global
      (λ (_ lo hi) (unit->range (chicken-random-unit) lo hi))
      (λ (_ lo hi)
         (check-integer-bounds 'rng-integer! lo hi)
         (let1 (l (min lo hi))
               (+ l (pseudo-random-integer (add1 (- (max lo hi) l))))))
      (λ (_)
         (hakaru10-error
           'parameter
           'rng-source-copy
           "the global state of (chicken random) cannot be copied")))))

mkStdGen :: Int -> StdGen -- why not Integer ?
mkStdGen s = mkStdGen32 $ fromIntegral s

mkStdGen32 :: Int32 -> StdGen
mkStdGen32 sMaybeNegative = StdGen (s1+1) (s2+1)
      where
	-- We want a non-negative number, but we can't just take the abs
	-- of sMaybeNegative as -minBound == minBound.
	s       = sMaybeNegative .&. maxBound
	(q, s1) = s `divMod` 2147483562
	s2      = q `mod` 2147483398

(define (test/seed->stdgen _)
  (let1 (g (seed->stdgen 17))
        (⊦ equal? '(18 1) (list (stdgen-s1 g) (stdgen-s2 g))))
  (⊦ equal?
     '((1 1) (2 1) (86 2) (82 2) (18 1))
     (map (λ (n)
             (let1 (g (seed->stdgen n)) (list (stdgen-s1 g) (stdgen-s2 g))))
          (list 0 1 2147483647 -5 (+ (expt 2 40) 17))))
  (⊦ eq? 'parameter (hakaru10-kind (τ (seed->stdgen 1.5)))))

((eta 0.0) (memory #(12582912 2241544 1048576)) (stdout "") (stderr ""))

stdNext :: StdGen -> (Int, StdGen)
-- Returns values in the range stdRange
stdNext (StdGen s1 s2) = (fromIntegral z', StdGen s1'' s2'')
	where	z'   = if z < 1 then z + 2147483562 else z
		z    = s1'' - s2''

		k    = s1 `quot` 53668
		s1'  = 40014 * (s1 - k * 53668) - k * 12211
		s1'' = if s1' < 0 then s1' + 2147483563 else s1'

		k'   = s2 `quot` 52774
		s2'  = 40692 * (s2 - k' * 52774) - k' * 3791
		s2'' = if s2' < 0 then s2' + 2147483399 else s2'

(define (test/stdNext _)
  (let*-values
    (((a g) (stdgen-next (seed->stdgen 17)))
     ((b g) (stdgen-next g))
     ((c g) (stdgen-next g)))
    (⊦ equal? '(679560 1394521907 656733876) (list a b c))
    (⊦ equal? '(612660577 2103410263) (list (stdgen-s1 g) (stdgen-s2 g))))
  (let1 (g (seed->stdgen 17))
        (stdgen-next g)
        (⊦ equal? '(18 1) (list (stdgen-s1 g) (stdgen-s2 g)))
        (⊦ equal?
           '(679560 1394521907 656733876)
           (let* ((a (stdgen-next! g))
                  (b (stdgen-next! g))
                  (c (stdgen-next! g)))
             (list a b c)))
        (⊦ equal? '(612660577 2103410263) (list (stdgen-s1 g) (stdgen-s2 g)))))

((eta 0.0) (memory #(12582912 2240736 1048576)) (stdout "") (stderr ""))

randomRFloating (l,h) g
    | l>h       = randomRFloating (h,l) g
    | otherwise = let (coef,g') = random g in
		  (2.0 * (0.5*l + coef * (0.5*h - 0.5*l)), g')  -- avoid overflow

instance Random Double where
  randomR = randomRFloating
  random rng     =
    case random rng of
      (x,rng') ->
          -- We use 53 bits of randomness corresponding to the 53 bit significand:
          ((fromIntegral (mask53 .&. (x::Int64)) :: Double)
	   /  fromIntegral twoto53, rng')
   where
    twoto53 = (2::Int64) ^ (53::Int64)
    mask53 = twoto53 - 1

(define (test/randomR-double _)
  (⊦ = 0.293899585024381 (rng-double! (make-stdgen-source 17) 0.0 1.0))
  (⊦ = 0.293899585024381 (rng-double! (make-stdgen-source) 0 1))
  (⊦ = -1.81007738610981 (rng-double! (make-stdgen-source 17) -3.5 2.25))
  (let*-values
    (((d g) (stdgen-double (seed->stdgen 17))))
    (⊦ = 0.293899585024381 d)
    (⊦ equal? '(612660577 2103410263) (list (stdgen-s1 g) (stdgen-s2 g)))))

((eta 0.0) (memory #(12582912 2242736 1048576)) (stdout "") (stderr ""))

randomIvalInteger :: (RandomGen g, Num a) => (Integer, Integer) -> g -> (a, g)
randomIvalInteger (l,h) rng
 | l > h     = randomIvalInteger (h,l) rng
 | otherwise = case (f n 1 rng) of (v, rng') -> (fromInteger (l + v `mod` k), rng')
     where
       k = h - l + 1
       -- ERROR: b here (2^31-87) represents a baked-in assumption about genRange:
       b = 2147483561
       n = iLogBase b k

       -- Here we loop until we've generated enough randomness to cover the range:
       f 0 acc g = (acc, g)
       f n' acc g =
          let
	   (x,g')   = next g
	  in
          -- We shift over the random bits generated thusfar (* b) and add in the new ones.
	  f (n' - 1) (fromIntegral x + acc * b) g'

(define (test/randomR-int _)
  (⊦ = 1 (rng-integer! (make-stdgen-source 17) 0 9))
  (⊦ = 1 (rng-integer! (make-stdgen-source 17) 9 0))
  (let*-values
    (((v g)
      (stdgen-ival-integer (- (expt 10 30)) (expt 10 30) (seed->stdgen 17))))
    (⊦ = -471988287518935644257797793060 v)
    (⊦ equal? '(1475456433 1872071452) (list (stdgen-s1 g) (stdgen-s2 g))))
  (let1 (src (make-stdgen-source 17))
        (⊦ = 5 (rng-integer! src 5 5))
        (let1 (g (rng-source-state src))
              (⊦ equal? '(720252 40692) (list (stdgen-s1 g) (stdgen-s2 g)))))
  (⊦ eq?
     'parameter
     (hakaru10-kind (τ (rng-integer! (make-stdgen-source) 0.0 9)))))

((eta 0.0) (memory #(12582912 2247624 1048576)) (stdout "") (stderr ""))

(define (test/fixnum-double-vs-reference _)
  (⊦ equal?
     '(0 0 0 0 0)
     (map (λ (seed) (stdgen-double-mismatches seed 100000))
          '(17 0 1 2147483647 -5))))

((eta 2.393) (memory #(12582912 2242632 1048576)) (stdout "") (stderr ""))

(define (test/copy-replays _)
  (let* ((src (make-stdgen-source 17))
         (_ (draws src 10))
         (copy (rng-source-copy src))
         (a (draws src 500))
         (b (draws copy 500)))
    (⊦ equal? a b)
    (⊦ equal?
       (draws (make-stdgen-source 17) 10)
       (draws (make-stdgen-source (seed->stdgen 17)) 10))))

((eta 0.001) (memory #(12582912 2309016 1048576)) (stdout "") (stderr ""))

(define (test/bounds _)
  (let1 (src (make-stdgen-source 17))
        (⊨ (all-draws?
             100000
             (τ (rng-double! src -2.5 4.0))
             (λ (x) (and (<= -2.5 x) (<= x 4.0)))))
        (⊨ (all-draws?
             1000
             (τ (rng-double! src 4.0 -2.5))
             (λ (x) (and (<= -2.5 x) (<= x 4.0)))))
        (⊨ (all-draws?
             1000
             (τ (rng-integer! src 9 -3))
             (λ (k) (and (<= -3 k) (<= k 9)))))
        (⊦ = 1.0 (rng-double! src 1.0 1.0)))
  (let*-values
    (((a _) (stdgen-range-double -2.5 4.0 (seed->stdgen 17)))
     ((b _) (stdgen-range-double 4.0 -2.5 (seed->stdgen 17))))
    (⊦ = a b)))

((eta 0.06) (memory #(12582912 2245328 1048576)) (stdout "") (stderr ""))

-- Perform an expensive logarithm on arbitrary-size integers by repeated division.
--
-- (NOTE: This actually returns ceiling(log(i) base b) except with an
--  incorrect result at iLogBase b b = 2.)
iLogBase :: Integer -> Integer -> Integer
iLogBase b i = if i < b then 1 else 1 + iLogBase b (i `div` b)

(define (test/ilog-base _)
  (⊦ equal?
     '(1 2 3 3)
     (map (λ (k) (ilog-base 2147483561 k))
          (list 10 2147483561 (expt 2147483561 2) (expt 2 64))))
  (let*-values
    (((x g) (stdgen-ival-integer 0 2147483560 (seed->stdgen 17)))
     ((z1 g1) (stdgen-next (seed->stdgen 17)))
     ((z2 g2) (stdgen-next g1)))
    (⊦ = (modulo (+ z2 (* (+ z1 (* 1 2147483561)) 2147483561)) 2147483561) x)
    (⊦ equal?
       (list (stdgen-s1 g2) (stdgen-s2 g2))
       (list (stdgen-s1 g) (stdgen-s2 g)))))

((eta 0.0) (memory #(12582912 2247576 1048576)) (stdout "") (stderr ""))

(define (test/mrg32k3a _)
  (let1 (src (make-mrg32k3a-source '(12345 12345 12345 12345 12345 12345)))
        (⊦ equal?
           '(0.127011122046577 0.318527565396794 0.30918601558327)
           (let* ((a (rng-double! src 0.0 1.0))
                  (b (rng-double! src 0.0 1.0))
                  (c (rng-double! src 0.0 1.0)))
             (list a b c)))
        (⊨ (all-draws?
             997
             (τ (rng-double! src 0.0 1.0))
             (λ (u) (and (< 0.0 u) (< u 1.0)))))
        (⊦ equal?
           (f64vector
             4239718941.0
             899640195.0
             1411745448.0
             2768972929.0
             343921931.0
             1471537888.0)
           (rng-source-state src)))
  (let* ((src (make-mrg32k3a-source 17)) (copy (rng-source-copy src)))
    (⊦ equal? (draws src 100) (draws copy 100))
    (⊨ (all-draws?
         10000
         (τ (rng-integer! src 9 0))
         (λ (k) (and (<= 0 k) (<= k 9)))))
    (⊨ (all-draws?
         10000
         (τ (rng-double! src 3.0 -1.0))
         (λ (x) (and (<= -1.0 x) (<= x 3.0))))))
  (⊦ equal?
     '(parameter parameter parameter)
     (map (λ (seed) (hakaru10-kind (τ (make-mrg32k3a-source seed))))
          (list '(0 0 0 1 1 1) '(1 1 4294967087 1 1 1) 'seventeen))))

((eta 0.017) (memory #(12582912 2250112 1048576)) (stdout "") (stderr ""))

(define (test/chicken-random _)
  (⊦ equal?
     '(parameter parameter parameter)
     (map (λ (seed) (hakaru10-kind (τ (make-chicken-random-source seed))))
          (list (make-bytevector 16 0) (bytevector 0) "17")))
  (let* ((a (draws (make-chicken-random-source (bytevector 17)) 5))
         (b (draws (make-chicken-random-source (bytevector 17 0 0 0 0 0 0 0))
                   5))
         (src (make-chicken-random-source (bytevector 1 2 3 4 5 6 7 8 9))))
    (⊦ equal? a b)
    (⊨ (all-draws?
         1000
         (τ (rng-double! src 2.0 0.0))
         (λ (x) (and (< 0.0 x) (< x 2.0)))))
    (⊨ (all-draws?
         1000
         (τ (rng-integer! src (expt 2 70) (+ (expt 2 70) 3)))
         (λ (k) (and (<= (expt 2 70) k) (<= k (+ (expt 2 70) 3))))))
    (⊦ eq? 'parameter (hakaru10-kind (τ (rng-source-copy src))))))

((eta 0.002) (memory #(12582912 2251976 1048576)) (stdout "") (stderr ""))

class Distribution d a | d -> a where
  logDensity :: d -> a -> LogLikelihood
  sample     :: RandomGen g => d -> g -> (a, LogLikelihood, g)

-- re-sampleable distribution. This is the place to specify
-- custom proposals
class Distribution d a => RDistribution d a | d -> a where
  resample :: RandomGen g =>
     d ->
     a -> LogLikelihood -> -- old value and its LL
     g ->
     ((a,              -- new sample
       LogLikelihood,  -- its likelihood
       g),             -- new seed
      LogLikelihood)   -- (fwd - bwd) for the proposal step.
  -- By default, resample is really sample again
  resample d _ llold g =
    let s@(_,ll,_) = sample d g in
    (s,ll-llold)

-- Make sure we report the real LL rather than up to the constant
-- (which may depend on the argument to |categorical|).
-- When logDensity is called during the scoring of a proposal,
-- the parameters of the distribution have likely changed

(module
  (aux hakaru10 distribution)
  *
  (import
    scheme
    (only (scheme base) exact-integer?)
    (chicken base)
    (chicken flonum)
    (chicken foreign)
    (only (chicken format) fprintf)
    (only (chicken module) reexport)
    (only srfi-1 every fold alist-copy make-list)
    (aux base)
    (aux hakaru10 random))
  (declare (local))
  (reexport
    (only (aux hakaru10 random)
          hakaru10-error
          hakaru10-error?
          hakaru10-error-kind))
  (define (->flonum who x)
    (cond ((flonum? x) x)
          ((real? x) (exact->inexact x))
          (else
           (hakaru10-error 'parameter who "a real number is expected" x))))
  (define-syntax-rule
    (let-flonums who (var ...) body ...)
    (let ((var (->flonum who var)) ...) body ...))
  (define-record
    rdistr
    name
    params
    log-density
    sample
    resample
    support
    independence?)
  (set-record-printer!
    rdistr
    (λ (d port)
       (fprintf port "#<rdistr ~a ~s>" (rdistr-name d) (rdistr-params d))))
  (define-record distk kind rd value)
  (set-record-printer!
    distk
    (λ (dk port)
       (let1 (rd (distk-rd dk))
             (fprintf port "#<distk ~a" (distk-kind dk))
             (when rd
                   (fprintf port " ~a ~s" (rdistr-name rd) (rdistr-params rd)))
             (unless
               (eq? (distk-kind dk) 'resampleable)
               (fprintf port " ~s" (distk-value dk)))
             (display ">" port))))
  (define ((independence-resample sample) src x llold)
    (receive (x* ll) (sample src) (values x* ll (fp- ll llold))))
  (define (make-resampleable
           name
           params
           log-density
           sample
           #!optional
           (resample #f)
           (support #f))
    (make-distk
      'resampleable
      (make-rdistr
        name
        params
        log-density
        sample
        (or resample (independence-resample sample))
        support
        (not resample))
      (void)))
  (define ((sample/log-density draw log-density) src)
    (let1 (x (draw src)) (values x (log-density x))))
  (define-syntax
    fp-horner
    (syntax-rules
      ()
      ((_ x a) a)
      ((_ x a b c ...) (fp-horner x (fp+ (fp* a x) b) c ...))))
  (define-syntax-rule
    (fp-horner/monic x b c ...)
    (fp-horner x (fp+ x b) c ...))
  (define (log-gamma x)
    (let-flonums
      'log-gamma
      (x)
      (cond ((fp<= x 0.0) +inf.0)
            ((fp> x 1e+308) +inf.0)
            ((fp< x 1.5)
             (receive
               (a b c)
               (if (fp< x 0.5)
                 (values (fpneg (fplog x)) (fp+ x 1.0) x)
                 (values 0.0 x (fp- x 1.0)))
               (fp+ a
                    (fp/ (fp* c
                              (fp-horner
                                b
                                3.13060547623
                                11.1667541262
                                -21.9698958928
                                -24.4387534237
                                -2.66685511495))
                         (fp-horner/monic
                           b
                           15.234687407
                           31.4690115749
                           11.9400905721
                           0.607771387771)))))
            ((fp< x 4.0)
             (fp/ (fp* (fp- x 2.0)
                       (fp-horner
                         x
                         4.16438922228
                         78.6994924154
                         137.519416416
                         -142.046296688
                         -78.3359299449))
                  (fp-horner/monic
                    x
                    43.3400022514
                    263.505074721
                    313.399215894
                    47.066876606)))
            ((fp< x 12.0)
             (fp/ (fp-horner
                    x
                    -2296.6072978
                    -40262.1119975
                    27464.7644705
                    230661.510616
                    -212159.572323)
                  (fp-horner/monic
                    x
                    -570.691009324
                    -24235.7409629
                    -146025.937511
                    -116328.495004)))
            (else
             (let* ((y (fplog x))
                    (k (fp+ (fp- (fp* x (fp- y 1.0)) (fp* 0.5 y))
                            0.918938533204673)))
               (if (fp> x 3000000.0)
                 k
                 (let* ((x1 (fp/ 1.0 x)) (x2 (fp* x1 x1)))
                   (fp+ k
                        (fp/ (fp* x1
                                  (fp-horner
                                    x2
                                    0.0692910599291889
                                    0.491731761050597
                                    0.279195317918525))
                             (fp-horner/monic
                               x2
                               6.0124592597641
                               3.3503438150223))))))))))
  (define libm-lgamma (foreign-lambda double "lgamma" double))
  (define (marsaglia src)
    (let* ((x (rng-double! src -1.0 1.0))
           (y (rng-double! src -1.0 1.0))
           (s (fp+ (fp* x x) (fp* y y))))
      (if (and (fp>= 1.0 s) (fp> s 0.0))
        (fp* x (fpsqrt (fp/ (fp* -2.0 (fplog s)) s)))
        (marsaglia src))))
  (define (normal-rng mu sd src)
    (if (fp> sd 0.0)
      (fp+ mu (fp* sd (marsaglia src)))
      (hakaru10-error 'parameter 'normal "normal: invalid parameters" mu sd)))
  (define (marsaglia-tsang shape scl src)
    (let* ((d (fp- shape (fp/ 1.0 3.0))) (c (fp/ 1.0 (fpsqrt (fp* 9.0 d)))))
      (let retry ()
        (let* ((v (let until ((v (normal-rng 1.0 c src)))
                    (if (fp> v 0.0) v (until (normal-rng 1.0 c src)))))
               (x (fp/ (fp- v 1.0) c))
               (sqr (fp* x x))
               (v3 (fp* (fp* v v) v))
               (u (rng-double! src 0.0 1.0)))
          (if (or (fp< u (fp- 1.0 (fp* 0.0331 (fp* sqr sqr))))
                  (fp< (fplog u)
                       (fp+ (fp* 0.5 sqr)
                            (fp* d (fp+ (fp- 1.0 v3) (fplog v3))))))
            (fp* (fp* scl d) v3)
            (retry))))))
  (define (gamma-rng shape scl src)
    (cond ((fp<= shape 0.0)
           (hakaru10-error
             'parameter
             'gamma
             "gamma: got a negative shape paramater"
             shape
             scl))
          ((fp<= scl 0.0)
           (hakaru10-error
             'parameter
             'gamma
             "gamma: got a negative scale paramater"
             shape
             scl))
          ((fp< shape 1.0)
           (let* ((g (gamma-rng (fp+ shape 1.0) 1.0 src))
                  (w (rng-double! src 0.0 1.0)))
             (fp* (fp* scl g) (fpexpt w (fp/ 1.0 shape)))))
          (else (marsaglia-tsang shape scl src))))
  (define (beta-rng a b src)
    (cond ((or (fp<= a 0.0) (fp<= b 0.0))
           (hakaru10-error
             'parameter
             'beta
             "beta: parameters must be positve"
             a
             b))
          ((and (fp<= a 1.0) (fp<= b 1.0))
           (let retry ()
             (let* ((u (rng-double! src 0.0 1.0))
                    (v (rng-double! src 0.0 1.0))
                    (x (fpexpt u (fp/ 1.0 a)))
                    (y (fpexpt v (fp/ 1.0 b))))
               (if (fp<= (fp+ x y) 1.0) (fp/ x (fp+ x y)) (retry)))))
          (else
           (let* ((ga (gamma-rng a 1.0 src)) (gb (gamma-rng b 1.0 src)))
             (fp/ ga (fp+ ga gb))))))
  (define (bern-with name p kernel)
    (let-flonums
      name
      (p)
      (let* ((lT (fplog p))
             (lF (fplog (fp- 1.0 p)))
             (log-density (λ (x) (if x lT lF))))
        (make-resampleable
          name
          (list p)
          log-density
          (λ (src)
             (if (fp<= (rng-double! src 0.0 1.0) p)
               (values #t lT)
               (values #f lF)))
          (kernel log-density)
          (τ `((#t unquote p) (#f unquote (fp- 1.0 p))))))))
  (define (bern p) (bern-with 'bern p (K #f)))
  (define ((flip-kernel log-density) src x ll)
    (let1 (x* (not x)) (values x* (log-density x*) 0.0)))
  (define (bern/flip p) (bern-with 'bern/flip p flip-kernel))
  (define (uniformly lst)
    (cond ((null? lst)
           (hakaru10-error 'parameter 'uniformly "uniformly: empty list"))
          ((not (list? lst))
           (hakaru10-error 'parameter 'uniformly "a list is expected" lst))
          (else
           (let* ((len (exact->inexact (length lst)))
                  (arr (list->vector lst))
                  (hi (sub1 (vector-length arr)))
                  (ll (fpneg (fplog len)))
                  (p (fp/ 1.0 len)))
             (make-resampleable
               'uniformly
               (list lst)
               (λ (x) (if (member x lst) ll -inf.0))
               (λ (src) (values (vector-ref arr (rng-integer! src 0 hi)) ll))
               #f
               (τ (map (λ (x) (cons x p)) lst)))))))
  (define (categorical-walk p alist)
    (let* ((entry (car alist)) (p0 (cdr entry)))
      (if (or (fp< p p0) (null? (cdr alist)))
        (values (car entry) (fplog p0))
        (categorical-walk (fp- p p0) (cdr alist)))))
  (define (categorical alist)
    (unless
      (and (pair? alist)
           (list? alist)
           (every (λ (e) (and (pair? e) (real? (cdr e)))) alist))
      (hakaru10-error
        'parameter
        'categorical
        "a non-empty alist ((value . weight) ...) is expected"
        alist))
    (let* ((weighted
             (map (λ (e) (cons (car e) (exact->inexact (cdr e)))) alist))
           (ptotal (fold (λ (e acc) (fp+ acc (cdr e))) 0.0 weighted))
           (normalized
             (map (λ (e) (cons (car e) (fp/ (cdr e) ptotal))) weighted)))
      (make-resampleable
        'categorical
        (list weighted)
        (λ (x) (cond ((assoc x normalized) => (o fplog cdr)) (else -inf.0)))
        (λ (src) (categorical-walk (rng-double! src 0.0 1.0) normalized))
        #f
        (τ (alist-copy normalized)))))
  (define (uniform lo hi)
    (let-flonums
      'uniform
      (lo hi)
      (let1 (l (fplog (fp/ 1.0 (fp- hi lo))))
            (make-resampleable
              'uniform
              (list lo hi)
              (λ (x)
                 (let-flonums
                   'uniform
                   (x)
                   (if (and (fp<= lo x) (fp<= x hi)) l -inf.0)))
              (λ (src) (values (rng-double! src lo hi) l))))))
  (define (normal mu sd)
    (let-flonums
      'normal
      (mu sd)
      (let* ((tau (fp/ 1.0 (fp* sd sd)))
             (lf (fplog (fp/ (fp/ tau 3.14159265358979) 2.0)))
             (log-density
               (λ (x)
                  (let-flonums
                    'normal
                    (x)
                    (let1 (y (fp- x mu))
                          (fp/ (fp+ (fpneg (fp* tau (fp* y y))) lf) 2.0))))))
        (make-resampleable
          'normal
          (list mu sd)
          log-density
          (sample/log-density (λ (src) (normal-rng mu sd src)) log-density)))))
  (define (gamma-log-density shape scl)
    (if (and (fp> shape 0.0) (fp> scl 0.0))
      (let ((shape-log-scl (fp* shape (fplog scl)))
            (log-g-shape (log-gamma shape))
            (shape-1 (fp- shape 1.0)))
        (λ (x)
           (let-flonums
             'gamma
             (x)
             (if (fp< x 0.0)
               -inf.0
               (fp- (fp+ (fp- (fpneg shape-log-scl) (fp/ x scl))
                         (fp* shape-1 (fplog x)))
                    log-g-shape)))))
      (K -inf.0)))
  (define (gamma shape scl)
    (let-flonums
      'gamma
      (shape scl)
      (let1 (log-density (gamma-log-density shape scl))
            (make-resampleable
              'gamma
              (list shape scl)
              log-density
              (sample/log-density
                (λ (src) (gamma-rng shape scl src))
                log-density)))))
  (define (beta-log-density a b)
    (if (or (fp<= a 0.0) (fp<= b 0.0))
      (λ (x)
         (hakaru10-error
           'parameter
           'beta
           "beta: parameters must be positve"
           a
           b))
      (let ((factor
              (fp- (fp- (log-gamma (fp+ a b)) (log-gamma a)) (log-gamma b)))
            (a-1 (fp- a 1.0))
            (b-1 (fp- b 1.0)))
        (λ (x)
           (let-flonums
             'beta
             (x)
             (if (or (fp< x 0.0) (fp> x 1.0))
               (hakaru10-error
                 'support
                 'beta
                 "beta: value must be between 0 and 1"
                 x)
               (fp+ (fp+ factor (fp* a-1 (fplog x)))
                    (fp* b-1 (fplog (fp- 1.0 x))))))))))
  (define (beta a b)
    (let-flonums
      'beta
      (a b)
      (let1 (log-density (beta-log-density a b))
            (make-resampleable
              'beta
              (list a b)
              log-density
              (sample/log-density (λ (src) (beta-rng a b src)) log-density)))))
  (define (poisson-log-density lam)
    (let1 (log-lam (fplog lam))
          (λ (x)
             (if (not (and (real? x) (integer? x)))
               -inf.0
               (let1 (x (inexact->exact x))
                     (cond ((= x 0) (fpneg lam))
                           ((fp= lam 0.0) -inf.0)
                           ((= x 1) (fp- log-lam lam))
                           ((= x 2) (fp- (fplog (fp/ (fp* lam lam) 2.0)) lam))
                           (else
                            (let1 (xf (exact->inexact x))
                                  (fp- (fp- (fp* xf log-lam) lam)
                                       (log-gamma (fp+ xf 1.0)))))))))))
  (define (knuth-poisson expl src)
    (let loop ((p (rng-double! src 0.0 1.0)) (k 1))
      (if (fp> p expl)
        (loop (fp* p (rng-double! src 0.0 1.0)) (add1 k))
        (sub1 k))))
  (define (poisson lam)
    (let-flonums
      'poisson
      (lam)
      (unless
        (fp>= lam 0.0)
        (hakaru10-error
          'parameter
          'poisson
          "In Poisson distribution, lambda cannot be negative"
          lam))
      (let ((expl (fpexp (fpneg lam))) (log-density (poisson-log-density lam)))
        (make-resampleable
          'poisson
          (list lam)
          log-density
          (if (fp= lam 0.0)
            (λ (src) (values 0 0.0))
            (sample/log-density
              (λ (src) (knuth-poisson expl src))
              log-density))))))
  (define (laplace mu sd)
    (let-flonums
      'laplace
      (mu sd)
      (let* ((log-2sd (fplog (fp* 2.0 sd)))
             (log-density
               (λ (x)
                  (let-flonums
                    'laplace
                    (x)
                    (fp- (fpneg log-2sd) (fp/ (fpabs (fp- x mu)) sd)))))
             (draw (λ (src)
                      (unless
                        (fp> sd 0.0)
                        (hakaru10-error
                          'parameter
                          'laplace
                          "laplace: invalid parameters"
                          mu
                          sd))
                      (let1 (u (rng-double! src 0.0 1.0))
                            (if (fp< u 0.5)
                              (fp+ mu (fp* sd (fplog (fp+ u u))))
                              (fp- mu (fp* sd (fplog (fp- (fp- 2.0 u) u)))))))))
        (make-resampleable
          'laplace
          (list mu sd)
          log-density
          (sample/log-density draw log-density)))))
  (define (dirichlet-gammas n a src)
    (if (zero? n)
      (values '() 0.0)
      (receive
        (xs total)
        (dirichlet-gammas (sub1 n) a src)
        (let1 (x (gamma-rng a 1.0 src)) (values (cons x xs) (fp+ x total))))))
  (define (dirichlet/symmetric n a)
    (unless
      (and (exact-integer? n) (positive? n))
      (hakaru10-error
        'parameter
        'dirichlet/symmetric
        "a positive exact integer is expected"
        n))
    (let-flonums
      'dirichlet/symmetric
      (a)
      (let* ((log-g-a (log-gamma a))
             (a-1 (fp- a 1.0))
             (log-g-sum (log-gamma (fold fp+ 0.0 (make-list n a))))
             (log-term
               (λ (y acc)
                  (fp+ acc
                       (fp- (fp* a-1 (fplog (exact->inexact y))) log-g-a))))
             (log-density
               (λ (xs)
                  (cond ((not (and (list? xs)
                                   (= (length xs) n)
                                   (every real? xs)))
                         -inf.0)
                        ((not (every positive? xs))
                         (hakaru10-error
                           'support
                           'dirichlet/symmetric
                           "dirichlet: all values must be between 0 and 1"
                           xs))
                        (else (fp+ (fold log-term 0.0 xs) log-g-sum))))))
        (make-resampleable
          'dirichlet/symmetric
          (list n a)
          log-density
          (sample/log-density
            (λ (src)
               (receive
                 (xs total)
                 (dirichlet-gammas n a src)
                 (map (λ (x) (fp/ x total)) xs)))
            log-density)))))
  (define (dirac x) (make-distk 'dirac #f x))
  (define (distk-sample dk src)
    (case (distk-kind dk)
      ((resampleable) ((rdistr-sample (distk-rd dk)) src))
      ((dirac) (values (distk-value dk) 0.0))
      ((observed)
       (let1 (x (distk-value dk))
             (values x ((rdistr-log-density (distk-rd dk)) x))))
      (else (hakaru10-error 'internal 'distk-sample "not a distk" dk))))
  (define (conditioned obs dk)
    (unless
      (distk? dk)
      (hakaru10-error
        'conditioning
        'condition
        "only distributions can be conditioned"
        dk))
    (case (distk-kind dk)
      ((resampleable) (make-distk 'observed (distk-rd dk) obs))
      ((observed)
       (hakaru10-error
         'conditioning
         'condition
         "an observed distribution cannot be conditioned again"
         obs
         dk))
      (else
       (hakaru10-error
         'conditioning
         'condition
         "a Dirac distribution cannot be conditioned"
         obs
         dk))))
  (define (condition obs d #!optional (internal? #f))
    (when (and internal? (internal? obs))
          (hakaru10-error
            'conditioning
            'condition
            "observation must be external to the model"
            obs))
    (if (procedure? d)
      (λ args (condition obs (apply d args)))
      (conditioned obs d)))
  (define (make-distribution
           name
           #!key
           log-density
           sample
           (resample #f)
           (support #f)
           (params '()))
    (define (check ok? what)
      (unless ok? (hakaru10-error 'construction 'make-distribution what name)))
    (check (and (procedure? log-density) (procedure? sample))
           "log-density: and sample: procedures are required")
    (check (or (not resample) (procedure? resample))
           "resample: must be a procedure or #f")
    (check (or (not support) (procedure? support))
           "support: must be a thunk or #f")
    (make-resampleable name params log-density sample resample support))
  (define (distribution-log-density dk x)
    (unless
      (distk? dk)
      (hakaru10-error 'parameter 'distribution-log-density "not a distk" dk))
    (if (eq? (distk-kind dk) 'dirac)
      (if (equal? x (distk-value dk)) 0.0 -inf.0)
      ((rdistr-log-density (distk-rd dk)) x)))
  (define (distribution-sample dk src)
    (unless
      (distk? dk)
      (hakaru10-error 'parameter 'distribution-sample "not a distk" dk))
    (distk-sample dk src))
  (define (distribution-support dk)
    (unless
      (distk? dk)
      (hakaru10-error 'parameter 'distribution-support "not a distk" dk))
    (if (eq? (distk-kind dk) 'dirac)
      (list (cons (distk-value dk) 1.0))
      (let1 (support (rdistr-support (distk-rd dk)))
            (and support (support))))))

data Normal = Normal !Double !Double !Double !Double
normal :: Double -> Double -> Normal
normal mu sd = Normal mu sd tau (log (tau / pi / 2))
 where tau = 1 / (sd*sd)

instance Distribution Normal Double where
  logDensity (Normal mu sd tau lf) x =
   (-tau * square (x - mu) + lf) / 2 where square y = y * y

(define (test/normal-density _)
  (⊦ (equal-approx? 1e-15) -0.923938533204673 (ld (normal 0 1) 0.1))
  (⊦ (equal-approx? 1e-15) -0.923938533204673 (ld (normal 0 1) 1/10))
  (⊦ (equal-approx? 1e-15)
     (- (+ (log 2.0) (* 0.5 (log (* 2.0 3.14159265358979)))))
     (ld (normal 1 2) 1.0)))

((eta 0.0) (memory #(8402704 2732752 1048576)) (stdout "") (stderr ""))

data Bern = Bern !Double !Double !Double
bern :: Double -> Bern
bern p = Bern p (log p) (log (1-p))

instance Distribution Bern Bool where
  logDensity (Bern _ lT lF) x = if x then lT else lF
  sample     (Bern p lT lF) g =
    let (t, g') = randomR (0, 1) g in
    if t <= p then (True, lT, g') else (False, lF, g')

(define (test/bern-density _)
  (let1 (d (bern 0.4))
        (⊦ (equal-approx? 1e-15) -0.916290731874155 (ld d #t))
        (⊦ (equal-approx? 1e-15) -0.510825623765991 (ld d #f))
        (⊦ equal?
           '(bern (0.4))
           (list (rdistr-name (distk-rd d)) (rdistr-params (distk-rd d)))))
  (⊦ equal? '(0.5) (rdistr-params (distk-rd (bern 1/2)))))

((eta 0.0) (memory #(8402704 2734752 1048576)) (stdout "") (stderr ""))

uniform lo hi = Uniform lo hi (log (recip (hi - lo)))

instance Distribution Uniform Double where
  logDensity (Uniform lo hi l) x =
    if lo <= x && x <= hi then l else minusInf

(define (test/uniform-density _)
  (let1 (d (uniform 10 20))
        (⊦ equal?
           (list -2.30258509299405 -2.30258509299405 -2.30258509299405 -inf.0)
           (map (λ (x) (ld d x)) '(15 10 20 9.99)))))

((eta 0.0) (memory #(8402704 2736880 1048576)) (stdout "") (stderr ""))

logGamma :: Double -> Double
logGamma x
    | x <= 0    = m_pos_inf
    -- Handle positive infinity. logGamma overflows before 1e308 so
    -- it's safe
    | x > 1e308 = m_pos_inf
    -- Normal cases
    | x < 1.5   = a + c *
                  ((((r1_4 * b + r1_3) * b + r1_2) * b + r1_1) * b + r1_0) /
                  ((((b + r1_8) * b + r1_7) * b + r1_6) * b + r1_5)
    | x < 4     = (x - 2) *
                  ((((r2_4 * x + r2_3) * x + r2_2) * x + r2_1) * x + r2_0) /
                  ((((x + r2_8) * x + r2_7) * x + r2_6) * x + r2_5)
    | x < 12    = ((((r3_4 * x + r3_3) * x + r3_2) * x + r3_1) * x + r3_0) /
                  ((((x + r3_8) * x + r3_7) * x + r3_6) * x + r3_5)
    | x > 3e6   = k
    | otherwise = k + x1 *
                  ((r4_2 * x2 + r4_1) * x2 + r4_0) /
                  ((x2 + r4_4) * x2 + r4_3)

(define (test/log-gamma _)
  (⊦ (equal-approx? 1e-14)
     '(1.09579799478688 3.1780538330717 7.53436424337661 39.3398841871995)
     (map log-gamma '(0.3 5.0 7.5 20.0)))
  (⊦ (equal-approx? 1e-08)
     (map libm-lgamma '(0.3 1.5 3.9 5.0 7.5 11.9 20.0))
     (map log-gamma '(0.3 1.5 3.9 5 7.5 11.9 20)))
  (⊦ (lambda (a b) (< (abs (- a b)) (* 1e-12 (abs a))))
     (libm-lgamma 10000000.0)
     (log-gamma 10000000.0))
  (⊦ equal? '(+inf.0 +inf.0 +inf.0) (map log-gamma '(0 -1.0 1.5e+308))))

((eta 0.0) (memory #(8402704 2738456 1048576)) (stdout "") (stderr ""))

gammaLogDensity :: Double -> Double -> Double -> Double
gammaLogDensity shape scl | shape > 0 && scl > 0 =
  let logShape =  log shape
      logGShape = logGamma shape in
  \x -> if x < 0 then minusInf else
          scl * logShape - scl * x + (shape - 1) * log x - logGShape
gammaLogDensity _ _ = \_ -> minusInf

(define (test/gamma-density _)
  (let1 (d (gamma 3 1/3))
        (⊦ (equal-approx? 1e-12) -0.397310314667487 (ld d 1))
        (⊦ (equal-approx? 1e-09)
           (- (+ (* 3.0 (log 3.0)) (* 2.0 (log 1.0))) 3.0 (log 2.0))
           (ld d 1.0))
        (⊦ = -inf.0 (ld d -0.5)))
  (⊦ = -inf.0 (ld (gamma 0 1) 1.0))
  (⊦ = -inf.0 (ld (gamma 2 -1) 1.0))
  (⊦ (equal-approx? 1e-12) (- (log 0.5) 0.5) (ld (gamma 1 2) 1.0))
  (⊦ (equal-approx? 1e-09)
     (- (+ (* 3.0 (log 3.0)) (* 2.0 (log 2.0))) 6.0 (log 2.0))
     (ld (gamma 3 1/3) 2.0))
  (⊦ (equal-approx? 1e-12) (- (- (log 2.0)) 1.5) (ld (gamma 1 2) 3.0)))

((eta 0.0) (memory #(8402704 2733560 1048576)) (stdout "") (stderr ""))

poisson :: Double -> Poisson
poisson lambda | lambda >= 0 = Poisson lambda (exp(-lambda))
poisson lambda = error "In Poisson distribution, lambda cannot be negative"

-- if lambda = 0 then poisson(0) = dirac(0)
instance Distribution Poisson Int where
  -- density is (lambda ^ x) * exp(-lambda) / x!
  logDensity (Poisson lambda expl) x
    | x == 0      = -lambda
    | lambda == 0 = minusInf
    | x == 1      = log lambda - lambda
    | x == 2      = log (lambda*lambda/2) - lambda
    | otherwise   = fromIntegral x * log lambda - lambda
                     - logGamma (fromIntegral x + 1)

(define (test/poisson-density _)
  (⊦ (equal-approx? 1e-12)
     '(-1.0 -1.0 -1.69314718055995 -2.79175946754953)
     (map (λ (k) (ld (poisson 1) k)) '(0 1 2 3)))
  (⊦ eqv? -0.0 (ld (poisson 0) 0))
  (⊦ = -inf.0 (ld (poisson 0) 1))
  (⊦ equal? '(-inf.0 -inf.0) (map (λ (k) (ld (poisson 1) k)) '(1.5 -1)))
  (⊦ = -1.69314718055995 (ld (poisson 1) 2.0))
  (⊦ eq? 'parameter (hakaru10-kind (τ (poisson -1))))
  (⊦ equal?
     "In Poisson distribution, lambda cannot be negative"
     (hakaru10-message (τ (poisson -1)))))

((eta 0.0) (memory #(8402704 2738944 1048576)) (stdout "") (stderr ""))

data Categorical a = Categorical ![(a,Double)]
categorical :: Eq a => [(a,Double)] -> Categorical a
categorical lst =
  let ptotal = sum $ map snd lst in
  Categorical $ ptotal `seq` map (\ (i,p) -> (i,p/ptotal)) lst

instance Eq a => Distribution (Categorical a) a where
  logDensity (Categorical lst) x = maybe minusInf log $ (lookup x lst)
  sample     (Categorical lst) g = loop p lst
    where
      (p, g1) = randomR (0, 1) g
      loop p ((x,p0):t) | p < p0 = (x,log p0,g1)
                        | otherwise = loop (p-p0) t

(define (test/categorical-density _)
  (let1 (d (categorical '((a . 1) (b . 3))))
        (⊦ (equal-approx? 1e-15)
           (list (log 0.25) (log 0.75) -inf.0)
           (map (λ (x) (ld d x)) '(a b c))))
  (let1 (d (categorical '((a . 1) (b . 2) (a . 1))))
        (⊦ = (log 0.25) (ld d 'a))
        (⊦ equal? '((a . 0.25) (b . 0.5) (a . 0.25)) (distribution-support d)))
  (⊦ equal?
     '((x . 1.0) (y . 3.0))
     (car (rdistr-params (distk-rd (categorical '((x . 1) (y . 3)))))))
  (⊦ eq? 'parameter (hakaru10-kind (τ (categorical '()))))
  (let1 (d (categorical '(((a) . 1) ("b" . 3))))
        (⊦ (equal-approx? 1e-15)
           (list (log 0.25) (log 0.75) -inf.0)
           (map (λ (x) (ld d x)) (list (list 'a) (string #\b) (list 'b))))))

((eta 0.0) (memory #(8402704 2739304 1048576)) (stdout "") (stderr ""))

uniformly :: Eq a => [a] -> Uniformly a
uniformly [] = error "uniformly: empty list"
uniformly lst = Uniformly (A.listArray (0,len-1) lst) (- log (fromIntegral len))
 where len = length lst

instance Eq a => Distribution (Uniformly a) a where
  logDensity (Uniformly a ll) x = if x `elem` a then ll else minusInf
  sample     (Uniformly a ll) g = let (i,g1) = randomR (A.bounds a) g in
                           (unsafeAt a i, ll, g1)

(define (test/uniformly _)
  (⊦ eq? 'parameter (hakaru10-kind (τ (uniformly '()))))
  (⊦ equal? "uniformly: empty list" (hakaru10-message (τ (uniformly '()))))
  (let1 (d (uniformly '(a b c)))
        (⊦ = (- (log 3.0)) (ld d 'b))
        (⊦ = -inf.0 (ld d 'z))
        (⊦ equal? '(a b c) (map car (distribution-support d))))
  (let1 (d (uniformly '((a) "b")))
        (⊦ equal?
           (list (- (log 2.0)) (- (log 2.0)) -inf.0)
           (map (λ (x) (ld d x)) (list (list 'a) (string #\b) (list 'b))))))

((eta 0.0) (memory #(8402704 2741680 1048576)) (stdout "") (stderr ""))

betaLogDensity :: Double -> Double -> Double -> Double
betaLogDensity a b | a <= 0 || b <= 0 =
  error "beta: parameters must be positve"
betaLogDensity a b =
  let factor = logGamma (a + b)
               - logGamma a
               - logGamma b
  in \x -> if x < 0 || x > 1 then error "beta: value must be between 0 and 1"
           else factor + (a - 1) * log x
                       + (b - 1) * log (1 - x)

(define (test/beta-density _)
  (⊦ (equal-approx? 1e-08) (log 1.5) (ld (beta 2 2) 0.5))
  (⊦ =
     (+ (+ (- (- (log-gamma 4.0) (log-gamma 2.0)) (log-gamma 2.0)) (log 0.5))
        (log 0.5))
     (ld (beta 2 2) 0.5))
  (⊦ eq? 'support (hakaru10-kind (τ (ld (beta 2 2) 1.5))))
  (⊦ equal?
     "beta: value must be between 0 and 1"
     (hakaru10-message (τ (ld (beta 2 2) 1.5))))
  (⊦ eq? 'parameter (hakaru10-kind (τ (ld (beta 0 2) 0.5))))
  (⊦ equal?
     "beta: parameters must be positve"
     (hakaru10-message (τ (ld (beta 0 2) 0.5)))))

((eta 0.0) (memory #(8402704 2746448 1048576)) (stdout "") (stderr ""))

(define (test/normalization _)
  (⊦ (equal-approx? 1e-12)
     1.0
     (+ (exp (ld (bern 0.3) #t)) (exp (ld (bern 0.3) #f))))
  (⊦ (equal-approx? 1e-12)
     1.0
     (fold +
           0.0
           (map (λ (x) (exp (ld (uniformly '(a b c d e f g)) x)))
                '(a b c d e f g))))
  (⊦ (equal-approx? 1e-12)
     1.0
     (fold +
           0.0
           (map (λ (x) (exp (ld (categorical '((a . 1) (b . 3) (c . 2.5))) x)))
                '(a b c))))
  (⊦ (equal-approx? 1e-12)
     '(1.0 1.0 1.0)
     (map support-mass
          (list (bern 0.3)
                (uniformly '(a b c d e f g))
                (categorical '((a . 1) (b . 3) (c . 2.5))))))
  (⊦ (equal-approx? 1e-09)
     1.0
     (fold + 0.0 (map (λ (k) (exp (ld (poisson 3.5) k))) (iota 61))))
  (⊦ (equal-approx? 0.0001)
     '(1.0 1.0 1.0 1.0)
     (list (trapezoid (normal 1 2) -19.0 21.0 20000)
           (trapezoid (gamma 3 1/3) 0.0 12.0 12000)
           (trapezoid (gamma 2 1) 0.0 40.0 40000)
           (trapezoid (beta 2 5) 0.0 1.0 10000))))

((eta 0.05) (memory #(8402704 2746624 1048576)) (stdout "") (stderr ""))

gamma_rng shape scl g | shape <  1.0  = (gvar2, g2)
                      where (gvar1, g1) = gamma_rng (shape + 1) scl g
                            (w,  g2) = randomR (0,1) g1
                            gvar2 = scl * gvar1 * (w ** recip shape) 

(define (test/sampler-moments _)
  (let* ((within (λ (tol expected got) (⊦ (equal-approx? tol) expected got)))
         (n 20000))
    (let1 (mv (mean+var (samples (normal 10 0.5) n)))
          (within 0.03 10.0 (car mv))
          (within 0.02 0.25 (cdr mv)))
    (let1 (mv (mean+var (samples (gamma 3 1/3) n)))
          (within 0.03 1.0 (car mv))
          (within 0.03 (/ 1.0 3.0) (cdr mv)))
    (let1 (mv (mean+var (samples (gamma 0.5 2) n)))
          (within 0.05 1.0 (car mv))
          (within 0.25 2.0 (cdr mv)))
    (within 0.01 (/ 2.0 7.0) (car (mean+var (samples (beta 2 5) n))))
    (let1 (mv (mean+var (samples (beta 0.5 0.5) n)))
          (within 0.02 0.5 (car mv))
          (within 0.01 0.125 (cdr mv)))
    (let1 (mv (mean+var (samples (poisson 3.5) n)))
          (within 0.06 3.5 (car mv))
          (within 0.15 3.5 (cdr mv)))
    (within 0.1 15.0 (car (mean+var (samples (uniform 10 20) n))))
    (within
      0.02
      0.3
      (/ (count identity (samples (bern 0.3) n)) (exact->inexact n)))))

((eta 0.509) (memory #(16805408 3229064 1048576)) (stdout "") (stderr ""))

  -- Knuth algorithm
  sample (Poisson lambda expl) g
    | lambda == 0 = (0,0,g)
    | otherwise = let (u, g1) = randomR (0, 1) g
                  in loop u 1 1 g1
    where
      -- acc is lamda^(k-1)/(k-1)!
      loop p k acc g | p > expl =
           let (u, g1) = randomR (0, 1) g
           in loop (p*u) (k+1) (acc*lambda/(fromIntegral k)) g1
                     | otherwise = (k-1, acc*expl, g)

(define (test/poisson-sample-ll _)
  (let1 (src (make-stdgen-source 17))
        (⊨ (every (λ (i)
                     (receive
                       (k ll)
                       (distribution-sample (poisson 3.5) src)
                       (= ll (ld (poisson 3.5) k))))
                  (iota 1000))))
  (⊦ equal?
     '(0 0.0)
     (receive
       (k ll)
       (distribution-sample (poisson 0) (make-stdgen-source 17))
       (list k ll))))

((eta 0.004) (memory #(16805408 2751736 1048576)) (stdout "") (stderr ""))

normal_rng mu sd g | sd > 0 = case marsaglia g of
                                ((x, _), g1) -> (mu + sd * x, g1)
normal_rng _ _ _ = error "normal: invalid parameters"

gamma_rng shape _   _ | shape <= 0.0  = error "gamma: got a negative shape paramater"
gamma_rng _     scl _ | scl <= 0.0  = error "gamma: got a negative scale paramater"

beta_rng a b g | a <= 1.0 && b <= 1.0 =
                 let (u, g1) = randomR (0.0, 1.0) g
                     (v, g2) = randomR (0.0, 1.0) g1
                     x = u ** (recip a)
                     y = v ** (recip b)
                 in  case (x+y) <= 1.0 of
                       True -> (x / (x + y), g2)
                       False -> beta_rng a b g2

(define (test/sampler-errors _)
  (⊦ eq?
     'parameter
     (hakaru10-kind
       (τ (distribution-sample (normal 0 0) (make-stdgen-source)))))
  (⊦ equal?
     "normal: invalid parameters"
     (hakaru10-message
       (τ (distribution-sample (normal 0 0) (make-stdgen-source)))))
  (⊦ eq?
     'parameter
     (hakaru10-kind
       (τ (distribution-sample (gamma -1 1) (make-stdgen-source)))))
  (⊦ equal?
     "gamma: got a negative shape paramater"
     (hakaru10-message
       (τ (distribution-sample (gamma -1 1) (make-stdgen-source)))))
  (⊦ equal?
     "gamma: got a negative scale paramater"
     (hakaru10-message
       (τ (distribution-sample (gamma 2 0) (make-stdgen-source)))))
  (⊦ eq? 'parameter (hakaru10-kind (τ (normal 'zero 1))))
  (let1 (src (make-stdgen-source 17))
        (⊦ equal?
           '(parameter parameter parameter)
           (map (λ (a b)
                   (hakaru10-kind (τ (distribution-sample (beta a b) src))))
                '(-0.5 0.5 0)
                '(1 -2 1)))
        (⊦ equal? '(18 1) (seeds src))))

((eta 0.0) (memory #(16805408 2756912 1048576)) (stdout "") (stderr ""))

-- resampling always flips, so fwd and bwd steps are 0
instance RDistribution Bern Bool where
{-
  resample d x _ g =
    let x' = not x in
    ((x', logDensity d x', g),0)
-}

(define (test/bern-flip _)
  (let* ((d (bern/flip 0.3))
         (rd (distk-rd d))
         (src (make-stdgen-source 17))
         (before (seeds src)))
    (⊦ equal?
       (list #f (log 0.7) 0.0)
       (receive
         (x ll fb)
         ((rdistr-resample rd) src #t (log 0.3))
         (list x ll fb)))
    (⊦ equal? before (seeds src)))
  (let* ((rd (distk-rd (bern 0.3))) (src (make-stdgen-source 17)))
    (⊦ equal?
       (list #t (log 0.3) (- (log 0.3) (log 0.7)))
       (receive
         (x ll fb)
         ((rdistr-resample rd) src #f (log 0.7))
         (list x ll fb)))))

((eta 0.0) (memory #(16805408 2756072 1048576)) (stdout "") (stderr ""))

conditioned :: a -> DistK KResampleable a -> DistK KObserved a
conditioned obs (Resampleable d) = Observed d obs

instance Condition (DistK KResampleable a) where
  type CondR (DistK KResampleable a) = DistK KObserved a
  type CondV (DistK KResampleable a) = a
  condition = conditioned

instance Condition b => Condition (a -> b) where
  type CondR (a -> b) = a -> CondR b
  type CondV (a -> b) = CondV b
  condition v f = \x -> condition v (f x)

sampleK :: DistK k a -> StdGen -> (a, LL, StdGen)
sampleK (Resampleable d) g = sample d g
sampleK (Dirac x) g = (x,0,g)
sampleK (Observed d x) g = (x,logDensity d x,g)

(define (test/condition _)
  (let* ((d (normal 0 1)) (o (condition 1.5 d)))
    (⊦ eq? 'observed (distk-kind o))
    (⊦ = 1.5 (distk-value o))
    (⊦ eq? (distk-rd d) (distk-rd o))
    (let1 (src (make-stdgen-source 17))
          (⊦ equal?
             (list 1.5 (ld d 1.5))
             (receive (x ll) (distk-sample o src) (list x ll)))
          (⊦ equal? '(18 1) (seeds src))))
  (let1 (o ((condition 1.0 normal) 0 1))
        (⊦ eq? 'observed (distk-kind o))
        (⊦ equal?
           '(normal (0.0 1.0))
           (list (rdistr-name (distk-rd o)) (rdistr-params (distk-rd o)))))
  (⊦ eq?
     'observed
     (distk-kind (((condition 2 (λ (a) (λ (b) (uniform a b)))) 0) 3)))
  (⊦ eq?
     'conditioning
     (hakaru10-kind (τ (condition #t (condition #f (bern 0.5))))))
  (⊦ eq? 'conditioning (hakaru10-kind (τ (condition 1 (dirac 1)))))
  (⊦ eq? 'conditioning (hakaru10-kind (τ ((condition 1 dirac) 1))))
  (⊦ eq? 'conditioning (hakaru10-kind (τ (condition 1 42))))
  (⊦ eq? 'conditioning (hakaru10-kind (τ (condition 'x (bern 0.5) symbol?))))
  (⊦ equal?
     "observation must be external to the model"
     (hakaru10-message (τ (condition 'x (bern 0.5) symbol?))))
  (let1 (d (dirac 3))
        (⊦ equal? '(dirac #f) (list (distk-kind d) (distk-rd d))))
  (⊦ equal?
     '(3 0.0)
     (receive
       (x ll)
       (distk-sample (dirac 3) (make-stdgen-source))
       (list x ll)))
  (⊦ equal? '(0.0 -inf.0) (list (ld (dirac 3) 3) (ld (dirac 3) 4))))

((eta 0.0) (memory #(16805408 2756552 1048576)) (stdout "") (stderr ""))

(define (test/replay _)
  (⊦ equal? '((#t #f #t) 2040924989 1624004889) (replay (bern 0.3)))
  (⊦ equal? '((b d c) 33947404 1403159235) (replay (uniformly '(a b c d e))))
  (⊦ equal?
     '((b b b) 2040924989 1624004889)
     (replay (categorical '((a . 1) (b . 3) (c . 2)))))
  (⊦ equal?
     '((12.9389958502438 13.273788866478 11.967275951981)
       2040924989
       1624004889)
     (replay (uniform 10 20)))
  (⊦ equal?
     '((9.39612042942084 9.41232243145482 9.44440922646125)
       238574758
       1847970949)
     (replay (normal 10 0.5)))
  (⊦ equal?
     '((0.380224876466386 1.17402250871286 0.695965406931075)
       1634784357
       1431524588)
     (replay (gamma 3 1/3)))
  (⊦ equal?
     '((1.66158483367182 5.89134661162112 3.29908537903117)
       1634784357
       1431524588)
     (replay (gamma 2.5 2)))
  (⊦ equal?
     '((0.230839286683393 0.364622341494618 5.99387735984082)
       507113220
       1216787979)
     (replay (gamma 0.5 2)))
  (⊦ equal?
     '((0.0859997609215855 0.152619827493028 0.323475141715806)
       470875069
       467317093)
     (replay (beta 2 5)))
  (⊦ equal?
     '((0.446268276733181 0.0880533825953212 0.0762135858315621)
       368601130
       810754135)
     (replay (beta 0.5 0.5)))
  (⊦ equal?
     '((0.160356098218376 0.29562975594083 0.260198785673769)
       609399710
       1534428535)
     (replay (beta 0.5 2)))
  (⊦ equal? '((2 3 8) 473751772 1482242513) (replay (poisson 3.5))))

((eta 0.007) (memory #(16805408 2759704 1048576)) (stdout "") (stderr ""))

laplace_rng :: (RandomGen g) => Double -> Double -> g -> (Double, g)
laplace_rng mu sd g = sample (randomR (0.0, 1.0) g)
   where sample (u, g1) = case u < 0.5 of
                            True  -> (mu + sd * log (u + u), g1)
                            False -> (mu - sd * log (2.0 - u - u), g1)

laplaceLogDensity :: Floating a => a -> a -> a -> a
laplaceLogDensity mu sd x = - log (2 * sd) - abs (x - mu) / sd

(define (test/laplace _)
  (let1 (d (laplace 1 2))
        (⊦ (equal-approx? 1e-15) (- (- (log 4.0)) 1.5) (ld d 4.0))
        (⊦ (equal-approx? 0.0001) 1.0 (trapezoid d -39.0 41.0 40000))
        (⊦ equal?
           '((-0.0627398739205491 0.153020152986878 -0.865576184156956)
             2040924989
             1624004889)
           (replay d))
        (let1 (mv (mean+var (samples d 20000)))
              (⊦ (equal-approx? 0.1) 1.0 (car mv))
              (⊦ (equal-approx? 0.6) 8.0 (cdr mv))))
  (⊦ eq?
     'parameter
     (hakaru10-kind
       (τ (distribution-sample (laplace 0 0) (make-stdgen-source))))))

((eta 0.067) (memory #(16805408 3343656 1048576)) (stdout "") (stderr ""))

dirichlet_rng :: (RandomGen g) => Int ->  Double -> g -> ([Double], g)
dirichlet_rng n' a g' = normalize (gammas g' n')
  where gammas g 0 = ([], 0, g)
        gammas g n = let (xs, total, g1) = gammas g (n-1)
                         ( x, g2) = gamma_rng a 1 g1
                     in ((x : xs), x+total, g2)
        normalize (b, total, h) = (map (/ total) b, h)

dirichletLogDensity :: [Double] -> [Double] -> Double
dirichletLogDensity a x | all (> 0) x = sum' (zipWith logTerm a x) + logGamma (sum a)
  where sum' = foldl' (+) 0
        logTerm b y = (b-1) * log y - logGamma b
dirichletLogDensity _ _ = error "dirichlet: all values must be between 0 and 1"

(define (test/dirichlet _)
  (let1 (d (dirichlet/symmetric 3 2))
        (⊦ (equal-approx? 1e-08) (log 3.6) (ld d (list 0.2 0.3 0.5)))
        (⊦ = -inf.0 (ld d (list 0.5 0.5)))
        (⊦ eq? 'support (hakaru10-kind (τ (ld d (list 0.5 0.5 0.0)))))
        (⊦ equal?
           '(((0.29568912068716 0.572538668183498 0.131772211129342))
             1901800541
             817484612)
           (let1 (r (replay d)) (cons (list (car (car r))) (cdr r))))
        (let* ((xs (samples d 20000))
               (mv1 (mean+var (map first xs)))
               (mv3 (mean+var (map third xs))))
          (⊨ (every (λ (v) (< (abs (- (fold + 0.0 v) 1.0)) 1e-12)) xs))
          (⊦ (equal-approx? 0.01) (/ 1.0 3.0) (car mv1))
          (⊦ (equal-approx? 0.01) (/ 1.0 3.0) (car mv3))
          (⊦ (equal-approx? 0.003) (/ 8.0 252.0) (cdr mv1))))
  (⊦ (equal-approx? 1e-08)
     (log 2.0)
     (ld (dirichlet/symmetric 3 1) (list 0.1 0.2 0.7))))

((eta 0.153) (memory #(16805408 5644072 1048576)) (stdout "") (stderr ""))

(define (test/make-distribution _)
  (let* ((d (make-distribution
              'exponential
              #:log-density
              (λ (x) (if (< x 0.0) -inf.0 (- x)))
              #:sample
              (λ (src)
                 (let1 (x (- (log (- 1.0 (rng-double! src 0.0 0.9)))))
                       (values x (- x))))
              #:params
              '(1.0)))
         (rd (distk-rd d))
         (src (make-stdgen-source 17)))
    (⊦ eq? 'resampleable (distk-kind d))
    (⊦ equal? '(exponential (1.0)) (list (rdistr-name rd) (rdistr-params rd)))
    (⊦ = -2.0 (ld d 2.0))
    (⊨ (not (distribution-support d)))
    (receive
      (x ll fb)
      ((rdistr-resample rd) src 3.0 -3.0)
      (⊦ = fb (+ ll 3.0))))
  (let* ((d (make-distribution
              'coin
              #:log-density
              (λ (x) (log 0.5))
              #:sample
              (λ (src) (values (= 0 (rng-integer! src 0 1)) (log 0.5)))
              #:resample
              (λ (src x ll) (values (not x) ll 0.0))
              #:support
              (τ '((#t . 0.5) (#f . 0.5)))))
         (rd (distk-rd d)))
    (⊦ equal?
       '(#f 0.0)
       (receive
         (x ll fb)
         ((rdistr-resample rd) (make-stdgen-source) #t (log 0.5))
         (list x fb)))
    (⊦ = 1.0 (support-mass d)))
  (⊦ eq?
     'construction
     (hakaru10-kind (τ (make-distribution 'broken #:sample (λ (src) 0))))))

((eta 0.001) (memory #(16805408 5648272 1048576)) (stdout "") (stderr ""))

(define (test/support-and-printing _)
  (⊦ equal? '((#t . 0.25) (#f . 0.75)) (distribution-support (bern 1/4)))
  (⊦ equal?
     '((#t . 0.25) (#f . 0.75))
     (distribution-support (condition #f (bern/flip 0.25))))
  (⊦ equal? '((7 . 1.0)) (distribution-support (dirac 7)))
  (⊦ equal?
     '(#f #f #f)
     (map distribution-support (list (normal 0 1) (poisson 2) (beta 1 1))))
  (let1 (d (uniformly '(x y z)))
        (⊨ (every (λ (x) (pair? (assq x (distribution-support d))))
                  (samples d 200))))
  (⊦ eq? 'parameter (hakaru10-kind (τ (distribution-support 'normal))))
  (⊦ equal?
     '("#<distk resampleable normal (0.0 1.0)>"
       "#<distk observed normal (0.0 1.0) 1.5>"
       "#<distk dirac (a b)>"
       "#<rdistr categorical (((\"x\" . 1.0)))>")
     (map (λ (x) (format "~a" x))
          (list (normal 0 1)
                (condition 1.5 (normal 0 1))
                (dirac '(a b))
                (distk-rd (categorical '(("x" . 1))))))))

((eta 0.0) (memory #(16805408 2772576 1048576)) (stdout "") (stderr ""))

(module
  (aux hakaru10 util)
  *
  (import
    scheme
    (chicken base)
    (chicken sort)
    (chicken flonum)
    (only (chicken string) ->string)
    (only srfi-1 count)
    srfi-69
    (aux base)
    (only (aux hakaru10 distribution) hakaru10-error))
  (define generic-orders
    (list (cons boolean? (λ (a b) (and (not a) b)))
          (cons number?
                (λ (a b)
                   (cond ((not (and (real? a) (real? b)))
                          (generic<?
                            (list (real-part a) (imag-part a))
                            (list (real-part b) (imag-part b))))
                         ((= a b) (and (exact? a) (inexact? b)))
                         (else (< a b)))))
          (cons char? char<?)
          (cons string? string<?)
          (cons symbol?
                (λ (a b) (string<? (symbol->string a) (symbol->string b))))
          (cons (λ (v) (or (null? v) (pair? v)))
                (λ (a b)
                   (cond ((null? a) (pair? b))
                         ((null? b) #f)
                         ((generic<? (car a) (car b)) #t)
                         ((generic<? (car b) (car a)) #f)
                         (else (generic<? (cdr a) (cdr b))))))
          (cons vector?
                (λ (a b) (generic<? (vector->list a) (vector->list b))))))
  (define (generic-order v)
    (let loop ((rank 0) (orders generic-orders))
      (cond ((null? orders) (values rank (λ (a b) #f)))
            (((caar orders) v) (values rank (cdar orders)))
            (else (loop (add1 rank) (cdr orders))))))
  (define (generic<? a b)
    (let-values
      (((ra less?) (generic-order a)) ((rb _) (generic-order b)))
      (or (< ra rb) (and (= ra rb) (less? a b)))))
  (define (histogram-sort alist)
    (sort alist (λ (p q) (generic<? (car p) (car q)))))
  (define (hash-table->histogram table)
    (histogram-sort (hash-table->alist table)))
  (define (hist xs)
    (let1 (table (make-hash-table equal?))
          (for-each (λ (x) (hash-table-update!/default table x add1 0)) xs)
          (hash-table->histogram table)))
  (define (frequency pred xs)
    (fp/ (exact->inexact (count pred xs)) (exact->inexact (length xs))))
  (define (momenta xs)
    (let* ((xs (map exact->inexact xs))
           (n (exact->inexact (length xs)))
           (av (fp/ (foldl fp+ 0.0 xs) n))
           (sqr (λ (x) (fp* x x))))
      (cons av
            (fp/ (foldr (λ (x acc) (fp+ (sqr (fp- x av)) acc)) 0.0 xs)
                 (fp- n 1.0)))))
  (define (to-bin step x)
    (let1 (b (fpfloor (fp/ x step))) (fp* (if (fp= b 0.0) 0.0 b) step)))
  (define (tabulate step xs)
    (let* ((step (exact->inexact step))
           (quantum (fp/ 1.0 (exact->inexact (length xs))))
           (insert!
             (λ (x table)
                (hash-table-update!/default
                  table
                  (to-bin step (exact->inexact x))
                  (λ (old) (fp+ quantum old))
                  0.0)
                table)))
      (hash-table->histogram (foldr insert! (make-hash-table equal?) xs))))
  (define (tabulate/bins nbins xs)
    (let* ((ys (map exact->inexact xs))
           (mn (foldl fpmin (car ys) (cdr ys)))
           (mx (foldl fpmax (car ys) (cdr ys)))
           (step-raw
             (if (fp= mn mx) 1.0 (fp/ (fp- mx mn) (exact->inexact nbins)))))
      (tabulate (fpexpt 10.0 (round (fp/ (fplog step-raw) (fplog 10.0)))) ys)))
  (define (kl p q)
    (let1 (q (alist->hash-table q equal?))
          (foldl (Λ (,sum (,i unquote pi))
                    (let1 (qi (hash-table-ref
                                q
                                i
                                (τ (hakaru10-error
                                     'support
                                     'kl
                                     (string-append
                                       "Undefined KL divergence - q is 0 when p is not at point "
                                       (->string i))
                                     i))))
                          (+ sum (* pi (log (/ pi qi))))))
                 0.0
                 (histogram-sort p))))
  (define (plot h #!optional (path "/tmp/plot.dat"))
    (with-output-to-file
      path
      (τ (print "x  y")
         (for-each (Λ ((,k unquote v)) (print k " " v)) (histogram-sort h))))
    path))

-- Very simple histogram, for a discrete distribution
hist :: Ord t => [t] -> M.Map t Int
hist = foldr (\e -> M.insertWith' (+) e 1) M.empty

(define (test/hist _)
  (⊦ equal? '((#f . 1) (#t . 2) (1 . 1) (2 . 2)) (hist '(#t #f #t 2 1 2)))
  (⊦ equal? '() (hist '()))
  (⊦ equal?
     '((#\a . 1)
       ("a" . 2)
       (a . 1)
       (b . 1)
       (() . 1)
       ((1) . 1)
       ((1 2) . 1)
       ((2) . 1)
       (#(1) . 1))
     (hist '(b (2) #(1) "a" (1 2) a () "a" (1) #\a))))

((eta 0.0) (memory #(6291456 2020976 1048576)) (stdout "") (stderr ""))

(define (test/generic<? _)
  (⊨ (generic<? #f #t))
  (⊭ (generic<? #t #f))
  (⊨ (generic<? #t -1))
  (⊨ (generic<? 1 1.5))
  (⊨ (generic<? 1 1.0))
  (⊭ (generic<? 1.0 1))
  (⊭ (generic<? 2 2))
  (⊨ (generic<? 'z '()))
  (⊨ (generic<? '(1 2) '(1 2 0)))
  (⊨ (generic<? '(1 2) '(1 3)))
  (⊨ (generic<? '(9) #(0)))
  (⊨ (generic<? #(1 2) #(1 2 0)))
  (⊭ (generic<? #(1 3) #(1 2 0))))

((eta 0.0) (memory #(6291456 2017120 1048576)) (stdout "") (stderr ""))

(define (test/frequency _)
  (⊦ equal? 0.75 (frequency even? '(2 4 5 6)))
  (⊦ equal? 0.0 (frequency symbol? '(1 2))))

((eta 0.0) (memory #(6291456 2018688 1048576)) (stdout "") (stderr ""))

-- compute the first two momenta
momenta :: [Double] -> (Double,Double)
momenta lst = (av,sigma2)
  where
    n   = fromIntegral $ length lst
    av  = sum lst / n
    sigma2 = (foldr (\x acc -> sqr (x - av) + acc) 0 lst) / (n - 1)
    sqr x = x * x

(define (test/momenta _)
  (⊦ equal? '(2.5 . 1.66666666666667) (momenta '(1.0 2.0 3.0 4.0)))
  (⊦ equal? '(2.5 . 1.66666666666667) (momenta '(1 2 3 4)))
  (⊦ equal? '(0.2602 . 0.0728702) (momenta '(0.1 0.2 0.3 0.7 0.001))))

((eta 0.0) (memory #(6291456 2020768 1048576)) (stdout "") (stderr ""))

-- Tabulate with the given bin size
tabulate :: Double -> [Double] -> Histogram
tabulate step lst =
  foldr (\x -> M.insertWith' (+) (to_bin x) quantum) M.empty lst
 where
   quantum = 1 / fromIntegral (length lst)
   to_bin x = fromInteger (floor (x / step)) * step

(define (test/tabulate _)
  (⊦ equal?
     '((-0.5 . 0.25) (0.0 . 0.5) (0.5 . 0.25))
     (tabulate 0.5 '(0.1 0.2 0.7 -0.1)))
  (⊦ equal? '((0.0 . 0.5) (1.0 . 0.5)) (tabulate 1 '(-0.0 0.0 1.5 1.0))))

((eta 0.0) (memory #(6291456 2022976 1048576)) (stdout "") (stderr ""))

-- The argument is the _hint_ to the number of bins
tabulate' :: Int -> [Double] -> Histogram
tabulate' nbins lst =
  foldr (\x -> M.insertWith' (+) (to_bin x) quantum) M.empty lst
 where
   mn = minimum lst
   mx = maximum lst
   quantum = 1 / fromIntegral (length lst)
   step_raw = if mn == mx then 1 else (mx - mn) / fromIntegral nbins
   step  = 10.0 ** (fromInteger . round $ log step_raw / log 10)
   to_bin x = fromInteger (floor (x / step)) * step

(define (test/tabulate/bins _)
  (⊦ equal?
     '((0.0 . 0.0909090909090909)
       (0.1 . 0.0909090909090909)
       (0.2 . 0.181818181818182)
       (0.4 . 0.0909090909090909)
       (0.5 . 0.181818181818182)
       (0.6 . 0.0909090909090909)
       (0.8 . 0.0909090909090909)
       (0.9 . 0.0909090909090909)
       (1.0 . 0.0909090909090909))
     (tabulate/bins 10 '(0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0)))
  (⊦ equal?
     '((0.0 . 0.909090909090909) (10.0 . 0.0909090909090909))
     (tabulate/bins 3 '(0.0 1.0 2.0 3.0 4.0 5.0 6.0 7.0 8.0 9.0 10.0)))
  (⊦ equal?
     '((0.0 . 0.333333333333333)
       (200.0 . 0.333333333333333)
       (1000.0 . 0.333333333333333))
     (tabulate/bins 4 '(0.0 250.0 1000.0)))
  (⊦ equal? '((3.0 . 1.0)) (tabulate/bins 5 '(3.0 3.0)))
  (⊦ equal?
     '((0.0 . 0.5) (300.0 . 0.5))
     (tabulate/bins 1 '(0.0 316.227766016838))))

((eta 0.0) (memory #(6291456 2025232 1048576)) (stdout "") (stderr ""))

kl :: Histogram -> Histogram -> Double
kl p q = M.foldlWithKey' f 0 p
 where f sum i pi | Just qi <- M.lookup i q = sum + pi * log (pi / qi)
                  | otherwise =
         error $ "Undefined KL divergence - q is 0 when p is not at point " ++
                 show i

(define (test/kl _)
  (define p (tabulate 0.5 '(0.1 0.2 0.7 -0.1)))
  (⊦ equal? 0.0 (kl p p))
  (⊦ equal?
     0.14384103622589
     (kl '((0.0 . 0.5) (1.0 . 0.5)) '((1.0 . 0.75) (0.0 . 0.25) (2.0 . 0.0))))
  (⊦ equal? 0.693147180559945 (kl '((0.0 . 1.0)) '((0.0 . 0.5) (0.0 . 0.25))))
  (⊦raises (exn hakaru10) (kl p '((0.0 . 0.5) (0.5 . 0.5))))
  (⊦ equal?
     '(support
        .
        "Undefined KL divergence - q is 0 when p is not at point -0.5")
     (condition-case
       (kl p '((0.0 . 0.5) (0.5 . 0.5)))
       (c (hakaru10)
          (cons (get-condition-property c 'hakaru10 'kind)
                (get-condition-property c 'exn 'message))))))

((eta 0.0) (memory #(6291456 2031560 1048576)) (stdout "") (stderr ""))

plot :: Histogram -> IO ()
plot dat = do
  withFile fname WriteMode write_data
 where
   fname = "/tmp/plot.dat"
   write_data h = do
     hPutStrLn h "x  y"
     mapM_ (\ (k,v) -> hPutStrLn h $ unwords [show k,show v]) $ M.toList dat

_ =
 let n = 10000 in plot (tabulate 0.2 $ mcmC (n `div` 2) (dist normal 10 2) ++ mcmC (n `div` 2) (dist gamma 3 (1/3)))

(define (test/plot _)
  (define path (create-temporary-file "dat"))
  (⊦ equal? path (plot '((0.5 . 0.25) (-0.5 . 0.25) (0.0 . 0.5)) path))
  (⊦ equal?
     '("x  y" "-0.5 0.25" "0.0 0.5" "0.5 0.25")
     (with-input-from-file path read-lines))
  (delete-file path))

((eta 0.001) (memory #(6291456 2033072 1048576)) (stdout "") (stderr ""))

(define (nnot p b) (if b p 1.0))

(define ((nor strength-x strength-y noise) x y)
  (bern (- 1.0
           (* (* (nnot (- 1.0 strength-x) x) (nnot (- 1.0 strength-y) y))
              (- 1.0 noise)))))

(define grass
  (do/monad
    (← rain (dist bern 0.3))
    (← sprinkler (dist bern 0.5))
    (← grass-is-wet (dist (condition #t (nor 0.9 0.8 0.1)) rain sprinkler))
    (return rain)))

(((mcmC 10 grass) (#t #t #t #f #f #f #f #f #f #f))
 ((exact-posterior grass) ((#f . 0.53152855727963) (#t . 0.46847144272037)))
 ((likelihood-weighting 2000 grass)
  ((#f . 0.53654028664459) (#t . 0.46345971335541))))

(define (pbern* bern) (dist bern 0.4))

(define pnorm (dist normal 10.0 0.5))

(define pdep1 (do/monad (← x pnorm) (diracN (sexp+ x 1.0))))

(define (deeper m) (do/monad (← x m) (dist normal x 3.0)))

(define (phierN n) (fold (λ (_ m) (deeper m)) (dist normal 0.0 1.0) (iota n)))

(define mixng
  (do/monad
    (← x (dist normal 0.0 1.0))
    (if_ (lift (λ (xv) (> xv 0.0)) x)
         (dist normal 10.0 2.0)
         (dist gamma 3.0 (/ 1.0 3.0)))))

(define (mixng′* bern)
  (do/monad
    (← x (dist bern 0.5))
    (if_ x (dist normal 10.0 2.0) (dist gamma 3.0 (/ 1.0 3.0)))))

(define (biased-coin c1 c2)
  (do/monad
    (← b (dist beta 1.0 1.0))
    (dist (condition c1 bern) b)
    (dist (condition c2 bern) b)
    (return b)))

(module
  (aux hakaru10 engine)
  *
  (import
    scheme
    (only (scheme base)
          vector-map
          vector-for-each
          make-parameter
          exact-integer?)
    (chicken base)
    (chicken fixnum)
    (chicken flonum)
    (only (chicken format) fprintf)
    (only (chicken port) with-output-to-port)
    (only (chicken sort) sort!)
    (only (chicken condition)
          make-property-condition
          signal
          condition-case
          with-exception-handler
          condition-predicate
          condition-property-accessor)
    (only srfi-1 any every find delete-duplicates append-map list-index iota)
    (only srfi-69 make-hash-table hash-table-exists? hash-table-set!)
    (aux base)
    (only (aux category state)
          category-state?
          state
          run-state
          eval-state
          exec-state
          get
          gets
          modify
          return
          >>=)
    (only (aux category monad state) do/monad ▷ when/monad)
    (aux hakaru10 random)
    (aux hakaru10 distribution))
  (declare (local))
  (define-record nstate val ll dst tstamp)
  (set-record-printer!
    nstate
    (λ (s port)
       (fprintf
         port
         "#<nstate ~s ll ~s t ~s>"
         (nstate-val s)
         (nstate-ll s)
         (nstate-tstamp s))))
  (define-record
    node
    chain
    addr
    kind
    role
    parents
    fn
    update
    scope
    deps
    branch
    state
    inactive
    delayed
    pending
    queued)
  (set-record-printer!
    node
    (λ (n port) (fprintf port "#<node ~a ~a>" (node-addr n) (node-kind n))))
  (define (node-value n) (nstate-val (node-state n)))
  (define (node-inactive? n) (fx> (node-inactive n) 0))
  (define-record
    chain
    rng
    semantics
    now
    nodes
    count
    llcum
    sampleable
    samp
    card-change
    queue
    pending
    journal
    scope
    step-hook)
  (set-record-printer!
    chain
    (λ (ch port)
       (fprintf
         port
         "#<chain ~a nodes, now ~a>"
         (chain-count ch)
         (chain-now ch))))
  (define (new-chain src semantics)
    (unless
      (memq semantics '(upstream extended))
      (hakaru10-error
        'construction
        'mcmC
        "semantics must be 'upstream or 'extended"
        semantics))
    (make-chain
      src
      semantics
      1
      (make-vector 64 #f)
      0
      0.0
      '()
      (vector)
      0
      #f
      '()
      #f
      '()
      #f))
  (define (chain-node ch a) (vector-ref (chain-nodes ch) a))
  (define (chain-store! ch n)
    (let ((a (node-addr n)) (v (chain-nodes ch)))
      (when (fx>= a (vector-length v))
            (chain-nodes-set!
              ch
              (vector-resize v (fx* 2 (vector-length v)) #f)))
      (vector-set! (chain-nodes ch) a n)
      (chain-count-set! ch a)))
  (define (chain-fold kons knil ch #!optional (lo 0) (hi (chain-count ch)))
    (let loop ((a (fx+ lo 1)) (acc knil))
      (if (fx> a hi) acc (loop (fx+ a 1) (kons (chain-node ch a) acc)))))
  (define (chain-freeze! ch)
    (chain-samp-set! ch (list->vector (chain-sampleable ch)))
    (chain-queue-set!
      ch
      (make-update-queue (make-vector (fx+ 1 (chain-count ch)) #f) 0))
    (chain-pending-set! ch '())
    (chain-scope-set! ch '()))
  (define (extended? ch) (eq? (chain-semantics ch) 'extended))
  (define (journal-begin! ch)
    (when (chain-journal ch) (journal-rollback! ch))
    (let ((now (chain-now ch)) (card (chain-card-change ch)))
      (chain-journal-set!
        ch
        (list (τ (chain-now-set! ch now) (chain-card-change-set! ch card))))))
  (define (journal-commit! ch) (chain-journal-set! ch #f))
  (define (journal-rollback! ch)
    (for-each (λ (undo) (undo)) (chain-journal ch))
    (chain-journal-set! ch #f)
    (chain-pending-set! ch '())
    (update-queue-clear! (chain-queue ch)))
  (define-syntax-rule
    (journaled-set! (getter setter) n v)
    (let* ((n* n)
           (ch (node-chain n*))
           (old (getter n*))
           (journal (chain-journal ch)))
      (when journal (chain-journal-set! ch (cons (τ (setter n* old)) journal)))
      (setter n* v)))
  (define (set-node-state! n s)
    (journaled-set! (node-state node-state-set!) n s))
  (define (set-inactive! n k)
    (journaled-set! (node-inactive node-inactive-set!) n k))
  (define (set-delayed! n flag)
    (unless
      (eq? flag (node-delayed n))
      (journaled-set! (node-delayed node-delayed-set!) n flag)))
  (define (set-pending! n flag)
    (unless
      (eq? flag (node-pending n))
      (journaled-set! (node-pending node-pending-set!) n flag)
      (when flag
            (let1 (ch (node-chain n))
                  (chain-pending-set! ch (cons n (chain-pending ch)))))))
  (define-record update-queue heap size)
  (define-syntax-rule (node<? a b) (fx< (node-addr a) (node-addr b)))
  (define (enqueue! q n)
    (unless
      (node-queued n)
      (node-queued-set! n #t)
      (let ((heap (update-queue-heap q)) (i (update-queue-size q)))
        (update-queue-size-set! q (fx+ i 1))
        (let up ((i i))
          (if (fx= i 0)
            (vector-set! heap 0 n)
            (let* ((p (fxshr (fx- i 1) 1)) (parent (vector-ref heap p)))
              (if (node<? n parent)
                (begin (vector-set! heap i parent) (up p))
                (vector-set! heap i n))))))))
  (define (dequeue! q)
    (let* ((heap (update-queue-heap q))
           (size (fx- (update-queue-size q) 1))
           (top (vector-ref heap 0))
           (last (vector-ref heap size)))
      (update-queue-size-set! q size)
      (node-queued-set! top #f)
      (let down ((i 0))
        (let* ((l (fx+ 1 (fx* 2 i))) (r (fx+ l 1)))
          (if (fx>= l size)
            (vector-set! heap i last)
            (let* ((c (if (and (fx< r size)
                               (node<? (vector-ref heap r) (vector-ref heap l)))
                        r
                        l))
                   (child (vector-ref heap c)))
              (if (node<? child last)
                (begin (vector-set! heap i child) (down c))
                (vector-set! heap i last))))))
      top))
  (define (update-queue-empty? q) (fx= 0 (update-queue-size q)))
  (define (update-queue-clear! q)
    (let ((heap (update-queue-heap q)))
      (do ((i 0 (fx+ i 1)))
          ((fx= i (update-queue-size q)))
        (node-queued-set! (vector-ref heap i) #f))
      (update-queue-size-set! q 0)))
  (define (enqueue-all! ch nodes)
    (let1 (q (chain-queue ch)) (for-each (λ (n) (enqueue! q n)) nodes)))
  (define-record nt leaves fn)
  (set-record-printer!
    nt
    (λ (e port)
       (fprintf port "#<nt ~s>" (map node-addr (vector->list (nt-leaves e))))))
  (define (sexp? x) (or (node? x) (nt? x)))
  (define (sexp-chain e)
    (node-chain (if (node? e) e (vector-ref (nt-leaves e) 0))))
  (define (sexp-leaves e)
    (cond ((node? e) (vector e)) ((nt? e) (nt-leaves e)) (else (vector))))
  (define (leaf-values leaves) (vector-map node-value leaves))
  (define-record deferred fn args)
  (set-record-printer! deferred (λ (d port) (fprintf port "#<deferred>")))
  (define (deferrable? e) (or (not (hakaru10-error? e)) (parameter-error? e)))
  (define (force-value x)
    (if (deferred? x) ((deferred-fn x) (deferred-args x)) x))
  (define (first-deferred vs)
    (let loop ((i 0))
      (cond ((fx= i (vector-length vs)) #f)
            ((deferred? (vector-ref vs i)) (vector-ref vs i))
            (else (loop (fx+ i 1))))))
  (define (forced-values vs)
    (if (first-deferred vs) (vector-map force-value vs) vs))
  (define (forced-list vs) (if (any deferred? vs) (map force-value vs) vs))
  (define (call/deferring g args)
    (condition-case
      (g args)
      (e (exn) (if (deferrable? e) (make-deferred g args) (signal e)))))
  (define (lazy-value g vs inactive?)
    (if inactive?
      (or (first-deferred vs) (call/deferring g vs))
      (g (forced-values vs))))
  (define lazy-lifting (make-parameter #f))
  (define (sexp-value e)
    (cond ((node? e) (force-value (node-value e)))
          ((nt? e) ((nt-fn e) (forced-values (leaf-values (nt-leaves e)))))
          (else e)))
  (define (check-not-model who x)
    (when (model? x)
          (hakaru10-error
            'construction
            who
            "a model where an SExp is expected; bind it with ← first"
            x)))
  (define (internal-part x)
    (let1 (r (internal-part/bounded x 100000))
          (cond ((not (fixnum? r)) r)
                ((fx>= r 0) #f)
                (else (internal-part/cyclic x)))))
  (define (internal-part/bounded x budget)
    (cond ((fx< budget 0) budget)
          ((or (number? x) (null? x) (boolean? x) (symbol? x)) budget)
          ((pair? x)
           (let1 (r (internal-part/bounded (car x) (fx- budget 1)))
                 (if (fixnum? r) (internal-part/bounded (cdr x) r) r)))
          ((vector? x)
           (let loop ((i 0) (b (fx- budget 1)))
             (if (fx= i (vector-length x))
               b
               (let1 (r (internal-part/bounded (vector-ref x i) b))
                     (if (fixnum? r) (loop (fx+ i 1) r) r)))))
          ((or (sexp? x) (model? x)) x)
          (else budget)))
  (define (internal-part/cyclic x)
    (let1 (seen (make-hash-table eq?))
          (let walk ((x x))
            (cond ((or (pair? x) (vector? x))
                   (and (not (hash-table-exists? seen x))
                        (begin
                          (hash-table-set! seen x #t)
                          (if (pair? x)
                            (or (walk (car x)) (walk (cdr x)))
                            (any walk (vector->list x))))))
                  ((or (sexp? x) (model? x)) x)
                  (else #f)))))
  (define (check-no-sexp who x)
    (let1 (y (internal-part x))
          (cond ((sexp? y)
                 (hakaru10-error
                   'construction
                   who
                   "an SExp inside a Scheme value; combine with collect, pair or lift"
                   y))
                ((model? y)
                 (hakaru10-error
                   'construction
                   who
                   "a model inside a Scheme value"
                   y)))))
  (define (check-lifted-value who v #!optional deep?)
    (cond ((or (number? v) (boolean? v)) v)
          ((sexp? v)
           (hakaru10-error
             'construction
             who
             "a lifted function answered an SExp"
             v))
          ((model? v)
           (hakaru10-error
             'construction
             who
             "a lifted function answered a model"
             v))
          ((and deep? (or (pair? v) (vector? v))) (check-no-sexp who v) v)
          (else v)))
  (define (lift* f es #!optional (checked? #t))
    (for-each (λ (e) (check-not-model 'lift e)) es)
    (if (any sexp? es)
      (lift/fuse f es checked?)
      (let1 (v (lift/apply f es))
            (if checked? (check-lifted-value 'lift v #t) v))))
  (define (lift/apply f vs)
    (cond ((not (lazy-lifting)) (apply f (forced-list vs)))
          ((find deferred? vs))
          (else (call/deferring (λ (args) (apply f args)) vs))))
  (define (lift/fuse f es checked?)
    (let* ((leaves
             (delete-duplicates
               (append-map (λ (e) (vector->list (sexp-leaves e))) es)
               eq?))
           (owner (node-chain (car leaves)))
           (position (λ (n) (list-index (λ (m) (eq? n m)) leaves)))
           (extractors (map (λ (e) (lift/extractor position e)) es)))
      (unless
        (every (λ (n) (eq? (node-chain n) owner)) leaves)
        (hakaru10-error
          'construction
          'lift
          "nodes of different model runs"
          es))
      (make-nt
        (list->vector leaves)
        (if checked?
          (let1 (first? #t)
                (λ (vs)
                   (let1 (v (check-lifted-value
                              'lift
                              (apply f (map (λ (x) (x vs)) extractors))
                              first?))
                         (set! first? #f)
                         v)))
          (λ (vs) (apply f (map (λ (x) (x vs)) extractors)))))))
  (define (lift/extractor position e)
    (cond ((node? e) (let1 (i (position e)) (λ (vs) (vector-ref vs i))))
          ((nt? e)
           (let* ((g (nt-fn e))
                  (ps (map position (vector->list (nt-leaves e)))))
             (if (equal? ps (iota (length ps)))
               g
               (let1 (ps (list->vector ps))
                     (λ (vs) (g (vector-map (λ (p) (vector-ref vs p)) ps)))))))
          ((deferred? e) (λ (vs) (force-value e)))
          (else (K e))))
  (define model? category-state?)
  (define (check-model m who)
    (unless
      (model? m)
      (hakaru10-error
        'construction
        who
        "a model (a State computation: do/monad, dist, diracN, if_ or return) is expected"
        m))
    m)
  (define-syntax-rule (modify! (ch) body ...) (modify (λ (ch) body ... ch)))
  (define (hakaru10-fail kind who msg . irritants)
    (state (λ (_) (apply hakaru10-error kind who msg irritants))))
  (define (lazy-model who thunk)
    (state (λ (ch) (run-state (check-model (thunk) who) ch))))
  (define (condition/external obs d) (condition obs d internal-part))
  (define (check-owner ch e who)
    (when (and (sexp? e) (not (eq? (sexp-chain e) ch)))
          (hakaru10-error
            'construction
            who
            "a node of another model run (a stale node) is used"
            e)))
  (define (check-distk who dk)
    (unless
      (distk? dk)
      (hakaru10-error
        'construction
        who
        "a distribution constructor must return a distribution"
        dk))
    dk)
  (define (check-observation n obs)
    (unless
      (equal? obs (node-value n))
      (hakaru10-error
        'conditioning
        'update
        "observation must be external to the model"
        n
        (node-value n)
        obs)))
  (define (inactive-scope? ch)
    (any (λ (p) (not (eq? (node-value (car p)) (cdr p)))) (chain-scope ch)))
  (define (parameter-error? e) (eq? (hakaru10-error-kind e) 'parameter))
  (define (build-guard ch thunk)
    (if (and (pair? (chain-scope ch)) (extended? ch) (inactive-scope? ch))
      (condition-case
        (thunk)
        (e (exn)
           (signal
             (if (deferrable? e)
               (make-property-condition 'hakaru10-rejected-build)
               e))))
      (thunk)))
  (define (scope-tail? s s*)
    (let tail? ((cs s*)) (or (eq? cs s) (and (pair? cs) (tail? (cdr cs))))))
  (define (check-branch-escape parent child)
    (unless
      (or (eq? (node-role child) 'exit)
          (scope-tail? (node-scope parent) (node-scope child)))
      (hakaru10-error 'branch 'if_ "node escapes its branch" parent child)))
  (define (add-dependency! parent child)
    (check-branch-escape parent child)
    (unless
      (memq child (node-deps parent))
      (node-deps-set! parent (cons child (node-deps parent)))))
  (define (check-active-update n leaves)
    (vector-for-each
      (λ (p)
         (when (node-inactive? p)
               (hakaru10-error
                 'dependency
                 'update
                 "the parent or the node itself are inactive"
                 p
                 n)))
      leaves)
    (when (node-inactive? n)
          (hakaru10-error
            'dependency
            'update
            "the parent or the node itself are inactive"
            n)))
  (define (leaves-tstamp leaves)
    (let loop ((i 0) (ts 0))
      (if (fx= i (vector-length leaves))
        ts
        (loop (fx+ i 1)
              (fxmax ts (nstate-tstamp (node-state (vector-ref leaves i))))))))
  (define (doupdate n)
    (let ((ch (node-chain n)) (self (node-state n)) (leaves (node-parents n)))
      (unless (extended? ch) (check-active-update n leaves))
      (unless
        (fx> (leaves-tstamp leaves) (nstate-tstamp self))
        (hakaru10-error
          'dependency
          'update
          "dependency violation!"
          n
          (leaves-tstamp leaves)))
      (let* ((dk (check-distk 'update ((node-fn n) (leaf-values leaves))))
             (kind (distk-kind dk)))
        (unless
          (eq? kind (node-kind n))
          (hakaru10-error
            'construction
            'update
            "a distribution constructor changed the kind of its node"
            n
            (node-kind n)
            kind))
        (values
          (if (eq? kind 'dirac)
            (make-nstate
              (check-lifted-value 'update (distk-value dk))
              0.0
              #f
              (chain-now ch))
            (let ((rd (distk-rd dk)) (v (nstate-val self)))
              (when (eq? kind 'observed)
                    (check-observation n (distk-value dk)))
              (make-nstate
                v
                ((rdistr-log-density rd) v)
                rd
                (nstate-tstamp self))))
          '()))))
  (define (new-nref n)
    (state (λ (ch)
              (chain-store! ch n)
              (when (eq? (node-kind n) 'resampleable)
                    (chain-sampleable-set! ch (cons n (chain-sampleable ch))))
              (chain-llcum-set!
                ch
                (fp+ (chain-llcum ch) (nstate-ll (node-state n))))
              (cons n ch))))
  (define (new-dependency parent child)
    (modify! (ch) (add-dependency! parent child)))
  (define (register-dependencies leaves child)
    (modify! (ch) (vector-for-each (λ (p) (add-dependency! p child)) leaves)))
  (define (create-node
           update
           dk
           #!optional
           (role 'root)
           (parents (vector))
           (fn dk))
    (state (λ (ch)
              (receive
                (v ll)
                (build-guard ch (τ (distk-sample dk (chain-rng ch))))
                (run-state
                  (new-nref
                    (make-node
                      ch
                      (fx+ 1 (chain-count ch))
                      (distk-kind dk)
                      role
                      parents
                      fn
                      update
                      (chain-scope ch)
                      '()
                      #f
                      (make-nstate v ll (distk-rd dk) (chain-now ch))
                      0
                      #f
                      #f
                      #f))
                  ch)))))
  (define (with-node-none dk) (create-node #f dk))
  (define (with-node-tree leaves fn #!optional (who 'dist))
    (state (λ (ch)
              (let1 (dk (check-distk
                          who
                          (build-guard ch (τ (fn (leaf-values leaves))))))
                    (when (eq? (distk-kind dk) 'dirac)
                          (check-no-sexp who (distk-value dk)))
                    (let* ((r (run-state
                                (create-node doupdate dk 'plain leaves fn)
                                ch))
                           (n (car r)))
                      (cons n
                            (exec-state
                              (register-dependencies leaves n)
                              (cdr r))))))))
  (define (with-node ndr f #!optional (who 'dist))
    (with-node-tree
      (vector ndr)
      (λ (vs) (f (force-value (vector-ref vs 0))))
      who))
  (define (lazy-tree leaves g df who)
    (do/monad
      (← ch get)
      (let self #f)
      (← n
         (with-node-tree
           leaves
           (λ (vs)
              (df (lazy-value
                    g
                    vs
                    (if self (node-inactive? self) (inactive-scope? ch)))))
           who))
      (let _ (set! self n))
      (return n)))
  (define (with-node-ctor e df #!optional (who 'dist) lazy)
    (check-not-model who e)
    (do/monad
      (gets (λ (ch) (check-owner ch e who)))
      (cond ((node? e)
             (if lazy
               (lazy-tree (vector e) (λ (vs) (vector-ref vs 0)) df who)
               (with-node e df who)))
            ((nt? e)
             (if lazy
               (lazy-tree (nt-leaves e) (nt-fn e) df who)
               (let1 (g (nt-fn e))
                     (with-node-tree
                       (nt-leaves e)
                       (λ (vs) (df (g (forced-values vs))))
                       who))))
            (else
             (check-no-sexp who e)
             (>>= (gets (λ (ch)
                           (check-distk
                             who
                             (build-guard ch (τ (df (force-value e)))))))
                  with-node-none)))))
  (define (to-node e #!optional (who 'to_node))
    (check-not-model who e)
    (cond ((node? e)
           (do/monad
             (gets (λ (ch)
                      (check-owner ch e who)
                      (unless
                        (scope-tail? (node-scope e) (chain-scope ch))
                        (hakaru10-error
                          'branch
                          who
                          "node escapes its branch"
                          e))))
             (return e)))
          ((nt? e) (with-node-ctor e dirac who #t))
          (else (check-no-sexp who e) (with-node-none (dirac e)))))
  (define (pair a b) (lift* cons (list a b) #f))
  (define (collect es) (lift* list es #f))
  (define-syntax-rule
    (ctor-apply applied? d args call)
    (check-distk
      'dist
      (if applied?
        call
        (let1 (dk (call/arity-check d (λ () call) args))
              (set! applied? #t)
              dk))))
  (define (call/arity-check d thunk args)
    (with-exception-handler
      (λ (e)
         (if (and (arity-condition? e) (memq d (exn-arguments e)))
           (hakaru10-error
             'construction
             'dist
             "wrong number of arguments for the distribution constructor"
             d
             args)
           (signal e)))
      thunk))
  (define arity-condition? (condition-predicate 'arity))
  (define exn-arguments (condition-property-accessor 'exn 'arguments '()))
  (define (dist d . args)
    (cond ((distk? d)
           (unless
             (null? args)
             (hakaru10-error
               'construction
               'dist
               "a distribution takes no arguments"
               d
               args)))
          ((procedure? d)
           (for-each
             (λ (a)
                (check-not-model 'dist a)
                (unless (sexp? a) (check-no-sexp 'dist a)))
             args))
          (else
           (hakaru10-error
             'construction
             'dist
             "a distribution or a distribution constructor is expected"
             d)))
    (let1 (applied? #f)
          (match/first
            args
            (()
             (if (distk? d)
               (with-node-none d)
               (with-node-ctor '() (λ (_) (ctor-apply applied? d '() (d))))))
            ((,x)
             (with-node-ctor x (λ (v) (ctor-apply applied? d (list v) (d v)))))
            ((,x ,y)
             (with-node-ctor
               (pair x y)
               (Λ ((,a unquote b))
                  (ctor-apply applied? d (list a b) (d a b)))))
            (else (with-node-ctor
                    (collect args)
                    (λ (vs) (ctor-apply applied? d vs (apply d vs))))))))
  (define (diracN e) (to-node e 'diracN))
  (define (if_* test th el)
    (check-not-model 'if_ test)
    (do/monad
      (← entry (with-node-ctor test (λ (v) (dirac (check-test 'if_ v))) 'if_))
      (modify! (ch) (node-role-set! entry 'entry))
      (ifnode
        entry
        (▷ (check-model th 'if_) to-node)
        (▷ (check-model el 'if_) to-node))))
  (define (model->trace m who) (▷ (check-model m who) to-node))
  (define-record branch entry then-result else-result then-range else-range)
  (set-record-printer!
    branch
    (λ (b port)
       (fprintf
         port
         "#<branch ~a ~a ~a>"
         (node-addr (branch-entry b))
         (node-addr (branch-then-result b))
         (node-addr (branch-else-result b)))))
  (define (branch-active-result b)
    (if (node-value (branch-entry b))
      (branch-then-result b)
      (branch-else-result b)))
  (define (check-test who v . irritants)
    (unless
      (boolean? v)
      (apply hakaru10-error
             'branch
             who
             "the test of if_ must be a boolean"
             v
             irritants))
    v)
  (define (move-activity!
           ch
           lo
           hi
           step
           switching
           observed-weight
           dirac-switched)
    (let loop ((a (fx+ lo 1)) (llc 0.0) (k 0))
      (if (fx> a hi)
        (values llc k)
        (let* ((n (chain-node ch a))
               (count (node-inactive n))
               (kind (node-kind n)))
          (when (fx< (fx+ count step) 0)
                (hakaru10-error
                  'internal
                  'activate
                  "activating an active node"
                  n))
          (set-inactive! n (fx+ count step))
          (if (fx= count switching)
            (loop (fx+ a 1)
                  (fp+ llc
                       (case kind
                         ((observed) (observed-weight n))
                         ((dirac) (dirac-switched n) 0.0)
                         (else 0.0)))
                  (if (eq? kind 'resampleable) (fx+ k 1) k))
            (loop (fx+ a 1) llc k))))))
  (define (deactivate! ch lo hi)
    (receive
      (llc k)
      (move-activity!
        ch
        lo
        hi
        1
        0
        (λ (n)
           (if (node-pending n)
             (begin (set-pending! n #f) 0.0)
             (nstate-ll (node-state n))))
        void)
      (chain-llcum-set! ch (fp- (chain-llcum ch) llc))
      (chain-card-change-set! ch (fx- (chain-card-change ch) k))))
  (define (activate! ch lo hi)
    (receive
      (llc k)
      (move-activity!
        ch
        lo
        hi
        -1
        1
        (if (extended? ch)
          (λ (n) (set-pending! n #t) 0.0)
          (λ (n) (nstate-ll (node-state n))))
        (λ (n)
           (when (deferred? (node-value n))
                 (chain-pending-set! ch (cons n (chain-pending ch))))))
      (chain-llcum-set! ch (fp+ (chain-llcum ch) llc))
      (chain-card-change-set! ch (fx+ (chain-card-change ch) k))))
  (define (settle-pending! ch)
    (let1 (pending (chain-pending ch))
          (when (pair? pending)
                (for-each
                  (λ (n)
                     (cond ((node-pending n)
                            (set-pending! n #f)
                            (chain-llcum-set!
                              ch
                              (fp+ (chain-llcum ch) (nstate-ll (node-state n)))))
                           ((and (eq? (node-kind n) 'dirac)
                                 (not (node-inactive? n)))
                            (let1 (s (node-state n))
                                  (when (deferred? (nstate-val s))
                                        (set-node-state!
                                          n
                                          (make-nstate
                                            (force-value (nstate-val s))
                                            (nstate-ll s)
                                            (nstate-dst s)
                                            (nstate-tstamp s))))))))
                  (sort! pending (λ (a b) (node<? a b))))
                (chain-pending-set! ch '()))))
  (define (switch-branches! ch s* extra newactive newinactive)
    (deactivate! ch (car newinactive) (cdr newinactive))
    (activate! ch (car newactive) (cdr newactive))
    (values
      s*
      (if (extended? ch)
        extra
        (chain-fold
          (λ (n deps)
             (if (node-delayed n)
               (begin (set-delayed! n #f) (cons n deps))
               deps))
          extra
          ch
          (car newactive)
          (cdr newactive)))))
  (define ((update-entry prev b) n)
    (let1 (old (node-value n))
          (receive
            (s* extra)
            (prev n)
            (let ((tv (nstate-val s*))
                  (ths (branch-then-range b))
                  (els (branch-else-range b)))
              (check-test 'if_ tv n)
              (cond ((eq? tv old) (values s* extra))
                    (tv (switch-branches! (node-chain n) s* extra ths els))
                    (else
                     (switch-branches! (node-chain n) s* extra els ths)))))))
  (define ((update-exit b) n)
    (let* ((self (node-state n))
           (ts (nstate-tstamp self))
           (e (node-state (branch-entry b)))
           (result (node-state (branch-active-result b))))
      (values
        (if (or (fx> (nstate-tstamp e) ts) (fx> (nstate-tstamp result) ts))
          (make-nstate
            (let1 (v (nstate-val result))
                  (if (node-inactive? n) v (force-value v)))
            (nstate-ll self)
            (nstate-dst self)
            (chain-now (node-chain n)))
          self)
        '())))
  (define (deactivate lo hi) (modify! (ch) (deactivate! ch lo hi)))
  (define (in-branch entry polarity m)
    (state (λ (ch)
              (let1 (scope (chain-scope ch))
                    (chain-scope-set! ch (cons (cons entry polarity) scope))
                    (let1 (r (if (inactive-scope? ch)
                               (parameterize
                                 ((lazy-lifting #t))
                                 (run-state m ch))
                               (run-state m ch)))
                          (chain-scope-set! (cdr r) scope)
                          r)))))
  (define (ifnode entry th el)
    (do/monad
      (when/monad
        (eq? (node-kind entry) 'resampleable)
        (hakaru10-fail
          'internal
          'if_
          "Entry node for if should not be sampleable"
          entry))
      (let testv (check-test 'if_ (node-value entry) entry))
      (let (,active ,inactive) (if testv (list th el) (list el th)))
      (← added-old (gets chain-count))
      (← inactive-nr (in-branch entry (not testv) inactive))
      (← added-inactive (gets chain-count))
      (deactivate added-old added-inactive)
      (← active-nr (in-branch entry testv active))
      (← added-active (gets chain-count))
      (let b (if
              testv
              (make-branch
                entry
                active-nr
                inactive-nr
                (cons added-inactive added-active)
                (cons added-old added-inactive))
              (make-branch
                entry
                inactive-nr
                active-nr
                (cons added-old added-inactive)
                (cons added-inactive added-active))))
      (← exit
         (create-node
           (update-exit b)
           (dirac (node-value active-nr))
           'exit
           (list->vector
             (delete-duplicates (list entry active-nr inactive-nr) eq?))
           #f))
      (modify! (ch) (node-branch-set! exit b))
      (new-dependency entry exit)
      (new-dependency active-nr exit)
      (new-dependency inactive-nr exit)
      (modify!
        (ch)
        (let1 (prev (node-update entry))
              (when prev (node-update-set! entry (update-entry prev b))))
        (node-branch-set! entry b))
      (return exit)))
  (define (modify-node! n s*)
    (let ((ch (node-chain n)) (old (node-state n)))
      (set-node-state! n s*)
      (unless
        (and (extended? ch)
             (eq? (node-kind n) 'observed)
             (or (node-inactive? n) (node-pending n)))
        (chain-llcum-set!
          ch
          (fp- (fp+ (chain-llcum ch) (nstate-ll s*)) (nstate-ll old))))
      (if (fx= (nstate-tstamp old) (nstate-tstamp s*)) '() (node-deps n))))
  (define (guarded-update upd n)
    (condition-case
      (upd n)
      (e (exn) (if (deferrable? e) (values #f '()) (signal e)))))
  (define (abort-propagation! ch)
    (chain-llcum-set! ch -inf.0)
    (chain-pending-set! ch '())
    (update-queue-clear! (chain-queue ch)))
  (define (propagate! ch)
    (let ((q (chain-queue ch)) (extended (extended? ch)) (now (chain-now ch)))
      (let loop ()
        (if (update-queue-empty? q)
          (settle-pending! ch)
          (let1 (n (dequeue! q))
                (cond ((and (not extended) (node-inactive? n))
                       (set-delayed! n #t)
                       (loop))
                      ((fx>= (nstate-tstamp (node-state n)) now)
                       (hakaru10-error
                         'dependency
                         'propagate
                         "dependency violation: updating the already updated"
                         n))
                      ((node-update n)
                       =>
                       (λ (upd)
                          (receive
                            (s* extra)
                            (if (and extended (node-inactive? n))
                              (guarded-update upd n)
                              (upd n))
                            (if s*
                              (begin
                                (enqueue-all! ch extra)
                                (enqueue-all! ch (modify-node! n s*))
                                (loop))
                              (abort-propagation! ch)))))
                      (else (loop))))))))
  (define hakaru10-debug (make-parameter #f))
  (define-syntax-rule
    (debug-trace arg ...)
    (when (hakaru10-debug)
          (with-output-to-port (current-error-port) (τ (print arg ...)))))
  (define (apply-proposal! n v ll fb)
    (let ((ch (node-chain n)) (s (node-state n)))
      (enqueue-all!
        ch
        (modify-node! n (make-nstate v ll (nstate-dst s) (chain-now ch))))
      fb))
  (define (resample-trace! ch)
    (let* ((samp (chain-samp ch))
           (hi (fx- (vector-length samp) 1))
           (src (chain-rng ch)))
      (let retry ()
        (let1 (n (vector-ref samp (rng-integer! src 0 hi)))
              (if (node-inactive? n)
                (retry)
                (let1 (s (node-state n))
                      (receive
                        (v ll fb)
                        ((rdistr-resample (nstate-dst s))
                         src
                         (nstate-val s)
                         (nstate-ll s))
                        (debug-trace
                          "resample "
                          (node-addr n)
                          ": "
                          (nstate-val s)
                          " -> "
                          v
                          " fb "
                          fb)
                        (apply-proposal! n v ll fb))))))))
  (define (log-acceptance-ratio ch sampleables card-old fb)
    (let1 (card-new (chain-card-change ch))
          (fp+ (fp- (chain-llcum ch) fb)
               (if (fx= card-new card-old)
                 0.0
                 (fplog (fp/ (exact->inexact (fx+ sampleables card-old))
                             (exact->inexact (fx+ sampleables card-new))))))))
  (define (rejected-ll? x) (and (not (nan? x)) (infinite? x)))
  (define (run-step-hook ch accepted? delta)
    (let1 (hook (chain-step-hook ch)) (when hook (hook ch accepted? delta))))
  (define (transit! ch)
    (journal-begin! ch)
    (chain-now-set! ch (fx+ (chain-now ch) 1))
    (chain-llcum-set! ch 0.0)
    (let ((card-old (chain-card-change ch))
          (sampleables (vector-length (chain-samp ch))))
      (if (fx= 0 (fx+ sampleables card-old))
        (begin (journal-commit! ch) (run-step-hook ch #t 0.0) #t)
        (let* ((fb (resample-trace! ch))
               (delta (begin (propagate! ch) (chain-llcum ch)))
               (ratio (and (not (rejected-ll? delta))
                           (log-acceptance-ratio ch sampleables card-old fb)))
               (accept?
                 (and ratio
                      (or (fp>= ratio 0.0)
                          (fp< (rng-double! (chain-rng ch) 0.0 1.0)
                               (fpexp ratio))))))
          (debug-trace
            "step "
            (chain-now ch)
            ": delta "
            delta
            " ratio "
            ratio
            (if accept? " accepted" " rejected"))
          (if accept? (journal-commit! ch) (journal-rollback! ch))
          (run-step-hook ch accept? delta)
          accept?))))
  (define (chain-extended-ll ch)
    (chain-fold
      (λ (n acc)
         (case (node-kind n)
           ((resampleable) (fp+ acc (nstate-ll (node-state n))))
           ((observed)
            (if (node-inactive? n) acc (fp+ acc (nstate-ll (node-state n)))))
           (else acc)))
      0.0
      ch))
  (define (build-chain prog src semantics)
    (check-model prog 'mcmC)
    (let* ((ch (new-chain src semantics))
           (result
             (condition-case
               (parameterize ((lazy-lifting #f)) (eval-state prog ch))
               (e (hakaru10-rejected-build) #f))))
      (chain-freeze! ch)
      (cond ((not result) (chain-llcum-set! ch -inf.0))
            ((extended? ch) (chain-llcum-set! ch (chain-extended-ll ch))))
      (values ch result)))
  (define (initial-step limit prog src semantics)
    (let retry ((limit limit))
      (if (<= limit 0)
        (values #f #f 0)
        (receive
          (ch result)
          (build-chain prog src semantics)
          (if (rejected-ll? (chain-llcum ch))
            (retry (- limit 1))
            (values ch result limit))))))
  (define (mcmc-generator
           limit
           prog
           #!key
           (seed 17)
           (rng #f)
           (semantics 'upstream)
           (step-hook #f)
           (build-hook #f))
    (unless
      (and (exact-integer? limit) (>= limit 0))
      (hakaru10-error
        'construction
        'mcmC
        "limit: must be a non-negative integer"
        limit))
    (receive
      (ch result limit*)
      (initial-step limit prog (or rng (make-stdgen-source seed)) semantics)
      (and ch
           (let1 (started #f)
                 (chain-step-hook-set! ch step-hook)
                 (when build-hook (build-hook ch))
                 (τ (if started (transit! ch) (set! started #t))
                    (force-value (node-value result)))))))
  (define (mcmc-fold kons knil limit prog . keys)
    (let1 (next (apply mcmc-generator limit prog keys))
          (if next
            (let loop ((i 0) (acc knil))
              (if (>= i limit) acc (loop (+ i 1) (kons (next) acc))))
            knil))))

(module
  (aux hakaru10)
  (dist diracN
        condition
        if_
        lift
        pair
        collect
        sexp-map
        sexp+
        sexp-
        sexp*
        sexp/
        return
        >>=
        >>
        ▷
        do/monad
        <$>
        <*>
        fold/monad
        traverse/monad
        sequence/monad
        for-each/monad
        replicate/monad
        bern
        bern/flip
        categorical
        uniformly
        uniform
        normal
        gamma
        beta
        poisson
        laplace
        dirichlet/symmetric
        make-distribution
        distribution-log-density
        distribution-sample
        mcmC
        mcmC-fold
        mcmC§
        model-graph
        exact-posterior
        model->search-tree
        likelihood-weighting
        weighted-samples
        make-stdgen-source
        make-mrg32k3a-source
        make-chicken-random-source
        hist
        frequency
        momenta
        tabulate
        tabulate/bins
        kl
        plot
        hakaru10-error?
        hakaru10-error-kind)
  (import
    scheme
    (only (scheme base) exact-integer?)
    (chicken base)
    (only srfi-1 every)
    (aux base)
    (only (aux stream) τ->§)
    (only (aux category state) return >>=)
    (only (aux category monad state)
          ▷
          do/monad
          >>
          <$>
          <*>
          fold/monad
          traverse/monad
          sequence/monad
          for-each/monad
          replicate/monad)
    (only (aux hakaru10 random)
          make-stdgen-source
          make-mrg32k3a-source
          make-chicken-random-source
          rng-source-copy)
    (only (aux hakaru10 distribution)
          hakaru10-error
          hakaru10-error?
          hakaru10-error-kind
          bern
          bern/flip
          categorical
          uniformly
          uniform
          normal
          gamma
          beta
          poisson
          laplace
          dirichlet/symmetric
          make-distribution
          distribution-log-density
          distribution-sample)
    (only (aux hakaru10 util)
          hist
          frequency
          momenta
          tabulate
          tabulate/bins
          kl
          plot)
    (only (aux hakaru10 engine)
          dist
          diracN
          if_*
          lazy-model
          model->trace
          pair
          collect
          lift*
          condition/external
          build-chain
          mcmc-generator
          mcmc-fold
          extended?
          chain-count
          chain-fold
          chain-now
          chain-card-change
          chain-extended-ll
          chain-step-hook-set!
          node-kind
          node-role
          node-state
          node-inactive
          node-inactive?
          node-delayed
          node-pending)
    (only (aux hakaru10 exact)
          exact-posterior
          model->search-tree
          likelihood-weighting
          weighted-samples
          construction-check
          reference-check))
  (define-syntax-rule
    (if_ test then else)
    (if_* test (lazy-model 'if_ (τ then)) (lazy-model 'if_ (τ else))))
  (define condition condition/external)
  (define (lift f . es) (lift* f es))
  (define sexp-map lift)
  (define (sexp+ . es) (lift* + es))
  (define (sexp- . es) (lift* - es))
  (define (sexp* . es) (lift* * es))
  (define (sexp/ . es) (lift* / es))
  (define-record step-snapshot now card-change ll nodes)
  (define (take-step-snapshot ch)
    (make-step-snapshot
      (chain-now ch)
      (chain-card-change ch)
      (chain-extended-ll ch)
      (chain-fold
        (λ (n acc)
           (cons (list (node-state n)
                       (node-inactive n)
                       (node-delayed n)
                       (node-pending n))
                 acc))
        '()
        ch)))
  (define (same-state? a b)
    (and (= (step-snapshot-now a) (step-snapshot-now b))
         (= (step-snapshot-card-change a) (step-snapshot-card-change b))
         (every (λ (m n) (every eq? m n))
                (step-snapshot-nodes a)
                (step-snapshot-nodes b))))
  (define (check-step before ch accepted? delta)
    (let1 (after (take-step-snapshot ch))
          (cond ((not accepted?)
                 (unless
                   (same-state? before after)
                   (hakaru10-error
                     'internal
                     'check
                     "a rejected step did not restore the state"
                     (step-snapshot-now before))))
                ((and (extended? ch)
                      (finite? (step-snapshot-ll before))
                      (finite? (step-snapshot-ll after)))
                 (let* ((ll0 (step-snapshot-ll before))
                        (expected (- (step-snapshot-ll after) ll0)))
                   (unless
                     (<= (abs (- delta expected))
                         (* 1e-09 (max 1.0 (abs ll0))))
                     (hakaru10-error
                       'internal
                       'check
                       "the step's delta is not the extended LL change"
                       (chain-now ch)
                       delta
                       expected)))))))
  (define (check-hooks check)
    (let ((before #f) (steps 0))
      (define (more?) (or (eq? check #t) (< steps check)))
      (define (verify! ch)
        (reference-check ch)
        (if (more?)
          (set! before (take-step-snapshot ch))
          (chain-step-hook-set! ch #f)))
      (values
        verify!
        (λ (ch accepted? delta)
           (check-step before ch accepted? delta)
           (add1! steps)
           (verify! ch)))))
  (define (with-chain-keys driver model seed rng semantics check)
    (unless
      (or (boolean? check) (and (exact-integer? check) (>= check 0)))
      (hakaru10-error
        'construction
        'mcmC
        "check: must be a boolean or a step count"
        check))
    (let1 (src (or rng (make-stdgen-source seed)))
          (if check
            (begin
              (construction-check
                model
                #:rng
                (rng-source-copy src)
                #:semantics
                semantics)
              (receive
                (build-hook step-hook)
                (check-hooks check)
                (driver
                  #:rng
                  src
                  #:semantics
                  semantics
                  #:build-hook
                  build-hook
                  #:step-hook
                  step-hook)))
            (driver #:rng src #:semantics semantics))))
  (define (mcmC-fold
           kons
           knil
           limit
           model
           #!key
           (seed 17)
           (rng #f)
           (semantics 'upstream)
           (check #f))
    (let1 (prog (model->trace model 'mcmC))
          (with-chain-keys
            (λ keys (apply mcmc-fold kons knil limit prog keys))
            model
            seed
            rng
            semantics
            check)))
  (define (mcmC limit model . keys)
    (reverse (apply mcmC-fold cons '() limit model keys)))
  (define (mcmC§
           model
           #!key
           (seed 17)
           (rng #f)
           (semantics 'upstream)
           (check #f)
           (retries 1000))
    (let* ((prog (model->trace model 'mcmC))
           (next (with-chain-keys
                   (λ keys (apply mcmc-generator retries prog keys))
                   model
                   seed
                   rng
                   semantics
                   check)))
      (if next (τ->§ next) '())))
  (define (model-graph model #!key (seed 17) (rng #f) (semantics 'upstream))
    (receive
      (ch result)
      (build-chain
        (model->trace model 'model-graph)
        (or rng (make-stdgen-source seed))
        semantics)
      (define (count-nodes node?)
        (chain-fold (λ (n k) (if (node? n) (add1 k) k)) 0 ch))
      (define ((kind? kind) n) (eq? (node-kind n) kind))
      `((nodes unquote (chain-count ch))
        (resampleable unquote (count-nodes (kind? 'resampleable)))
        (observed unquote (count-nodes (kind? 'observed)))
        (dirac unquote (count-nodes (kind? 'dirac)))
        (inactive unquote (count-nodes node-inactive?))
        (entries unquote (count-nodes (λ (n) (eq? (node-role n) 'entry))))))))

(define (test/language/graph/grass _)
  (⊦ equal?
     '((nodes . 3)
       (resampleable . 2)
       (observed . 1)
       (dirac . 0)
       (inactive . 0)
       (entries . 0))
     (model-graph grass)))

((eta 0.0) (memory #(12582912 6198944 1048576)) (stdout "") (stderr ""))

(define (normals means std) (traverse/monad (λ (m) (dist normal m std)) means))

(define (dot xs ays)
  (reduce-right sexp+ #f (map (λ (x ay) (lift (λ (a) (* a x)) ay)) xs ays)))

(define (regression xsy)
  (do/monad
    (let mu (make-list (length (caar xsy)) 0.0))
    (← w-mean (normals mu 2.0))
    (← w (normals w-mean 1.0))
    (← noise-sd (<$> (cut lift (λ (g) (/ 1.0 g)) <>) (dist gamma 0.5 0.5)))
    (traverse/monad
      (Λ ((,xs unquote y)) (dist (condition y normal) (dot xs w) noise-sd))
      xsy)
    (return (collect w))))

-- normals means std
-- returns several independently distributed normal random variables
-- whose means are given by the list means, with the same std
normals :: [SExp Double] -> SExp Double -> MCMCM [SExp Double]
normals means std = mapM (\m -> dist normal m std) means


dot :: [Double] -> [SExp Double] -> SExp Double
dot xs ays = foldr1 (+) $ zipWith (\x ay -> (*x) <$> ay) xs ays

collect :: [SExp a] -> SExp [a]
collect = sequenceA

type RegrDatum = ([Double],Double) -- Xs and Y
model :: [RegrDatum] -> Model [Double]
model xsy = do
  let mu = replicate dimK 0
  w_mean   <- normals mu 2
  w        <- normals w_mean 1
  noise_sd <- (1/) <$> dist gamma 0.5 0.5
  let make_cond (xs,y) = dist (y `condition` normal) (dot xs w) noise_sd
  mapM make_cond xsy
  return $ collect w

(define (test/language/graph/regression _)
  (⊦ equal?
     '((nodes . 512)
       (resampleable . 11)
       (observed . 500)
       (dirac . 1)
       (inactive . 0)
       (entries . 0))
     (model-graph (regression (regression-data 5 500)))))

((eta 0.02) (memory #(12582912 7805176 1048576)) (stdout "") (stderr ""))

(define (opposite-color color) (if (eq? color 'blue) 'green 'blue))

(define (observed-color color)
  (categorical `((,color . 0.8) (,(opposite-color color) . 0.2))))

(define max-balls 8)

(define (balls-prior n)
  (do/monad
    (← balls (replicate/monad n (dist uniformly '(blue green))))
    (return (map cons (iota n 1) balls))))

(define (balls-observe balls obs nballs)
  (>> (for-each/monad
        (λ (color)
           (do/monad
             (← b (dist uniformly (iota nballs 1)))
             (fold-right
               (λ (i rest)
                  (if_ (lift (λ (bv) (= bv i)) b)
                       (dist (condition color observed-color)
                             (alist-ref i balls))
                       rest))
               (return color)
               (iota max-balls 1))))
        obs)
      (return '())))

(define (cballs-model obs)
  (do/monad
    (← balls (balls-prior max-balls))
    (← nballs (dist uniformly (iota max-balls 1)))
    (fold-right
      (λ (i rest)
         (if_ (lift (λ (n) (= n i)) nballs) (balls-observe balls obs i) rest))
      (return '())
      (iota max-balls 1))
    (return nballs)))

-- Submodel: Given balls color and the observations, perform conditioning
balls_observe :: M.IntMap (SExp Color) -> [Color] -> Int -> Model ()
balls_observe balls obs nballs = mapM_ observe_ball obs >> return (pure ())
 where
  -- The same ball can be drawn twice as we return the drawn balls
  -- into the urn. A ball keeps its true color, no matter how many
  -- times it is drawn.
   observe_ball color = do
        b <- dist uniformly (pure [1..nballs])
        -- It is more efficient to create if statements, rather than
        -- create an SExp(IntMap Color).
        let observe i =
              let true_color = balls ! i in
              if_ ((== i) <$> b) $
                dist (color `condition` observed_color) true_color
        foldr observe (return (pure color)) [1..maxBalls]

-- The model takes a list of observations as an argument
cballs_model :: [Color] -> Model Int
cballs_model obs = do
  balls  <- balls_prior maxBalls
  nballs <- dist uniformly (pure [1..maxBalls])
  let obs_number i =
        if_ ((== i) <$> nballs) $ balls_observe balls obs i
  foldr obs_number (return (pure())) [1..maxBalls]
  return nballs

(define (test/language/graph/colored-balls _)
  (let1 (graph (model-graph (cballs-model (make-list 10 'blue))))
        (⊦ equal?
           '(2114 89 1994)
           (map (λ (key) (alist-ref key graph))
                '(nodes resampleable inactive)))))

((eta 0.01) (memory #(12582912 7894904 1048576)) (stdout "") (stderr ""))

(define ordered-pair
  (<*> (<$> (λ (a) (cut pair a <>)) (dist normal 0.0 1.0))
       (dist normal 5.0 1.0)))

(define ordered-pair*
  (do/monad
    (← a (dist normal 0.0 1.0))
    (← b (dist normal 5.0 1.0))
    (return (pair a b))))

(define swapped-pair
  (do/monad
    (← b (dist normal 5.0 1.0))
    (← a (dist normal 0.0 1.0))
    (return (pair a b))))

(define (test/language/applicative _)
  (⊦ equal? (mcmC 50 ordered-pair*) (mcmC 50 ordered-pair))
  (⊦≠ (mcmC 50 swapped-pair) (mcmC 50 ordered-pair))
  (⊦ equal?
     '(-3.0 0.25 2 0.5 4 #f)
     (list (sexp- 3.0)
           (sexp/ 4.0)
           (sexp- 3 1)
           (sexp/ 1.0 2.0)
           (sexp-map add1 3)
           (sexp? 3)))
  (⊦ equal?
     '((-3.0) (0.25))
     (list (mcmC 1 (return (sexp- 3.0))) (mcmC 1 (return (sexp/ 4.0)))))
  (⊦ = 4 (sexp-map (λ (f v) (f v)) add1 3))
  (for-each
    (Λ ((,x ,x? ,nx? ,neg ,inv ,diff ,quot ,mapped ,applied))
       (⊦ equal?
          (list #t #t (- x) (/ x) (- 10.0 x) (/ x 2.0) (- x) (- x))
          (list x? nx? neg inv diff quot mapped applied)))
    (mcmC 20
          (do/monad
            (← x (dist normal 0.0 1.0))
            (return
              (lift list
                    x
                    (sexp? x)
                    (sexp? (sexp- x))
                    (sexp- x)
                    (sexp/ x)
                    (sexp- 10.0 x)
                    (sexp/ x 2.0)
                    (sexp-map - x)
                    (sexp-map (λ (f v) (f v)) - x)))))))

((eta 0.001) (memory #(12582912 6201232 1048576)) (stdout "") (stderr ""))

-- NB: if x is just a NoderRef, (diracN x) returns this reference (shares)
-- rather than creating a copy of it!
diracN :: Show a => SExp a -> Model a
diracN (Val x)  = fmap NRef $ create_node Nothing (diracS x)
diracN x@NRef{} = return x
diracN x = fmap NRef $ with_node_ctor x diracS

(define (test/language/sharing _)
  (⊦ equal?
     '(#t)
     (mcmC 1
           (do/monad
             (← x (dist normal 0.0 1.0))
             (← y (diracN x))
             (return (eq? x y)))))
  (⊦ equal?
     '((nodes . 1)
       (resampleable . 0)
       (observed . 0)
       (dirac . 1)
       (inactive . 0)
       (entries . 0))
     (model-graph (diracN 3.0)))
  (⊦ =
     2
     (alist-ref
       'nodes
       (model-graph
         (do/monad (← x (dist normal 0.0 1.0)) (diracN (sexp+ x x))))))
  (receive
    (ch result)
    (build-chain
      (model->trace
        (do/monad (← x (dist normal 0.0 1.0)) (diracN (lift + x x x)))
        'test)
      (make-stdgen-source 17)
      'upstream)
    (⊦ = 1 (vector-length (node-parents result)))))

((eta 0.0) (memory #(12582912 6205384 1048576)) (stdout "") (stderr ""))

{- Type error
biased_coinX c1 c2 = do
  b <- dist beta 1 1
  dist (c1 `condition`  (c2 `condition` bern)) b
  dist (c2 `condition` bern) b
  return b
-}

{-
biased_coin_ill_typed c1 c2 = do
  b <- dist bern 0.5
  dist (c1 `condition` bern) b
  dist (c2 `condition` bern) b
  return b
-}

{-
biased_coin_ill_typed_too c1 c2 = do
  b <- dist bern 0.5
  dist (b  `condition` bern) 0.4
  return b
-}

(define (test/language/errors _)
  (define captured #f)
  (define leaked/sibling
    (do/monad
      (← c (dist bern 0.5))
      (let leak #f)
      (if_ c
           (do/monad (← h (dist bern 0.9)) (let _ (set! leak h)) (return h))
           (return #f))
      (← y (if_ #t (return leak) (return #f)))
      (return (pair c y))))
  (define leaked/result
    (do/monad
      (← c (dist bern 0.5))
      (let leak #f)
      (if_ c
           (do/monad
             (← h (dist normal 10.0 1.0))
             (let _ (set! leak h))
             (return h))
           (return 0.0))
      (return leak)))
  (define cases
    `((condition-on-a-random-choice
        conditioning
        ,(τ (mcmC 3
                  (do/monad
                    (← b (dist bern 0.5))
                    (dist (condition b bern) 0.4)
                    (return b)))))
      (condition-on-a-model
        conditioning
        ,(τ (mcmC 3
                  (do/monad
                    (← b (dist bern 0.5))
                    (dist (condition (dist bern 0.5) bern) 0.5)
                    (return b)))))
      (condition-on-an-sexp-inside-a-value
        conditioning
        ,(τ (mcmC 3
                  (do/monad
                    (← c (dist bern 0.5))
                    (dist (condition (list c) uniformly) '((#t) (#f)))
                    (return c)))))
      (condition-twice
        conditioning
        ,(τ (condition #t (condition #t (bern 0.5)))))
      (condition-twice-applied
        conditioning
        ,(τ (mcmC 3
                  (do/monad
                    (← b (dist beta 1.0 1.0))
                    (dist (condition #t (condition #f bern)) b)
                    (return b)))))
      (condition-on-a-constructor-argument
        conditioning
        ,(τ (mcmC 50
                  (do/monad
                    (← x (dist bern 0.5))
                    (dist (λ (xv) (condition xv (bern 0.8))) x)
                    (return x)))))
      (dist-in-a-lifted-function
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← x (dist normal 0.0 1.0))
                    (return (lift (λ (xv) (dist normal xv 1.0)) x))))))
      (model-from-a-lifted-function-at-update
        construction
        ,(τ (mcmC 200
                  (do/monad
                    (← x (dist bern 0.5))
                    (diracN (lift (λ (xv) (if xv 0.5 (dist bern 0.5))) x))))))
      (sexp-from-a-lifted-function-at-update
        construction
        ,(τ (mcmC 200
                  (do/monad
                    (← x (dist bern 0.5))
                    (diracN (lift (λ (xv) (if xv 0.5 (sexp+ x 1))) x))))))
      (model-from-a-lifted-function-as-a-parameter
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← c (dist bern 0.5))
                    (dist normal (lift (λ (cv) pnorm) c) 1.0)))))
      (model-inside-a-lifted-value
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← c (dist bern 0.5))
                    (dist uniformly (lift (λ (cv) (list pnorm 1)) c))))))
      (a-model-as-a-lift-argument
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← x (dist bern 0.5))
                    (return (lift (λ (xv m) xv) x pnorm))))))
      (a-model-as-the-only-lift-argument
        construction
        ,(τ (lift (λ (m) 1) pnorm)))
      (a-model-as-a-dist-argument
        construction
        ,(τ (mcmC 3 (dist normal pnorm 1.0))))
      (a-model-as-the-result construction ,(τ (mcmC 3 (return pnorm))))
      (a-model-inside-a-value construction ,(τ (mcmC 3 (diracN (list pnorm)))))
      (not-a-model construction ,(τ (mcmC 3 42)))
      (not-a-model/exact construction ,(τ (exact-posterior 42)))
      (a-thunk-is-not-a-model construction ,(τ (mcmC 3 (λ () 1))))
      (a-branch-that-is-not-a-model construction ,(τ (mcmC 3 (if_ #t 1 2))))
      (a-container-of-sexps-as-the-result
        construction
        ,(τ (mcmC 3 (do/monad (← x (dist bern 0.5)) (return (list x))))))
      (sexps-nested-in-the-result
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← x (dist bern 0.5))
                    (← y (dist bern 0.5))
                    (return (cons 1 (vector x y)))))))
      (an-sexp-inside-a-diracN-value
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← x (dist bern 0.5))
                    (diracN (vector 1 (list x)))))))
      (an-sexp-inside-a-dist-argument
        construction
        ,(τ (mcmC 3
                  (do/monad (← x (dist bern 0.5)) (dist uniformly (list x))))))
      (an-sexp-inside-a-dist-argument-next-to-an-sexp
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← x (dist bern 0.5))
                    (← s (dist bern 0.5))
                    (dist (λ (xs sv) (uniformly xs)) (list x) s)))))
      (a-model-inside-a-dist-argument-next-to-sexps
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← s (dist bern 0.5))
                    (dist (λ (ms sv mu) (normal mu 1.0))
                          (list pnorm)
                          s
                          0.0)))))
      (ctor-not-a-distribution
        construction
        ,(τ (mcmC 3
                  (do/monad (← x (dist normal 0.0 1.0)) (dist (λ (m) 42) x)))))
      (a-distribution-with-arguments construction ,(τ (dist (bern 0.5) 1)))
      (not-a-distribution construction ,(τ (dist 42)))
      (negative-limit construction ,(τ (mcmC -1 (dist bern 0.5))))
      (negative-limit/fold
        construction
        ,(τ (mcmC-fold cons '() -1 (dist bern 0.5))))
      (negative-limit/stream
        construction
        ,(τ (mcmC§ (dist bern 0.5) #:retries -1)))
      (fractional-limit construction ,(τ (mcmC 2.5 (dist bern 0.4))))
      (rational-limit construction ,(τ (mcmC 1/2 (dist bern 0.4))))
      (inexact-limit/fold
        construction
        ,(τ (mcmC-fold cons '() 3.0 (dist bern 0.5))))
      (infinite-limit/stream
        construction
        ,(τ (mcmC§ (dist bern 0.5) #:retries +inf.0)))
      (check-not-a-boolean-or-a-count
        construction
        ,(τ (mcmC 3 (dist bern 0.5) #:check 'yes)))
      (check-a-negative-count
        construction
        ,(τ (mcmC 3 (dist bern 0.5) #:check -1)))
      (non-boolean-test branch ,(τ (mcmC 3 (if_ 1 (return 1) (return 2)))))
      (non-boolean-test-at-update
        branch
        ,(τ (mcmC 200
                  (do/monad
                    (← x
                       (dist categorical '((#t . 0.5) (#f . 0.25) (0 . 0.25))))
                    (if_ (lift identity x) (return 1) (return 2))))))
      (a-node-of-a-previous-chain
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← x (dist normal 0.0 1.0))
                    (let _ (set! captured x))
                    (return x)))
            (mcmC 3 (dist normal captured 1.0))))
      (a-node-of-a-previous-chain-lifted-with-a-fresh-one
        construction
        ,(τ (mcmC 3
                  (do/monad
                    (← x (dist normal 0.0 1.0))
                    (let _ (set! captured x))
                    (return x)))
            (mcmC 3
                  (do/monad
                    (← x (dist normal 0.0 1.0))
                    (dist normal (sexp+ x captured) 1.0)))))
      (branch-escape
        branch
        ,(τ (mcmC 3
                  (do/monad
                    (← c (dist bern 0.5))
                    (let leak #f)
                    (if_ c
                         (do/monad
                           (← l (dist normal 0.0 1.0))
                           (let _ (set! leak l))
                           (return l))
                         (return 0.0))
                    (dist normal leak 1.0)))))
      (branch-escape/as-the-result-of-a-sibling-if_
        branch
        ,(τ (mcmC 3 leaked/sibling)))
      (branch-escape/as-the-result-of-a-sibling-if_/exact
        branch
        ,(τ (exact-posterior leaked/sibling)))
      (branch-escape/as-the-model-result branch ,(τ (mcmC 3 leaked/result)))
      (kind-change
        construction
        ,(τ (mcmC 200
                  (do/monad
                    (← x (dist bern 0.5))
                    (dist (λ (xv) (if xv (bern 0.3) (condition #t (bern 0.3))))
                          x)))))
      (sampleable-entry
        internal
        ,(τ (eval-state
              (do/monad
                (← c (dist bern 0.5))
                (ifnode c (to-node 1) (to-node 2)))
              (new-chain (make-stdgen-source 17) 'upstream))))))
  (for-each
    (Λ ((,name ,kind ,thunk))
       (⊦ equal? (list name kind) (list name (hakaru10-kind thunk))))
    cases)
  (mcmC 3
        (do/monad (← x (dist bern 0.5)) (let _ (set! captured x)) (return x)))
  (⊦ equal?
     '((construction if_) (construction if_) (construction if_))
     (list (raised/location (τ (mcmC 3 (if_ captured (return 1) (return 2)))))
           (raised/location
             (τ (mcmC 3 (if_ (list (return 1)) (return 1) (return 2)))))
           (raised/location
             (τ (mcmC 3
                      (do/monad
                        (← c (dist bern 0.5))
                        (if_ (list c) (return 1) (return 2))))))))
  (let1 (inactive-typo
          (do/monad
            (← c (dist bern 0.5))
            (if_ #f (dist normal 0.0) (return 1))
            (return c)))
        (⊦ equal?
           (make-list 6 '(construction dist))
           (list (raised/location (τ (mcmC 3 (dist normal 0.0))))
                 (raised/location (τ (mcmC 3 (dist (λ (m) (normal m 1.0))))))
                 (raised/location
                   (τ (mcmC 3 (dist (λ (m s) (normal m s)) 0.0 1.0 2.0))))
                 (raised/location (τ (mcmC 3 inactive-typo)))
                 (raised/location
                   (τ (mcmC 3 inactive-typo #:semantics 'extended)))
                 (raised/location (τ (exact-posterior inactive-typo))))))
  (⊦ eq? 'exn (raised (τ (mcmC 3 (dist (λ (m) (normal m)) 0.0)))))
  (let1 (fresh-list
          (do/monad
            (← x (dist bern 0.5))
            (dist (λ (xv)
                     (condition
                       (list 1)
                       (uniformly (if xv '((1) (2)) '((1) (3) (4))))))
                  x)
            (return x)))
        (⊦ (equal-approx? 1e-12)
           '((#f . 0.4) (#t . 0.6))
           (exact-posterior fresh-list))
        (⊦ equal? (mcmC 500 fresh-list) (mcmC 500 fresh-list #:check #t)))
  (let1 (fresh-flonum
          (do/monad
            (← x (dist normal 0.0 1.0))
            (dist (λ (xv) (condition (exact->inexact 1) (normal xv 1.0))) x)
            (return x)))
        (⊦ equal? (mcmC 500 fresh-flonum) (mcmC 500 fresh-flonum #:check #t)))
  (let1 (k 0)
        (⊦ eq?
           'exn
           (raised
             (τ (mcmC 3
                      (do/monad
                        (← x (dist bern 0.5))
                        (add1! k)
                        (return x)))))))
  (⊦ equal?
     (mcmC 50
           (do/monad
             (← a (dist normal 0.0 1.0))
             (dist (λ (m s) (normal (- m 10.0) s)) a 1.0)))
     (mcmC 50
           (do/monad
             (← a (dist normal 0.0 1.0))
             (dist (λ (m s k) (normal (- m k) s)) a 1.0 10.0)))))

((eta 0.012) (memory #(12582912 6222856 1048576)) (stdout "") (stderr ""))

(define (test/language/constant-chain _)
  (⊦ equal? '(3 3 3 3 3) (mcmC 5 (return 3))))

((eta 0.0) (memory #(12582912 6224432 1048576)) (stdout "") (stderr ""))

(define (test/language/no-active-resampleable _)
  (⊦ equal?
     (make-list 10 3)
     (mcmC 10 (if_ #f (dist normal 0.0 1.0) (return 3)))))

((eta 0.0) (memory #(12582912 6226008 1048576)) (stdout "") (stderr ""))

(define (test/language/stream _)
  (let1 (s (mcmC§ pcond4))
        (⊦ equal? (mcmC 100 pcond4) (§->list (take§ 100 s)))
        (⊦ equal? (cadr (mcmC 2 pcond4)) (car§ (cdr§ s))))
  (⊦ equal? '() (mcmC§ (dist (condition 5.0 uniform) 0.0 1.0)))
  (let1 (rare (do/monad
                (← c (dist bern 0.1))
                (dist (condition #t bern) (lift (λ (cv) (if cv 1.0 0.0)) c))
                (return c)))
        (⊦ equal? '() (mcmC 6 rare))
        (⊦ equal? (make-list 7 #t) (mcmC 7 rare))
        (⊦ equal? '() (mcmC§ rare #:retries 6))
        (⊦ equal?
           (make-list 3 #t)
           (§->list (take§ 3 (mcmC§ rare #:retries 7))))))

((eta 0.004) (memory #(12582912 6219544 1048576)) (stdout "") (stderr ""))

(define (test/language/fold _)
  (⊦ =
     (count-true (mcmC 2000 grass))
     (mcmC-fold (λ (rain n) (if rain (add1 n) n)) 0 2000 grass))
  (⊦ equal? (mcmC 100 pcond1 #:rng (make-stdgen-source 17)) (mcmC 100 pcond1))
  (⊦≠ (mcmC 100 pcond1 #:seed 18) (mcmC 100 pcond1)))

((eta 0.015) (memory #(12582912 6224328 1048576)) (stdout "") (stderr ""))

(define (test/language/three-interpretations _)
  (let1 (exact (exact-posterior grass))
        (⊦ (equal-approx? 1e-12)
           '((#f . 0.53152855727963) (#t . 0.46847144272037))
           exact)
        (⊦ (equal-approx? 0.02) exact (likelihood-weighting 20000 grass))
        (⊦ (equal-approx? 0.04) exact (proportions (mcmC 20000 grass)))))

((eta 0.174) (memory #(12582912 6700624 1048576)) (stdout "") (stderr ""))

(define (search-tree-paths m)
  (let paths ((tree (model->search-tree m)))
    (fold (λ (entry n)
             (+ n
                (match/first
                  (car entry)
                  ((V ,_) 1)
                  ((C ,thunk) (paths (thunk))))))
          0
          tree)))

(define (test/language/search-tree-paths _)
  (⊦ equal? '(4 4 3 4) (map search-tree-paths (list grass pcond1 pcond4 ph2))))

((eta 0.0) (memory #(12582912 6222888 1048576)) (stdout "") (stderr ""))

(define procedure-choice
  (do/monad
    (← f (dist categorical (list (cons car 0.5) (cons cadr 0.5))))
    (return (lift (λ (g) (g '(1 2))) f))))

(define procedure-constant
  (do/monad
    (← c (dist bern 0.5))
    (← f (diracN car))
    (return (lift (λ (g cv) (if cv (g '(1 2)) 0)) f c))))

(define procedure-built
  (do/monad
    (← k (dist uniformly '(1 2 3)))
    (← scale (diracN (lift (λ (kv) (λ (x) (* kv x))) k)))
    (return (lift (λ (f) (f 10)) scale))))

(define circular-value
  (do/monad
    (← b (dist bern 0.5))
    (← cyc (diracN (lift (λ (bv) (circular-list bv)) b)))
    (return (lift car cyc))))

(define distribution-value
  (do/monad
    (← c (dist bern 0.5))
    (← d
       (diracN
         (lift (λ (cv) (if cv (uniformly '(1 2)) (uniformly '(1 2 3 4)))) c)))
    (← x (dist identity d))
    (dist (condition #t bern) (lift (λ (xv) (if (odd? xv) 0.8 0.3)) x))
    (return (pair c x))))

(define record-value
  (do/monad
    (← k (dist uniformly '(1 2 3)))
    (← s (diracN (lift (λ (kv) (make-scaler (λ (x) (* kv x)))) k)))
    (← p (diracN (lift (λ (kv) (delay (* kv 10))) k)))
    (return (lift (λ (sv pv) (list ((scaler-f sv) 10) (force pv))) s p))))

(define (test/language/procedure-values _)
  (⊦ (equal-approx? 1e-12)
     '((1 . 0.5) (2 . 0.5))
     (exact-posterior procedure-choice))
  (⊦ (equal-approx? 1e-12)
     '((0 . 0.5) (1 . 0.5))
     (exact-posterior procedure-constant))
  (⊦ equal? '((1 . 974) (2 . 1026)) (hist (mcmC 2000 procedure-choice)))
  (⊦ equal? '((0 . 1026) (1 . 974)) (hist (mcmC 2000 procedure-constant)))
  (⊦ equal?
     '(50 50)
     (list (length (mcmC 50 procedure-choice #:check #t))
           (length (mcmC 50 procedure-constant #:check #t))))
  (⊦ equal?
     '((1 2) (0 1))
     (list (map car (likelihood-weighting 100 procedure-choice))
           (map car (likelihood-weighting 100 procedure-constant))))
  (⊦ (equal-approx? 1e-12)
     (map (λ (v) (cons v (/ 1.0 3.0))) '(10 20 30))
     (exact-posterior procedure-built))
  (⊦ (equal-approx? 1e-12)
     '((#f . 0.5) (#t . 0.5))
     (exact-posterior circular-value))
  (⊦ (equal-approx? 1e-12)
     (map (Λ ((,v unquote w)) (cons v (/ w 11.0)))
          '(((#f . 1) . 2.0)
            ((#f . 2) . 0.75)
            ((#f . 3) . 2.0)
            ((#f . 4) . 0.75)
            ((#t . 1) . 4.0)
            ((#t . 2) . 1.5)))
     (exact-posterior distribution-value))
  (⊦ (equal-approx? 1e-12)
     (map (λ (v) (cons (list v v) (/ 1.0 3.0))) '(10 20 30))
     (exact-posterior record-value))
  (let1 (scales
          (do/monad
            (← k (dist uniformly '(1 2 3)))
            (diracN (lift (λ (kv) (λ (x) (* kv x))) k))))
        (⊦ equal?
           (map (λ (f) (f 10)) (mcmC 50 scales))
           (map (λ (f) (f 10)) (mcmC 50 scales #:check #t))))
  (for-each
    (λ (m)
       (⊦ equal? (mcmC 50 m) (mcmC 50 m #:check #t))
       (⊦ equal?
          (mcmC 50 m #:semantics 'extended)
          (mcmC 50 m #:semantics 'extended #:check #t)))
    (list procedure-built circular-value distribution-value record-value)))

((eta 0.451) (memory #(12582912 6233600 1048576)) (stdout "") (stderr ""))

(define pnorm (dist normal 10.0 0.5))

(define pdep1 (do/monad (← x pnorm) (diracN (sexp+ x 1.0))))

pnorm = dist normal 10 0.5

pdep1 = do
  x <- pnorm
  diracN (x+1)

-- Literally the same samples
pdep_run = mcmC 100 (dist normal 11 0.5) == mcmC 100 pdep1
-- True

(define (test/replay/pdep _)
  (⊦ equal? (mcmC 100 (dist normal 11.0 0.5)) (mcmC 100 pdep1)))

((eta 0.001) (memory #(12582912 6205976 1048576)) (stdout "") (stderr ""))

(define pdepR (do/monad (← x pnorm) (diracN x)))

(define pdepR′ (do/monad (← x pnorm) (return x)))

pdepR = do
  x <- pnorm
  diracN x

pdepR_run = mcmC 100 pdepR

pdepR' = do
  x <- pnorm
  return x

(define (test/replay/pdepR _)
  (let1 (reference (mcmC 100 pnorm))
        (⊦ equal? reference (mcmC 100 pdepR))
        (⊦ equal? reference (mcmC 100 pdepR′))
        (⊦ < 10 (length (delete-duplicates reference)))))

((eta 0.001) (memory #(12582912 6207360 1048576)) (stdout "") (stderr ""))

(define (nnot p b) (if b p 1.0))

(define ((nor strength-x strength-y noise) x y)
  (bern (- 1.0
           (* (* (nnot (- 1.0 strength-x) x) (nnot (- 1.0 strength-y) y))
              (- 1.0 noise)))))

(define grass
  (do/monad
    (← rain (dist bern 0.3))
    (← sprinkler (dist bern 0.5))
    (← grass-is-wet (dist (condition #t (nor 0.9 0.8 0.1)) rain sprinkler))
    (return rain)))

grass = do
  rain        <- dist bern 0.3
  sprinkler   <- dist bern 0.5
  --  dist (nor 0.9 0.8 0.1) rain sprinkler
  grass_is_wet <- dist (True `condition` nor 0.9 0.8 0.1) rain sprinkler
  return rain

-- noisy-or function
nor :: Double -> Double -> Double -> Bool -> Bool -> Distribution Bool
nor strengthX strengthY noise x y =
  bern $ 1 - nnot (1-strengthX) x * nnot (1-strengthY) y * (1-noise)
  
-- noisy not function
nnot :: Num a => a -> Bool -> a
nnot p True  = p
nnot p False = 1

grass_run = let ns = 20000
                r  = mcmC ns grass
            in
             ((fromIntegral . length . filter id $ r) /
              (fromIntegral ns :: Double))
-- 0.46815 for 20000 -- with old resampling
-- 0.4684714427203697  -- theoretical

(define (test/replay/grass _) (⊦ = 9363 (count-true (mcmC 20000 grass))))

((eta 0.069) (memory #(12582912 7162792 1048576)) (stdout "") (stderr ""))

(define (is1 x) (lift (λ (xv) (= xv 1)) x))

(define pcond1
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (is1 x)
         (dist categorical '((10 . 0.5) (11 . 0.5)))
         (dist categorical '((20 . 0.5) (21 . 0.5))))))

pcond1 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (dist categorical (pure [(10,0.5), (11,0.5)]))
     (dist categorical (pure [(20,0.5), (21,0.5)]))

pcond1r = hist $ mcmC 10000 pcond1
-- fromList [(10,246),(11,222),(20,240),(21,292)]
-- fromList [(10,2481),(11,2394),(20,2565),(21,2560)]

(define (test/replay/pcond1 _)
  (⊦ equal?
     '((10 . 246) (11 . 222) (20 . 240) (21 . 292))
     (hist (mcmC 1000 pcond1)))
  (⊦ equal?
     '((10 . 2481) (11 . 2394) (20 . 2565) (21 . 2560))
     (hist (mcmC 10000 pcond1))))

((eta 0.031) (memory #(12582912 6216672 1048576)) (stdout "") (stderr ""))

(define pcond2
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (is1 x) (return x) (dist categorical '((20 . 0.5) (21 . 0.5))))))

pcond2 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (return x)
     (dist categorical (pure [(20,0.5), (21,0.5)]))

pcond2r = hist $ mcmC 10000 pcond2
-- fromList [(1,500),(20,214),(21,286)]
-- fromList [(1,5000),(20,2321),(21,2679)]

(define (test/replay/pcond2 _)
  (⊦ equal? '((1 . 500) (20 . 214) (21 . 286)) (hist (mcmC 1000 pcond2)))
  (⊦ equal? '((1 . 5000) (20 . 2321) (21 . 2679)) (hist (mcmC 10000 pcond2))))

((eta 0.037) (memory #(12582912 6213776 1048576)) (stdout "") (stderr ""))

(define pcond3
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (is1 x)
         (return x)
         (>> (dist normal 0.0 1.0)
             (dist categorical '((20 . 0.5) (21 . 0.5)))))))

pcond3 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (return x)
     (do
       dist normal 0 1
       dist categorical (pure [(20,0.5), (21,0.5)]))

pcond3r = hist $ mcmC 10000 pcond3
-- fromList [(1,5011),(20,2507),(21,2482)]
-- fromList [(1,5040),(20,2509),(21,2451)] -- new

(define (test/replay/pcond3 _)
  (⊦ equal? '((1 . 5040) (20 . 2509) (21 . 2451)) (hist (mcmC 10000 pcond3))))

((eta 0.034) (memory #(12582912 6220784 1048576)) (stdout "") (stderr ""))

(define pcond4
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (is1 x)
         (return x)
         (>> (dist (condition #t bern) 0.5)
             (dist categorical '((20 . 0.5) (21 . 0.5)))))))

pcond4 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (return x)
     (do
       dist (True `condition` bern) 0.5
       dist categorical (pure [(20,0.5), (21,0.5)]))

_ = hist $ mcmC 1000 pcond4
{-
*Examples1> hist $ mcmC 1000 pcond4
fromList [(1,648),(20,174),(21,178)]
*Examples1> hist $ mcmC 10000 pcond4
fromList [(1,6697),(20,1641),(21,1662)]
-}
-- Exact:
--      [(0.166666666666666657, V 21); (0.166666666666666657, V 20);
--       (0.66666666666666663, V 1)])

(define (test/replay/pcond4 _)
  (⊦ equal? '((1 . 648) (20 . 174) (21 . 178)) (hist (mcmC 1000 pcond4)))
  (⊦ equal? '((1 . 6697) (20 . 1641) (21 . 1662)) (hist (mcmC 10000 pcond4))))

((eta 0.04) (memory #(12582912 6218096 1048576)) (stdout "") (stderr ""))

(define (pcorr* pfuzz result)
  (define ((fuzzy2 pfuzz op) x y)
    (let1 (v (op x y))
          (categorical
            `((,v unquote pfuzz) (,(not v) unquote (- 1.0 pfuzz))))))
  (define (xor y x) (not (eq? x y)))
  (do/monad
    (← x (dist bern 0.5))
    (← y (dist bern 0.5))
    (dist (condition #t (fuzzy2 pfuzz xor)) x y)
    (return (result x y))))

pcorr pfuzz = do
  x <- dist bern 0.5
  y <- dist bern 0.5
  dist (True `condition` fuzzy2 pfuzz xor) x y
  return y
 where
   -- Creating a custom distribution
   fuzzy2 pfuzz op x y = let v = op x y in
     categorical [(v,pfuzz),(not v,1-pfuzz)]
   xor y x = not (x == y)

_ = hist $ mcmC 1000 (pcorr 0.1)
-- fromList [(False,419),(True,581)]

(define (test/replay/pcorr _)
  (⊦ equal? '((#f . 419) (#t . 581)) (hist (mcmC 1000 (pcorr 0.1)))))

((eta 0.004) (memory #(12582912 6214952 1048576)) (stdout "") (stderr ""))

(define girl
  (let ()
    (define ((fuzzy2 pfuzz op) x y)
      (let1 (v (op x y))
            (categorical
              `((,v unquote (- 1.0 pfuzz)) (,(not v) unquote pfuzz)))))
    (do/monad
      (← girl1 (dist bern 0.5))
      (← girl2 (dist bern 0.5))
      (dist (condition #t (fuzzy2 0.01 (λ (a b) (or a b)))) girl1 girl2)
      (return girl1))))

girl = do
  girl1 <- dist bern 0.5
  girl2 <- dist bern 0.5
  dist (True `condition` fuzzy2 0.01 (||)) girl1 girl2
  return girl1
 where
   -- Creating a custom distribution
   fuzzy2 pfuzz op x y = let v = op x y in
     categorical [(v,1-pfuzz),(not v,pfuzz)]

_ = length . filter id $ mcmC 100 girl
-- 77, without the bias

(define (test/replay/girl _) (⊦ = 77 (count-true (mcmC 100 girl))))

((eta 0.001) (memory #(12582912 6217896 1048576)) (stdout "") (stderr ""))

(define (post-bias c)
  (do/monad
    (← coin (dist bern 0.5))
    (if_ coin
         (dist (condition c normal) 0.0 1.0)
         (dist (condition c normal) 100.0 1.0))
    (return coin)))

post_bias c = do
  coin <- dist bern 0.5
  if_ coin (dist (c `condition` normal) 0 1)
           (dist (c `condition` normal) 100 1)
  return coin

_ = mcmC 10 (post_bias 1)
-- All true

(define (test/replay/post-bias _)
  (⊦ equal? (make-list 10 #t) (mcmC 10 (post-bias 1.0)))
  (⊦ equal? (make-list 100 #t) (mcmC 100 (post-bias 1.0))))

((eta 0.0) (memory #(12582912 6218880 1048576)) (stdout "") (stderr ""))

(define (biased-coin c1 c2)
  (do/monad
    (← b (dist beta 1.0 1.0))
    (dist (condition c1 bern) b)
    (dist (condition c2 bern) b)
    (return b)))

biased_coin c1 c2 = do
  b <- dist beta 1 1
  dist (c1 `condition` bern) b
  dist (c2 `condition` bern) b
  return b

biased_coin22 =
  momenta . drop 1000 $
   mcmC 10000 (biased_coin True False)
-- beta 2 2, av = 0.5, sigma2 = 5e-2
-- original Metropolis
-- (0.49433040918945736,5.013851472980666e-2)
-- incremental
-- (0.49910027237934973,4.97358737659099e-2)

(define (test/replay/biased-coin _)
  (⊦ equal?
     '(0.49910027237935 . 0.0497358737659099)
     (momenta (drop (mcmC 10000 (biased-coin #t #f)) 1000))))

((eta 0.052) (memory #(12582912 6357056 1048576)) (stdout "") (stderr ""))

(define mixng
  (do/monad
    (← x (dist normal 0.0 1.0))
    (if_ (lift (λ (xv) (> xv 0.0)) x)
         (dist normal 10.0 2.0)
         (dist gamma 3.0 (/ 1.0 3.0)))))

mixng = do
  x <- dist normal 0 1
  if_ ((>0) <$> x)
           (dist normal 10 2)
           (dist gamma 3 (1/3))

_  =
 let n = 5000 in kl (tabulate 0.5 (mcmC n mixng)) (tabulate 0.5 $ mcmC (n `div` 2) (dist normal 10 2) ++ mcmC (n `div` 2) (dist gamma 3 (1/3)))
{-
5000 samples 1.3563495645465825e-2
1.3896442111895212e-2 now
-}

(define (test/replay/mixng-interleaved _)
  (⊦ =
     0.0138964421118952
     (kl (tabulate 0.5 (mcmC 5000 mixng))
         (tabulate
           0.5
           (append
             (mcmC 2500 (dist normal 10.0 2.0))
             (mcmC 2500 (dist gamma 3.0 (/ 1.0 3.0))))))))

((eta 0.031) (memory #(12582912 6222736 1048576)) (stdout "") (stderr ""))

(define (deeper m) (do/monad (← x m) (dist normal x 3.0)))

(define (phierN n) (fold (λ (_ m) (deeper m)) (dist normal 0.0 1.0) (iota n)))

phier n = 
  (iterate (\m -> do {x<-m; dist normal x 3}) $ dist normal 0 1) !! n

main = do
  [n',nchain'] <- getArgs
  let n = read n' :: Int
  let nchain = read nchain' :: Integer
  print . momenta $ mcmC nchain (phier n)

{-
1000
(3.4117556376053577,32.27027981263249)

(define (test/replay/phier _)
  (⊦ equal?
     '(3.41175563760536 . 32.2702798126325)
     (momenta (mcmC 1000 (phierN 10)))))

((eta 0.004) (memory #(12582912 6284280 1048576)) (stdout "") (stderr ""))

cbr = hist $ mcmC 10000 (cballs_model (replicate 10 Blue))

-- fromList [(1,4265),(2,2203),(3,1244),(4,805),(5,486),(6,291),(7,565),(8,141)]
-- Good agreement

(define (test/replay/colored-balls _)
  (⊦ equal?
     '((1 . 4265)
       (2 . 2203)
       (3 . 1244)
       (4 . 805)
       (5 . 486)
       (6 . 291)
       (7 . 565)
       (8 . 141))
     (hist (mcmC 10000 (cballs-model (make-list 10 'blue))))))

((eta 0.696) (memory #(25165824 6220976 1048576)) (stdout "") (stderr ""))

(define (pbern* bern) (dist bern 0.4))

(define (ph2* bern)
  (do/monad
    (← x (dist bern 0.2))
    (dist bern (lift (λ (xv) (if xv 0.1 0.7)) x))))

(define (ph21* bern)
  (do/monad
    (← x (dist bern 0.2))
    (← y (dist bern (lift (λ (xv) (if xv 0.1 0.7)) x)))
    (return x)))

(define (mixng′* bern)
  (do/monad
    (← x (dist bern 0.5))
    (if_ x (dist normal 10.0 2.0) (dist gamma 3.0 (/ 1.0 3.0)))))

-- Simple models
pbern = dist bern 0.4

pbern_run = mcmC 10 pbern

ph2 = do
  x <- dist bern 0.2
  dist bern ((\xv -> if xv then 0.1 else 0.7) <$> x)
-- 0.2 * 0.1 + 0.8*0.7 = 0.58

ph2r = length . filter id $ mcmC 10000 ph2
-- 5809

ph21 = do
  x <- dist bern 0.2
  y <- dist bern ((\xv -> if xv then 0.1 else 0.7) <$> x)
  return x

ph21r = length . filter id $ mcmC 10000 ph21
-- 2026

mixng' = do
  x <- dist bern 0.5
  if_ x
           (dist normal 10 2)
           (dist gamma 3 (1/3))

_  =
 let n = 5000 in kl (tabulate 0.5 (mcmC n mixng)) (tabulate 0.5 (mcmC n mixng'))
{-
5000 samples 2.275866271810348e-2
5000 samples 1.1301558326703463e-2 now

-- resampling always flips, so fwd and bwd steps are 0
instance RDistribution Bern Bool where
{-
  resample d x _ g =
    let x' = not x in
    ((x', logDensity d x', g),0)
-}

(define (test/replay/always-flip _)
  (let1 (bern bern/flip)
        (⊦ = 4 (count-true (mcmC 10 (pbern* bern))))
        (⊦ = 5809 (count-true (mcmC 10000 (ph2* bern))))
        (⊦ = 2026 (count-true (mcmC 10000 (ph21* bern))))
        (⊦ =
           0.0113015583267035
           (kl (tabulate 0.5 (mcmC 5000 mixng))
               (tabulate 0.5 (mcmC 5000 (mixng′* bern)))))))

((eta 0.071) (memory #(25165824 6227456 1048576)) (stdout "") (stderr ""))

(define tac1
  (do/monad
    (← c (dist bern 0.5))
    (if_ c (dist normal 0.0 1.0) (dist uniform 10.0 20.0))
    (return c)))

(define tac2
  (do/monad
    (← c (dist bern 0.5))
    (if_ c (>> (diracN #t) (dist normal 0.0 1.0)) (dist uniform 10.0 20.0))
    (return c)))

(define tac3
  (do/monad
    (← c (dist bern 0.5))
    (if_ c
         (>> (diracN #t) (dist normal 0.0 1.0))
         (>>= (dist uniform 10.0 20.0) diracN))
    (return c)))

-- Programs from my messages about the problem with conditionals
tac1 = do
  c <- dist bern 0.5
  if_ c
    (dist normal 0 1)
    (dist uniform 10 20)
  return c

tac1_run = length . filter id $ mcmC 1000 tac1
-- 620 True

tac2 = do
  c <- dist bern 0.5
  if_ c
    (diracN (Val True) >> dist normal 0 1)
    (dist uniform 10 20)
  return c

tac2_run = length . filter id $ mcmC 1000 tac2
-- 523 True

tac3 = do
  c <- dist bern 0.5
  if_ c
    (diracN (Val True) >> dist normal 0 1)
    (dist uniform 10 20 >>= diracN)
  return c

tac3_run = length . filter id $ mcmC 1000 tac3
-- 532 True

-- tac2 and tac3 produce identical results, as expected:
-- since diracN is unit, e >>= diracN is indeed the same as diracN

(define (test/replay/tac _)
  (let ((t1 (mcmC 1000 tac1)) (t2 (mcmC 1000 tac2)) (t3 (mcmC 1000 tac3)))
    (⊦ equal? t1 t2)
    (⊦ equal? t2 t3)
    (⊦ = 505 (count-true t1))))

((eta 0.008) (memory #(25165824 6320328 1048576)) (stdout "") (stderr ""))

(define (test/laws/left-unit _)
  (⊦ equal?
     (mcmC 200 (dist normal 5.0 1.0))
     (mcmC 200 (do/monad (← x (diracN 5.0)) (dist normal x 1.0)))))

((eta 0.001) (memory #(25165824 6212720 1048576)) (stdout "") (stderr ""))

{-
diractest = mcmC 5 $ do 
     a <- bernN (Val 0.5)
     diracN a
-}
diractest = mcmc 10 $
  with_node_none (bernS 0.5) `next` (\a ->
  with_node a (\a -> diracS a))
-- The result should show both True and False, in the same proportion

(define (test/laws/diractest _)
  (⊦ (equal-approx? 0.03)
     0.5
     (frequency
       identity
       (mcmC 10000
             (do/monad (← x (dist bern 0.5)) (diracN (lift identity x)))))))

((eta 0.022) (memory #(25165824 6438560 1048576)) (stdout "") (stderr ""))

(define (test/laws/p2 _)
  (⊦ <=
     9990
     (length
       (hist (mcmC 10000
                   (do/monad
                     (← x (dist uniform 0.0 1.0))
                     (diracN (lift identity x))))))))

((eta 0.053) (memory #(25165824 6842280 1048576)) (stdout "") (stderr ""))

(define pmarg (do/monad (← x pbern) (← y pnorm) (return y)))

pmarg = do
  x <- pbern
  y <- pnorm
  return y

pmarg_run = mcmC 100 pmarg

(define (test/laws/marginal _)
  (let1 (m (momenta (mcmC 10000 pmarg)))
        (⊦ (equal-approx? 0.05) '(10.0 . 0.25) m)))

((eta 0.017) (memory #(25165824 6315008 1048576)) (stdout "") (stderr ""))

(define pjoin (do/monad (← x pbern) (← y pnorm) (return (pair x y))))

pjoin = do
  x <- pbern
  y <- pnorm
  return (pair x y)

_ = mcmC 100 pjoin

(define (test/laws/joint _)
  (let* ((samples (mcmC 10000 pjoin))
         (xs (map (λ (s) (if (car s) 1.0 0.0)) samples))
         (ys (map cdr samples))
         (mx (mean xs))
         (my (mean ys))
         (cov (mean (map (λ (x y) (* (- x mx) (- y my))) xs ys)))
         (corr (/ cov (sqrt (* (cdr (momenta xs)) (cdr (momenta ys)))))))
    (⊦ (equal-approx? 0.03) 0.4 mx)
    (⊦ < (abs corr) 0.05)))

((eta 0.038) (memory #(25165824 7391856 1048576)) (stdout "") (stderr ""))

(define (test/laws/branch-unit _) (⊦ equal? (mcmC 1000 tac2) (mcmC 1000 tac3)))

((eta 0.005) (memory #(25165824 6258408 1048576)) (stdout "") (stderr ""))

(define (test/laws/exact _)
  (let ((coin (dist bern 0.4))
        (pmarg/coin
          (do/monad (← x (dist bern 0.4)) (← y (dist bern 0.3)) (return y))))
    (⊦ equal?
       (exact-posterior coin)
       (exact-posterior (do/monad (← x (dist bern 0.4)) (diracN x))))
    (⊦ equal?
       (exact-posterior (dist bern 0.3))
       (exact-posterior (do/monad (← p (diracN 0.3)) (dist bern p))))
    (⊦ (equal-approx? 1e-15)
       (exact-posterior (dist bern 0.3))
       (exact-posterior pmarg/coin))
    (⊦ (equal-approx? 1e-15)
       (exact-posterior ph2)
       (exact-posterior
         (do/monad
           (← x (dist bern 0.2))
           (← q (diracN (lift (λ (xv) (if xv 0.1 0.7)) x)))
           (dist bern q))))))

((eta 0.0) (memory #(25165824 6263624 1048576)) (stdout "") (stderr ""))

So, when it comes to computing the acceptance ratio for a program
with branches, we take

p4 = do
     c <- bernN (Val 0.5)
     if_ c (do e1 e2 ...)
           (do e3 e4 e5 ...)

to be

p4' = do
      c <- bernN (Val 0.5)
      if_ c e_then e_else

That is, we merge (at least for the purpose of presentation) e1, e2 ...
into one node e_then, whose ll is the sum of ll of constituents. Ditto
for e_else.

When we switch the branches, we remove the ll of one branch from
llcum and add ll of the other branch. In our terms, we remove the ll
of the deactivated nodes and add ll of the activated nodes.

This algorithm is resilient with respect to adding replacing a node e
with e >>= diracN -- the desired behavior.

(define pswitch
  (do/monad
    (← c (dist bern 0.5))
    (← m (dist uniformly '(0 1)))
    (if_ c
         (dist categorical (lift (λ (mv) `((,mv . 0.9) (,(- 1 mv) . 0.1))) m))
         (diracN (lift identity m)))
    (return c)))

(define (test/semantics/pswitch/upstream _)
  (⊦ (equal-approx? 0.03) 0.3773 (frequency identity (mcmC 40000 pswitch))))

((eta 0.137) (memory #(25165824 7160488 1048576)) (stdout "") (stderr ""))

(define (test/semantics/pswitch/extended _)
  (⊦ (equal-approx? 0.02)
     0.5
     (frequency identity (mcmC 40000 pswitch #:semantics 'extended))))

((eta 0.154) (memory #(25165824 7160976 1048576)) (stdout "") (stderr ""))

(define (test/semantics/extended/examples _)
  (let1 (run (λ (m) (proportions (mcmC 10000 m #:semantics 'extended))))
        (⊦ (equal-approx? 0.03)
           '((10 . 0.25) (11 . 0.25) (20 . 0.25) (21 . 0.25))
           (run pcond1))
        (⊦ (equal-approx? 0.04)
           '((1 . 0.5) (20 . 0.25) (21 . 0.25))
           (run pcond2))
        (⊦ (equal-approx? 0.05)
           '((1 . 0.5) (20 . 0.25) (21 . 0.25))
           (run pcond3))
        (⊦ (equal-approx? 0.05)
           `((1 unquote (/ 2.0 3.0))
             (20 unquote (/ 1.0 6.0))
             (21 unquote (/ 1.0 6.0)))
           (run pcond4)))
  (let1 (xs (mcmC 10000 mixng #:semantics 'extended))
        (⊦ (equal-approx? 0.4) 5.5 (mean xs))
        (⊦ (equal-approx? 0.05) 0.5 (frequency (λ (x) (> x 5.0)) xs)))
  (let1 (balls (proportions
                 (mcmC 10000
                       (cballs-model (make-list 10 'blue))
                       #:semantics
                       'extended
                       #:seed
                       17)))
        (⊦ = 0.4265 (alist-ref 1 balls))
        (⊦ > (alist-ref 1 balls) (alist-ref 2 balls) (alist-ref 3 balls))))

((eta 1.36) (memory #(25165824 6442552 1048576)) (stdout "") (stderr ""))

(define guarded-car/constant
  (do/monad
    (let xs '())
    (← b (dist bern 0.5))
    (if_ (lift null? xs) (return b) (diracN (lift car xs)))))

(define guarded-car/node
  (do/monad
    (← b (dist bern 0.5))
    (let xs (lift (λ (bv) (if bv '() '(1))) b))
    (if_ (lift null? xs) (return 'none) (diracN (lift car xs)))))

(define act-update
  (do/monad
    (← x (dist uniformly '(0 1 2)))
    (if_ (lift (λ (v) (< v 2)) x)
         (dist (condition 0 categorical)
               (lift (λ (v)
                        (case v
                          ((0) '((0 . 0.9) (1 . 0.1)))
                          ((1) '((0 . 0.3) (1 . 0.7)))
                          (else '((1 . 1.0)))))
                     x))
         (return 'none))
    (return x)))

(define act-update/finite
  (do/monad
    (← x (dist uniformly '(0 1 2)))
    (if_ (lift (λ (v) (< v 2)) x)
         (dist (condition 0 categorical)
               (lift (λ (v)
                        (case v
                          ((0) '((0 . 0.9) (1 . 0.1)))
                          ((1) '((0 . 0.3) (1 . 0.7)))
                          (else '((0 . 0.5) (1 . 0.5)))))
                     x))
         (return 'none))
    (return x)))

(define pending-abort
  (do/monad
    (← x (dist uniformly '(0 1 2)))
    (if_ (lift (λ (v) (< v 2)) x)
         (dist (condition 0 categorical)
               (lift (λ (v)
                        `((0 unquote (if (= v 0) 0.9 0.3))
                          (1 unquote (if (= v 0) 0.1 0.7))))
                     x))
         (return 'none))
    (if_ (lift (λ (v) (not (= v 0))) x)
         (dist uniformly (lift (λ (v) (if (= v 0) '() '(a b))) x))
         (return 'z))
    (return x)))

(define nan-initial
  (do/monad
    (← b (dist uniformly '(1 2 3)))
    (dist (condition 0.0 normal)
          0.0
          (lift (λ (bv) (case bv ((3) 0.0) ((2) 1.0) (else 2.0))) b))
    (← c (dist bern 0.5))
    (return (pair b c))))

(define nan-mean
  (do/monad
    (let mean-of (λ (kv) (/ (exact->inexact (* kv 3)) (exact->inexact kv))))
    (← k (dist uniformly '(0 1 2)))
    (← mean (diracN (lift mean-of k)))
    (← z
       (dist categorical
             (lift (λ (kv) `((,(mean-of kv) . 0) (a . 1) (b . 1))) k)))
    (return (pair k z))))

(define (test/semantics/check _)
  (for-each
    (λ (semantics)
       (for-each
         (Λ ((,name ,m))
            (⊦ equal?
               (list name semantics 300)
               (list name
                     semantics
                     (length (mcmC 300 m #:semantics semantics #:check #t)))))
         `((phier3 ,(phierN 3))
           (mhex ,mhex)
           (exif2 ,exif2)
           (pcond4 ,pcond4)
           (colored-balls ,(cballs-model '(blue green blue)))
           (pswitch ,pswitch)
           (regression ,(regression (regression-data 2 10))))))
    '(upstream extended))
  (for-each
    (λ (semantics)
       (for-each
         (Λ ((,name ,m))
            (⊦ equal?
               (list name semantics (mcmC 50 m #:semantics semantics))
               (list name
                     semantics
                     (mcmC 50 m #:semantics semantics #:check #t))))
         `((guarded-car/constant ,guarded-car/constant)
           (guarded-car/node ,guarded-car/node))))
    '(upstream extended))
  (for-each
    (Λ ((,name ,m))
       (⊦ equal?
          (list name 2000)
          (list name (length (mcmC 2000 m #:semantics 'extended #:check #t)))))
    `((act-update ,act-update)
      (act-update/finite ,act-update/finite)
      (pending-abort ,pending-abort)))
  (⊦ (equal-approx? 0.02)
     `((1 unquote (/ 0.3 1.3)) (2 unquote (/ 1.0 1.3)))
     (proportions (mcmC 20000 pending-abort #:semantics 'extended)))
  (let1 (far-inactive
          (do/monad
            (← c (dist bern 0.5))
            (if_ (lift not c)
                 (do/monad (dist (condition 1000.0 normal) 0.0 0.1) (return 1))
                 (return 0))))
        (for-each
          (λ (semantics)
             (⊦ equal?
                (list semantics #t 100)
                (list semantics
                      (construction-check far-inactive #:semantics semantics)
                      (length
                        (mcmC 100
                              far-inactive
                              #:semantics
                              semantics
                              #:check
                              #t)))))
          '(upstream extended)))
  (let1 (exact (exact-posterior act-update))
        (⊦ (equal-approx? 1e-12)
           `((0 unquote (/ 0.9 2.2))
             (1 unquote (/ 0.3 2.2))
             (2 unquote (/ 1.0 2.2)))
           exact)
        (⊦ (equal-approx? 0.02)
           exact
           (proportions (mcmC 20000 act-update #:semantics 'extended))))
  (for-each
    (λ (semantics)
       (let1 (xs (mcmC 200 nan-initial #:semantics semantics))
             (⊦ equal?
                (list semantics #t xs)
                (list semantics
                      (construction-check nan-initial #:semantics semantics)
                      (mcmC 200 nan-initial #:semantics semantics #:check #t)))
             (⊦ equal?
                (list semantics '(3))
                (list semantics (delete-duplicates (map car xs))))
             (⊦ equal?
                (list semantics 2)
                (list semantics (length (delete-duplicates (map cdr xs)))))))
    '(upstream extended))
  (for-each
    (λ (semantics)
       (⊦ equal?
          (list semantics #t (mcmC 300 nan-mean #:semantics semantics) 0)
          (list semantics
                (construction-check
                  nan-mean
                  #:seed
                  107336
                  #:semantics
                  semantics)
                (mcmC 300 nan-mean #:semantics semantics #:check #t)
                (car (car (mcmC 1
                                nan-mean
                                #:seed
                                107336
                                #:semantics
                                semantics
                                #:check
                                #t))))))
    '(upstream extended))
  (⊦ equal?
     '()
     (filter-map
       (Λ ((,name unquote m))
          (let1 (residual (kernel-residual m))
                (and (not (< residual 1e-12)) (cons name residual))))
       `((act-update unquote act-update)
         (act-update/finite unquote act-update/finite)))))

((eta 1.051) (memory #(25165824 6207944 1048576)) (stdout "") (stderr ""))

(define (test/semantics/check/negative _)
  (let ((k 0)
        (j 0)
        (u 0)
        (runs 0)
        (calls 0)
        (skips 0)
        (shifts 0)
        (victim #f)
        (shifted #f)
        (tamper! #f))
    (define (drifting)
      (let1 (d (normal 0.0 1.0))
            (make-distribution
              'drifting
              #:log-density
              (λ (x)
                 (add1! calls)
                 (- (distribution-log-density d x) (* 0.001 calls)))
              #:sample
              (λ (src) (distribution-sample d src)))))
    (define (unjournaled)
      (let1 (d (normal 0.0 1.0))
            (make-distribution
              'unjournaled
              #:log-density
              (λ (x) (distribution-log-density d x))
              #:sample
              (λ (src)
                 (when (node? victim)
                       (let1 (s (node-state victim))
                             (node-state-set!
                               victim
                               (make-nstate
                                 (nstate-val s)
                                 (nstate-ll s)
                                 (nstate-dst s)
                                 (nstate-tstamp s)))))
                 (distribution-sample d src)))))
    (define behind-the-journal
      (do/monad
        (← x (dist (condition 0.5 normal) 0.0 1.0))
        (let _ (set! victim x))
        (← s (dist unjournaled))
        (dist (condition 3.0 normal) s 0.5)
        (return s)))
    (define (restate-node! n f)
      (let1 (s (node-state n))
            (node-state-set!
              n
              (f (nstate-val s)
                 (nstate-ll s)
                 (nstate-dst s)
                 (nstate-tstamp s)))))
    (define (restate! ch addr f) (restate-node! (chain-node ch addr) f))
    (define shifted-normal
      (do/monad (← x (dist normal 0.0 1.0)) (diracN (sexp+ x 1.0))))
    (define (corrupted-message
             corrupt!
             #!optional
             (m shifted-normal)
             (semantics 'upstream))
      (internal-message
        (τ (receive
             (ch result)
             (build-chain
               (model->trace m 'test)
               (make-stdgen-source 17)
               semantics)
             (corrupt! ch)
             (reference-check ch)))))
    (define (shift-inactive-ll! ch)
      (restate-node!
        (chain-fold
          (λ (n found)
             (or found
                 (and (node-inactive? n) (eq? (node-kind n) 'resampleable) n)))
          #f
          ch)
        (λ (v ll d t) (make-nstate v (+ 1.0 ll) d t))))
    (define (tampering)
      (let1 (d (normal 0.0 1.0))
            (make-distribution
              'tampering
              #:log-density
              (λ (x) (distribution-log-density d x))
              #:sample
              (λ (src)
                 (when tamper! (let1 (t! tamper!) (set! tamper! #f) (t!)))
                 (distribution-sample d src)))))
    (define tampered
      (do/monad
        (← x (dist (condition 0.5 normal) 0.0 1.0))
        (let _ (set! shifted x))
        (dist tampering)))
    (define counter/k (>>= (return #f) (λ (_) (add1! k) (return k))))
    (define counter/j
      (>>= (return #f)
           (λ (_)
              (add1! j)
              (dist (condition 0.0 normal) 0.0 (exact->inexact j)))))
    (define counter/bound
      (>>= (return #f)
           (λ (_)
              (add1! u)
              (if_ #f
                   (dist (condition 5.0 uniform) 0.0 (if (= u 1) 1.0 10.0))
                   (return 0.0)))))
    (define counter/deferred
      (do/monad
        (← c (dist bern 0.5))
        (>>= (return #f)
             (λ (_)
                (add1! k)
                (if_ #t
                     (return c)
                     (diracN (lift (λ (kv) (car (list-tail '() kv))) k)))))
        (return c)))
    (define switching
      (do/monad
        (← x (diracN 1.0))
        (← y (diracN 1.0))
        (>>= (return #f)
             (λ (_) (add1! runs) (diracN (sexp+ (if (odd? runs) x y) 0.0))))))
    (define (third-skips)
      (let1 (d (normal 0.0 1.0))
            (make-distribution
              'third-skips
              #:log-density
              (λ (x) (distribution-log-density d x))
              #:sample
              (λ (src)
                 (add1! skips)
                 (when (zero? (modulo skips 3)) (distribution-sample d src))
                 (distribution-sample d src)))))
    (define (third-shifts)
      (let1 (d (normal 0.0 1.0))
            (make-distribution
              'third-shifts
              #:log-density
              (λ (x) (distribution-log-density d x))
              #:sample
              (λ (src)
                 (add1! shifts)
                 (receive
                   (x ll)
                   (distribution-sample d src)
                   (values x (if (zero? (modulo shifts 3)) (+ ll 1.0) ll)))))))
    (define (third-shifts-message obs semantics)
      (set! shifts 0)
      (internal-message
        (τ (construction-check
             (do/monad
               (← x (dist third-shifts))
               (if_ #f obs (return 0))
               (return x))
             #:semantics
             semantics))))
    (define (tampered-message check k)
      (set! tamper! #f)
      (internal-message
        (τ (let loop ((s (mcmC§ tampered #:semantics 'extended #:check check))
                      (i 0))
             (car§ s)
             (when (= i (sub1 k))
                   (set! tamper!
                     (τ (restate-node!
                          shifted
                          (λ (v ll d t) (make-nstate v (+ 1.0 ll) d t))))))
             (when (< i 20) (loop (cdr§ s) (add1 i)))))))
    (⊦ equal?
       "a rebuild and the initial trace differ in value"
       (internal-message (τ (construction-check counter/k))))
    (⊦ equal?
       "a rebuild and the initial trace differ in LL"
       (internal-message (τ (construction-check counter/j))))
    (⊦ equal?
       "a rebuild and the initial trace differ in value"
       (internal-message (τ (mcmC 10 counter/k #:check #t))))
    (⊦ equal?
       "a rebuild and the initial trace differ in LL"
       (internal-message (τ (mcmC 10 counter/j #:check #t))))
    (⊦ equal?
       "a rebuild and the initial trace differ in LL"
       (internal-message
         (τ (construction-check counter/bound #:semantics 'extended))))
    (⊦ equal?
       (make-list 3 "a rebuild and the initial trace differ in value")
       (list (internal-message (τ (construction-check counter/deferred)))
             (internal-message
               (τ (construction-check counter/deferred #:semantics 'extended)))
             (internal-message (τ (mcmC 10 counter/deferred #:check #t)))))
    (⊦ equal?
       "a rebuild and the initial trace differ in graph"
       (internal-message (τ (construction-check switching))))
    (⊦ equal?
       "an address-order redraw and the initial trace differ in value"
       (internal-message (τ (construction-check (dist third-skips)))))
    (⊦ equal?
       "an address-order redraw and the initial trace differ in LL"
       (internal-message (τ (construction-check (dist third-shifts)))))
    (let ((impossible (dist (condition 5.0 (uniform 0.0 1.0))))
          (far (dist (condition 1000000000000.0 normal) 0.0 1.0)))
      (⊦ equal?
         (make-list
           2
           "an address-order redraw and the initial trace differ in LL")
         (map (cut third-shifts-message <> 'extended) (list impossible far)))
      (⊦ equal?
         '(none none)
         (map (cut third-shifts-message <> 'upstream) (list impossible far)))
      (set! skips 0)
      (⊦ equal?
         "an address-order redraw and the initial trace differ in value"
         (internal-message
           (τ (construction-check
                (do/monad
                  (← x (dist third-skips))
                  (if_ #f impossible (return 0))
                  (return x)))))))
    (⊦ equal?
       '("an LL differs from its recomputation" none)
       (map (λ (semantics)
               (corrupted-message shift-inactive-ll! tac1 semantics))
            '(extended upstream)))
    (for-each
      (Λ ((,expected ,corrupt!))
         (⊦ equal? expected (corrupted-message corrupt!)))
      `(("an LL differs from its recomputation"
         ,(λ (ch)
             (restate! ch 1 (λ (v ll d t) (make-nstate v (+ 1.0 ll) d t)))))
        ("an LL differs from its recomputation"
         ,(λ (ch) (restate! ch 1 (λ (v ll d t) (make-nstate v -inf.0 d t)))))
        ("a value differs from its recomputation"
         ,(λ (ch) (restate! ch 2 (λ (v ll d t) (make-nstate 99.0 ll d t)))))
        ("a distribution differs from its recomputation"
         ,(λ (ch)
             (restate!
               ch
               1
               (λ (v ll d t)
                  (make-nstate v ll (distk-rd (normal 0.0 2.0)) t)))))
        ("a tstamp is after now"
         ,(λ (ch) (restate! ch 1 (λ (v ll d t) (make-nstate v ll d 99)))))
        ("a deactivation count differs from its recomputation"
         ,(λ (ch) (node-inactive-set! (chain-node ch 1) 1)))
        ("an update is delayed on an active node or in 'extended mode"
         ,(λ (ch) (node-delayed-set! (chain-node ch 1) #t)))
        ("a node is still pending or queued after the step"
         ,(λ (ch) (node-pending-set! (chain-node ch 1) #t)))
        ("the card change differs from its recomputation"
         ,(λ (ch) (chain-card-change-set! ch 1)))))
    (⊦ equal?
       "an LL differs from its recomputation"
       (internal-message (τ (mcmC 50 (dist drifting) #:check #t))))
    (⊦ = 50 (length (mcmC 50 behind-the-journal)))
    (for-each
      (λ (semantics)
         (⊦ equal?
            (list semantics "a rejected step did not restore the state")
            (list semantics
                  (internal-message
                    (τ (mcmC 50
                             behind-the-journal
                             #:check
                             #t
                             #:semantics
                             semantics))))))
      '(upstream extended))
    (⊦ equal?
       "the step's delta is not the extended LL change"
       (tampered-message #t 3))
    (⊦ equal?
       "the step's delta is not the extended LL change"
       (tampered-message 5 5))
    (⊦ eq? 'none (tampered-message 5 6))
    (⊦ eq? 'none (tampered-message #f 3))))

((eta 0.002) (memory #(25165824 6207952 1048576)) (stdout "") (stderr ""))

(define (test/semantics/check/escaped-error _)
  (let ((armed #f))
    (define flaky
      (do/monad
        (← x (dist bern 0.1))
        (← y
           (dist normal
                 (lift (λ (xv)
                          (if xv
                            (if armed
                              (begin (set! armed #f) (error "flaky"))
                              5.0)
                            0.0))
                       x)
                 1.0))
        (return (pair x y))))
    (let ((s (mcmC§ flaky #:check #t #:seed 3)))
      (set! armed #t)
      (let1 (reads (let loop ((i 0) (acc '()))
                     (if (= i 12)
                       (reverse acc)
                       (loop (add1 i)
                             (cons (condition-case
                                     (let1 (v (car§ s)) (set! s (cdr§ s)) v)
                                     (e (exn)
                                        ((condition-property-accessor 'exn 'message) e)))
                                   acc)))))
            (⊦ equal? '("flaky") (filter string? reads))
            (⊦ = 11 (count pair? reads))))))

((eta 0.001) (memory #(25165824 6214696 1048576)) (stdout "") (stderr ""))

(define invalid-inactive
  (do/monad
    (← x (dist normal 0.0 1.0))
    (if_ (lift positive? x) (dist gamma x 1.0) (return 0.0))))

(define invalid-inactive/beta
  (do/monad
    (← x (dist normal 0.0 1.0))
    (if_ (lift positive? x) (dist beta x 1.0) (return 0.0))))

(define (invalid-inactive/branch p)
  (do/monad
    (← c (dist bern p))
    (← x (dist bern 0.5))
    (if_ c
         (return 3)
         (if_ (lift (λ (cv xv) (if cv 0 xv)) c x) (return 1) (return 2)))
    (return c)))

(define active-invalid
  (do/monad
    (← b (dist bern 0.9))
    (dist (condition 0.5 beta) (lift (λ (bv) (if bv 1.0 -1.0)) b) 1.0)
    (return b)))

(define (test/semantics/invalid-inactive _)
  (let1 (xs (mcmC 2000 invalid-inactive #:semantics 'extended))
        (⊦ = 2000 (length xs))
        (⊦ every positive? xs))
  (⊦ eq? 'parameter (hakaru10-kind (τ (mcmC 2000 invalid-inactive))))
  (let1 (xs (mcmC 2000 invalid-inactive/beta #:semantics 'extended))
        (⊦ = 2000 (length xs))
        (⊦ every positive? xs))
  (⊦ eq? 'parameter (hakaru10-kind (τ (mcmC 2000 invalid-inactive/beta))))
  (⊦ eq?
     'branch
     (hakaru10-kind
       (τ (mcmC 1 (invalid-inactive/branch 0.5) #:semantics 'extended))))
  (⊦ eq?
     'branch
     (hakaru10-kind
       (τ (mcmC 2000 (invalid-inactive/branch 0.05) #:semantics 'extended))))
  (⊦ eq? 'branch (hakaru10-kind (τ (mcmC 1 (invalid-inactive/branch 0.5)))))
  (let1 (xs (mcmC 2000 (invalid-inactive/branch 0.05)))
        (⊦ = 2000 (length xs))
        (⊦ any identity xs))
  (⊦ equal?
     '(branch branch branch)
     (map (λ (oracle)
             (hakaru10-kind (τ (oracle (invalid-inactive/branch 0.5)))))
          (list exact-posterior
                (cut likelihood-weighting 100 <>)
                (cut weighted-samples 100 <>))))
  (⊦ (equal-approx? 1e-12)
     '((#f . 0.95) (#t . 0.05))
     (exact-posterior (invalid-inactive/branch 0.05)))
  (⊦ equal?
     '(none none)
     (map (λ (oracle)
             (hakaru10-kind (τ (oracle 100 (invalid-inactive/branch 0.05)))))
          (list likelihood-weighting weighted-samples)))
  (for-each
    (λ (semantics)
       (⊦ equal?
          (list semantics 'parameter 'parameter)
          (list semantics
                (hakaru10-kind
                  (τ (mcmC 10 (dist normal 0.0 0.0) #:semantics semantics)))
                (hakaru10-kind
                  (τ (mcmC 2000 active-invalid #:semantics semantics))))))
    '(extended upstream))
  (⊦ = 1 (length (mcmC 1 active-invalid #:semantics 'extended))))

((eta 0.029) (memory #(25165824 6224216 1048576)) (stdout "") (stderr ""))

(define guarded-car
  (do/monad
    (← k (dist uniformly '(0 1 2)))
    (← xs (diracN (lift (λ (kv) (list-tail '(a b) kv)) k)))
    (← empty (diracN (lift null? xs)))
    (if_ empty (return 'none) (diracN (lift car xs)))))

(define nested-guarded-car
  (do/monad
    (← k (dist uniformly '(0 1 2)))
    (← xs (diracN (lift (λ (kv) (list-tail '(a b) kv)) k)))
    (← c (dist bern 0.5))
    (← r
       (if_ (lift null? xs)
            (return 'none)
            (if_ c (diracN (lift car xs)) (return 'other))))
    (return (pair c r))))

(define guarded-car/reversed
  (do/monad
    (← k (dist uniformly '(2 1 0)))
    (← xs (diracN (lift (λ (kv) (list-tail '(a b) kv)) k)))
    (← empty (diracN (lift null? xs)))
    (if_ empty (return 'none) (diracN (lift car xs)))))

(define guarded-inside-lift
  (do/monad
    (← b (dist uniformly '(0 1 2)))
    (← c (dist bern 0.5))
    (return
      (lift (λ (bv cv)
               (handle-exceptions
                 e
                 'caught
                 (if (and cv (= bv 0))
                   (error "guarded-inside-lift")
                   (list bv cv))))
            b
            c))))

(define (strict-reader xs)
  (do/monad
    (← c (dist bern 0.5))
    (if_ (lift null? xs)
         (return c)
         (do/monad
           (← h (diracN (lift car xs)))
           (← mu (diracN (lift cdr (pair h 1.0))))
           (dist normal mu 1.0)))))

(define strict-reader/node
  (do/monad
    (← c (dist bern 0.5))
    (let xs (lift (λ (cv) (if cv '() '(1))) c))
    (if_ (lift null? xs)
         (return 'empty)
         (do/monad
           (← h (diracN (lift car xs)))
           (← mu (diracN (lift cdr (pair h 1.0))))
           (dist normal mu 1.0)))))

(define ignoring-reader
  (do/monad
    (← c (dist bern 0.5))
    (if_ #t
         (return c)
         (do/monad
           (let d (lift car '()))
           (← mu (diracN (lift (λ (x) 1.0) d)))
           (dist normal mu 1.0)))))

(define strict-active
  (do/monad
    (← c (dist bern 0.5))
    (← xs (diracN (lift (λ (cv) '()) c)))
    (if_ #t (>> (diracN (lift car xs)) (return c)) (return c))))

(define strict-active/val
  (do/monad
    (← c (dist bern 0.5))
    (if_ #t (>> (diracN (lift car '())) (return c)) (return c))))

(define (test/semantics/deferred _)
  (let1 (thirds (map (λ (v) (cons v (/ 1.0 3.0))) '(a b none)))
        (⊦ (equal-approx? 1e-12) thirds (exact-posterior guarded-car))
        (for-each
          (λ (semantics)
             (⊦ (equal-approx? 0.02)
                thirds
                (proportions (mcmC 6000 guarded-car #:semantics semantics)))
             (⊦ equal?
                (mcmC 300 guarded-car #:semantics semantics)
                (mcmC 300 guarded-car #:semantics semantics #:check #t)))
          '(upstream extended)))
  (⊦ equal? '(none none b none a) (mcmC 5 guarded-car))
  (let1 (exact (exact-posterior nested-guarded-car))
        (⊦ (equal-approx? 1e-12)
           (map cons
                '((#f . none) (#f . other) (#t . a) (#t . b) (#t . none))
                (map (λ (w) (/ w 6.0)) '(1.0 2.0 1.0 1.0 1.0)))
           exact)
        (for-each
          (λ (semantics)
             (⊦ (equal-approx? 0.03)
                (list semantics exact)
                (list semantics
                      (proportions
                        (mcmC 20000
                              nested-guarded-car
                              #:semantics
                              semantics))))
             (⊦ equal?
                (mcmC 300 nested-guarded-car #:semantics semantics)
                (mcmC 300
                      nested-guarded-car
                      #:semantics
                      semantics
                      #:check
                      #t)))
          '(upstream extended)))
  (⊦ equal?
     '()
     (filter-map
       (Λ ((,name unquote m))
          (let1 (residual (kernel-residual m))
                (and (not (< residual 1e-12)) (cons name residual))))
       `((guarded-car unquote guarded-car)
         (guarded-car/reversed unquote guarded-car/reversed)
         (nested-guarded-car unquote nested-guarded-car))))
  (⊦ equal?
     '(a none)
     (list (car (mcmC 1 guarded-car/reversed)) (car (mcmC 1 guarded-car))))
  (let1 (sixths
          (map (λ (v) (cons v (/ 1.0 6.0)))
               '(caught (0 #f) (1 #f) (1 #t) (2 #f) (2 #t))))
        (⊦ (equal-approx? 1e-12) sixths (exact-posterior guarded-inside-lift)))
  (⊦ equal?
     '(exn none exn)
     (map (λ (m) (raised (τ (mcmC 3 m))))
          (list (strict-reader '()) (strict-reader '(5)) strict-reader/node)))
  (⊦ equal? '() (mcmC 3 (strict-reader '()) #:semantics 'extended))
  (⊦ every flonum? (mcmC 300 strict-reader/node #:semantics 'extended))
  (for-each
    (λ (m)
       (⊦ equal?
          '(exn exn exn)
          (list (raised (τ (mcmC 3 m)))
                (raised (τ (mcmC 3 m #:semantics 'extended)))
                (raised (τ (exact-posterior m))))))
    (list strict-active strict-active/val))
  (⊦ equal?
     '(exn ())
     (list (raised (τ (mcmC 3 ignoring-reader)))
           (mcmC 3 ignoring-reader #:semantics 'extended)))
  (let1 (inner #f)
        (mcmC 1
              (do/monad
                (← c (dist bern 0.5))
                (if_ #t
                     (return c)
                     (diracN
                       (lift (λ (k)
                                (set! inner
                                  (raised
                                    (τ (exact-posterior
                                         (do/monad
                                           (← xs (return '()))
                                           (let _ (lift car xs))
                                           (return 'leaked))))))
                                k)
                             1)))))
        (⊦ eq? 'exn inner)))

((eta 0.339) (memory #(25165824 6222072 1048576)) (stdout "") (stderr ""))

(define pdep
  (do/monad
    (← x (dist normal 0.0 1.0))
    (← y (dist normal x 1.0))
    (← z (dist normal y 1.0))
    (return (sexp+ x y z))))

(define (update-traces n model . keys)
  (let ((updated '()) (steps '()))
    (define (spy! ch)
      (chain-fold
        (λ (node _)
           (let1 (update (node-update node))
                 (when update
                       (node-update-set!
                         node
                         (λ (n) (push! (node-addr n) updated) (update n))))))
        (void)
        ch))
    (define (collect! ch accepted? delta)
      (push! (reverse updated) steps)
      (set! updated '()))
    (let1 (next (apply mcmc-generator
                       n
                       (model->trace model 'update-traces)
                       #:build-hook
                       spy!
                       #:step-hook
                       collect!
                       keys))
          (do ((i 0 (add1 i))) ((= i n)) (next)))
    (reverse steps)))

(define upp1
  (do/monad
    (← a (dist uniform 0.0 1.0))
    (← b (diracN (sexp+ a 2.0)))
    (← c (diracN (sexp+ a 3.0)))
    (← d (dist uniform (sexp+ a b) c))
    (return d)))

(define upp2
  (do/monad
    (← a (dist uniform 0.0 1.0))
    (← b (diracN (sexp+ a 2.0)))
    (← c (diracN (sexp+ b 3.0)))
    (← d (dist uniform (sexp+ a b) c))
    (return d)))

(define upp1/raw
  (do/monad
    (← a (with-node-none (uniform 0.0 1.0)))
    (← b (with-node a (λ (av) (dirac (+ av 2.0)))))
    (← c (with-node a (λ (av) (dirac (+ av 3.0)))))
    (← ab
       (with-node-tree
         (vector a b)
         (λ (vs) (dirac (+ (vector-ref vs 0) (vector-ref vs 1))))))
    (with-node-tree
      (vector ab c)
      (λ (vs) (uniform (vector-ref vs 0) (vector-ref vs 1))))))

(define (test/semantics/incremental _)
  (for-each
    (λ (semantics)
       (⊦ equal?
          (list semantics '((2 4) (3 4) (4)))
          (list semantics
                (map car
                     (hist (update-traces 1000 pdep #:semantics semantics)))))
       (for-each
         (λ (m)
            (⊦ equal?
               (list semantics '(() (2 3 4)))
               (list semantics
                     (map car
                          (hist (update-traces 1000 m #:semantics semantics)))))
            (⊦ = 300 (length (mcmC 300 m #:semantics semantics #:check #t))))
         (list upp1 upp2))
       (⊦ equal?
          (list semantics '(() (2 3 4 5)))
          (list semantics
                (map car
                     (hist (update-traces
                             1000
                             upp1/raw
                             #:semantics
                             semantics)))))
       (⊦ = 300 (length (mcmC 300 upp1/raw #:semantics semantics #:check #t))))
    '(upstream extended)))

((eta 0.056) (memory #(25165824 6234936 1048576)) (stdout "") (stderr ""))

(define (test/statistics/pbern _)
  (for-each
    (λ (seed)
       (⊦≈ seed 0.025 0.4 (frequency identity (mcmC 10000 pbern #:seed seed))))
    seeds))

((eta 0.012) (memory #(25165824 6437536 1048576)) (stdout "") (stderr ""))

(define (test/statistics/pnorm _)
  (for-each
    (λ (seed)
       (⊦≈ seed 0.03 '(10.0 . 0.25) (momenta (mcmC 10000 pnorm #:seed seed))))
    seeds))

((eta 0.02) (memory #(25165824 6347008 1048576)) (stdout "") (stderr ""))

(define ph1 (do/monad (← x (dist normal 0.0 1.0)) (dist normal x 3.0)))

(define (test/statistics/ph1 _)
  (for-each
    (λ (seed)
       (let1 (m (momenta (mcmC 20000 ph1 #:seed seed)))
             (⊦≈ seed 0.3 0.0 (car m))
             (⊦≈ seed 1.5 10.0 (cdr m))))
    seeds))

((eta 0.056) (memory #(25165824 6295680 1048576)) (stdout "") (stderr ""))

(define (test/statistics/phier3 _)
  (for-each
    (λ (seed)
       (let1 (m (momenta (mcmC 20000 (phierN 3) #:seed seed)))
             (⊦≈ seed 1.5 0.0 (car m))
             (⊦≈ seed 8.0 28.0 (cdr m))))
    seeds))

((eta 0.065) (memory #(25165824 6297192 1048576)) (stdout "") (stderr ""))

phier_kl n =
  kl (tabulate 0.5 (mcmC n phier))
     (tabulate 0.5 (mcmC n (dist normal 0 (sqrt <$> 91))))
     

_ =  phier_kl 10

{-
*Examples1> phier_kl 1000
0.7862745014080637
*Examples1> phier_kl 10000
0.23618257021796316
*Examples1> phier_kl 100000
3.441636543766045e-2
*Examples1> phier_kl 1000000
2.1573977909591603e-3
-}

((kl (tabulate 0.5 (mcmC 40000 (phierN 10)))
     (normal-bins 0.5 (sqrt 91.0) bins))
 0.0770140470263418)

(define (test/statistics/phier-kl _)
  (for-each
    (λ (seed)
       (let1 (p (tabulate 0.5 (mcmC 40000 (phierN 10) #:seed seed)))
             (⊦< seed (kl p (normal-bins 0.5 (sqrt 91.0) (map car p))) 0.25)))
    seeds))

((eta 0.297) (memory #(25165824 6223368 1048576)) (stdout "") (stderr ""))

(define (test/statistics/ph2 _)
  (for-each
    (λ (seed)
       (⊦≈ seed 0.05 0.58 (frequency identity (mcmC 10000 ph2 #:seed seed)))
       (⊦≈ seed 0.05 0.2 (frequency identity (mcmC 10000 ph21 #:seed seed))))
    seeds))

((eta 0.038) (memory #(25165824 6446432 1048576)) (stdout "") (stderr ""))

(define (test/statistics/pcond _)
  (for-each
    (λ (seed)
       (let1 (run (λ (m) (proportions (mcmC 10000 m #:seed seed))))
             (⊦≈ seed
                 0.03
                 '((10 . 0.25) (11 . 0.25) (20 . 0.25) (21 . 0.25))
                 (run pcond1))
             (⊦≈ seed 0.04 '((1 . 0.5) (20 . 0.25) (21 . 0.25)) (run pcond2))
             (⊦≈ seed 0.05 '((1 . 0.5) (20 . 0.25) (21 . 0.25)) (run pcond3))
             (⊦≈ seed
                 0.05
                 `((1 unquote (/ 2.0 3.0))
                   (20 unquote (/ 1.0 6.0))
                   (21 unquote (/ 1.0 6.0)))
                 (run pcond4))))
    seeds))

((eta 0.132) (memory #(25165824 6449016 1048576)) (stdout "") (stderr ""))

(define (test/statistics/mixng _)
  (for-each
    (λ (seed)
       (for-each
         (λ (m)
            (let1 (xs (mcmC 10000 m #:seed seed))
                  (⊦≈ seed 0.4 5.5 (mean xs))
                  (⊦≈ seed 0.05 0.5 (frequency (λ (x) (> x 5.0)) xs))))
         (list mixng mixng′)))
    seeds))

((eta 0.069) (memory #(25165824 6526144 1048576)) (stdout "") (stderr ""))

{-
test_multiple_conditions13 =
  momenta . drop 1000 $
   mcmC 10000 (prog_multiple_conditions False False)
-- beta 1 3: 0.25, sigma2 = 3.75e-2
-- original Metropolis
-- (0.25144728141545775,3.782788408489756e-2)
-- incremental
-- (0.2445212550806874,3.651432172401049e-2)
-}

(define (test/statistics/biased-coin _)
  (for-each
    (λ (seed)
       (let ((tf (momenta
                   (drop (mcmC 10000 (biased-coin #t #f) #:seed seed) 1000)))
             (ff (momenta
                   (drop (mcmC 10000 (biased-coin #f #f) #:seed seed) 1000))))
         (⊦≈ seed 0.02 0.5 (car tf))
         (⊦≈ seed 0.008 0.05 (cdr tf))
         (⊦≈ seed 0.02 0.25 (car ff))
         (⊦≈ seed 0.008 0.0375 (cdr ff))))
    seeds))

((eta 0.104) (memory #(25165824 6377440 1048576)) (stdout "") (stderr ""))

(define (nice-prog-mult-conditions c1 c2)
  (do/monad
    (← b (dist beta 1.0 1.0))
    (dist (condition c1 bern) b)
    (dist (condition c2 bern) b)
    (diracN b)))

nice_prog_mult_conditions c1 c2 = do
  b <- dist beta 1 1
  dist (c1 `condition` bern) b
  dist (c2 `condition` bern) b
  diracN b

niceexmcr11 = sum $ mcmC 1000 (nice_prog_mult_conditions True False)
-- 499.799988743311

niceexmcr12 = sum $ mcmC 1000 (nice_prog_mult_conditions False False)
-- 257.2928205920028

(define (test/statistics/nice-prog-mult-conditions _)
  (for-each
    (λ (seed)
       (let1 (average
               (λ (c1 c2)
                  (/ (fold +
                           0.0
                           (mcmC 1000
                                 (nice-prog-mult-conditions c1 c2)
                                 #:seed
                                 seed))
                     1000.0)))
             (⊦≈ seed 0.05 0.5 (average #t #f))
             (⊦≈ seed 0.05 0.25 (average #f #f))))
    seeds))

((eta 0.01) (memory #(25165824 6244456 1048576)) (stdout "") (stderr ""))

(define niceex1
  (do/monad
    (← c (dist uniform 0.0 1.0))
    (← d (dist bern 0.5))
    (← e (dist bern c))
    (diracN (lift (λ (ev dv) (and ev dv)) e d))))

niceex1 = do
  c <- dist uniform 0 1
  d <- dist bern 0.5
  e <- dist bern c
  diracN ((&&) <$> e <*> d)

niceex1r = mcmC 7 niceex1
-- All True

(define (test/statistics/niceex1 _)
  (for-each
    (λ (seed)
       (⊦≈ seed
           0.05
           0.25
           (frequency identity (mcmC 10000 niceex1 #:seed seed))))
    seeds))

((eta 0.025) (memory #(25165824 6454392 1048576)) (stdout "") (stderr ""))

(define (prog-mixture1 cv)
  (do/monad
    (← c (dist bern 0.5))
    (if_ c
         (dist (condition cv normal) 1.0 1.0)
         (dist (condition cv uniform) 0.0 3.0))
    (return c)))

(define tac10
  (do/monad
    (← c (dist bern 0.5))
    (if_ c
         (dist (condition 0.0 normal) 0.0 1.0)
         (dist (condition 0.0 uniform) 10.0 20.0))
    (return c)))

(define tac11
  (do/monad
    (← c (dist bern 0.5))
    (if_ c
         (dist (condition 0.0 normal) 0.0 1.0)
         (do/monad (← d (diracN 10.0)) (dist (condition 0.0 uniform) d 20.0)))
    (return c)))

prog_mixture1 cv = do
  c <- dist bern 0.5
  if_ c (dist (cv `condition` normal)  1 1)
        (dist (cv `condition` uniform) 0 3)
  return c

prog_mixture1_run = mcmC 20 (prog_mixture1 (-2))
-- all true

tac10 = do
  c <- dist bern 0.5
  if_ c
    -- then
    (dist (0 `condition` normal) 0 1)
    -- else
    (dist (0 `condition` uniform) 10 20)
  return c

tac10_run = length . filter id $ mcmC 100 tac10
-- 100  -- all True

tac11 = do
  c <- dist bern 0.5
  if_ c
    -- then
    (dist (0 `condition` normal) 0 1)
    -- else
    (do
        d <- diracN 10
        dist (0 `condition` uniform) d 20)
  return c

tac11_run = length . filter id $ mcmC 100 tac11
-- 100  -- all True

(define (test/statistics/impossible-branches _)
  (for-each
    (λ (seed)
       (⊦ equal?
          (list seed (make-list 100 #t))
          (list seed (mcmC 100 (prog-mixture1 -2.0) #:seed seed)))
       (⊦ equal?
          (list seed (make-list 100 #t))
          (list seed (mcmC 100 tac10 #:seed seed)))
       (⊦ equal?
          (list seed (make-list 100 #t))
          (list seed (mcmC 100 tac11 #:seed seed))))
    seeds)
  (⊦ equal?
     (make-list 3 '((#t . 1.0)))
     (map exact-posterior (list (prog-mixture1 -2.0) tac10 tac11))))

((eta 0.002) (memory #(25165824 6218208 1048576)) (stdout "") (stderr ""))

(define (test/statistics/tac1 _)
  (for-each
    (λ (seed)
       (⊦≈ seed 0.05 0.5 (frequency identity (mcmC 20000 tac1 #:seed seed))))
    seeds))

((eta 0.054) (memory #(25165824 6697184 1048576)) (stdout "") (stderr ""))

(define exif1
  (do/monad
    (← c (dist bern 0.5))
    (← d
       (if_ c
            (dist normal (lift (λ (cv) (if cv 1.0 0.0)) c) 1.0)
            (diracN 5.0)))
    (← e (if_ (lift not c) (dist normal 20.0 1.0) (diracN d)))
    (diracN (sexp+ d e))))

(define exif2
  (let ()
    (define (cl1 x) (do/monad (← d (dist bern 0.5)) (if_ d cl3 (diracN x))))
    (define cl3
      (do/monad
        (← a (dist uniform 0.0 1.0))
        (← b (dist uniform 0.0 1.0))
        (diracN (sexp+ a b))))
    (do/monad (← c (dist bern 0.5)) (if_ c (cl1 10.0) (cl1 20.0)))))

exif1 = do
  c <- dist bern 0.5
  d <- if_ c (dist normal ((fromIntegral . fromEnum) <$> c) 1)
             (diracN 5)
  e <- if_ (not <$> c) (dist normal 20 1) (diracN d)
  diracN (d + e)

exif1r = mcmC 7 exif1

exif2 = do
  c <- dist bern 0.5
  if_ c (cl1 10) (cl1 20)
 where
   cl1 x = do
     d <- dist bern 0.5
     if_ d cl3 (diracN x)
     
   cl3 = do
     a <- dist uniform 0 1
     b <- dist uniform 0 1
     diracN (a + b)
     

exif2r = mcmC 17 exif2

(define (test/statistics/exif _)
  (for-each
    (λ (seed)
       (let1 (xs (mcmC 20000 exif1 #:seed seed))
             (⊦≈ seed 0.8 13.5 (mean xs))
             (⊦≈ seed 0.05 0.5 (frequency (λ (x) (> x 10.0)) xs)))
       (let1 (xs (mcmC 20000 exif2 #:seed seed))
             (⊦≈ seed 0.04 0.25 (frequency (λ (x) (eqv? x 10.0)) xs))
             (⊦≈ seed 0.04 0.25 (frequency (λ (x) (eqv? x 20.0)) xs))
             (⊦≈ seed 0.04 0.5 (frequency (λ (x) (<= x 2.0)) xs))))
    seeds))

((eta 0.199) (memory #(25165824 6772096 1048576)) (stdout "") (stderr ""))

(define (test/statistics/pcorr _)
  (for-each
    (λ (seed)
       (let1 (xys (mcmC 20000 (pcorr* 0.1 pair) #:seed seed))
             (⊦≈ seed 0.03 0.9 (frequency (Λ ((,x unquote y)) (eq? x y)) xys))
             (⊦≈ seed 0.1 0.5 (frequency cdr xys))))
    seeds))

((eta 0.116) (memory #(25165824 6984120 1048576)) (stdout "") (stderr ""))

(define (phierN′ n start obs)
  (fold/monad
    (λ (x k) (k x))
    start
    (append
      (make-list (sub1 n) (λ (x) (dist normal x 3.0)))
      (list (λ (x) (dist (condition obs normal) x 3.0))))))

(define condhier
  (do/monad (← x (dist normal 0.0 1.0)) (phierN′ 9 x 1.0) (return x)))

phierN' :: Int -> SExp Double -> Double -> Model Double
phierN' n start obs =
  foldl (>>=) (pure start) $
  replicate (n-1) (\x -> dist normal x 3) ++
  [(\x -> dist (obs `condition` normal) x 3)]

condhier = do
  x <- dist normal 0 1
  phierN' 9 x 1
  return x

(define (test/statistics/condhier _)
  (for-each
    (λ (seed)
       (let1 (m (momenta (mcmC 20000 condhier #:seed seed)))
             (⊦≈ seed 0.2 (/ 1.0 82.0) (car m))
             (⊦≈ seed 0.25 (/ 81.0 82.0) (cdr m))))
    seeds))

((eta 0.073) (memory #(25165824 6261240 1048576)) (stdout "") (stderr ""))

(define grass/if
  (do/monad
    (← rain (dist bern 0.3))
    (← sprinkler (dist bern 0.5))
    (let obs (λ (p) (dist (condition #t bern) p)))
    (if_ rain (obs 0.9) (if_ sprinkler (obs 0.8) (obs 0.1)))
    (return rain)))

(define grass/wet
  (do/monad
    (← rain (dist bern 0.3))
    (← sprinkler (dist bern 0.5))
    (dist (nor 0.9 0.8 0.1) rain sprinkler)))

{-
grass = do
  rain        <- dist bern 0.3
  sprinkler   <- dist bern 0.5
  let obs = dist (True `condition` bern)
  if_ rain
    (obs 0.9)
    (if_ sprinkler
        (obs 0.8)
        (obs 0.1))
  return rain
-}

(define (test/statistics/grass-girl _)
  (for-each
    (λ (seed)
       (⊦≈ seed
           0.04
           0.4684714
           (frequency identity (mcmC 20000 grass #:seed seed)))
       (⊦≈ seed
           0.05
           0.6644295
           (frequency identity (mcmC 10000 girl #:seed seed)))
       (⊦≈ seed
           0.04
           (/ 6.0 13.0)
           (frequency identity (mcmC 20000 grass/if #:seed seed)))
       (⊦≈ seed
           0.05
           0.6058
           (frequency identity (mcmC 40000 grass/wet #:seed seed))))
    seeds)
  (⊦ (equal-approx? 1e-12)
     `((#f unquote (/ 7.0 13.0)) (#t unquote (/ 6.0 13.0)))
     (exact-posterior grass/if))
  (⊦ (equal-approx? 1e-12)
     '((#f . 0.3942) (#t . 0.6058))
     (exact-posterior grass/wet)))

((eta 0.276) (memory #(25165824 6224024 1048576)) (stdout "") (stderr ""))

(define mhex
  (do/monad
    (← x (dist normal 0.0 1.0))
    (← y (dist normal x 2.0))
    (if_ (lift (λ (xv) (> xv 0.5)) x)
         (return x)
         (do/monad (← z (dist beta 1.0 1.0)) (return (sexp+ y z))))))

mhex = do
  x <- dist normal 0 1
  y <- dist normal x 2
  if_ ((>0.5) <$> x)
        (return x)
        (do {z <- dist beta 1 1; return (y + z)})

(define (test/statistics/mhex _)
  (for-each
    (λ (seed) (⊦≈ seed 0.1 0.345731 (mean (mcmC 20000 mhex #:seed seed))))
    seeds))

((eta 0.093) (memory #(25165824 6917160 1048576)) (stdout "") (stderr ""))

(define (test/statistics/post-bias _)
  (for-each
    (λ (seed)
       (⊦ equal?
          (list seed (make-list 80 #t))
          (list seed (drop (mcmC 100 (post-bias 1.0) #:seed seed) 20))))
    seeds))

((eta 0.001) (memory #(25165824 6230600 1048576)) (stdout "") (stderr ""))

     -- the parent nodes must have been updated
     when (tstamp (Mem.dref entry) <= tstamp self &&
           tstamp nd <= tstamp self) $
       fail $ unwords ["the parents of the exit node", show (Mem.addr ns),
                       "namely", show [Mem.addr entry, Mem.addr ndr],
                       "must have been updated"]

(define shared-inactive
  (do/monad
    (← c (dist bern 0.5))
    (← d (dist normal 0.0 1.0))
    (if_ c (return d) (dist normal 5.0 1.0))))

(define (test/statistics/if/shared-inactive _)
  (for-each
    (λ (seed)
       (⊦≈ seed
           0.06
           0.5
           (frequency
             (λ (x) (< x 2.5))
             (mcmC 10000 shared-inactive #:seed seed))))
    seeds))

((eta 0.029) (memory #(25165824 6539568 1048576)) (stdout "") (stderr ""))

type 'a vc = V of 'a | C of (unit -> 'a pV) and 'a pV = (prob * 'a vc) list

So, when it comes to computing the acceptance ratio for a program
with branches, we take

p4 = do
     c <- bernN (Val 0.5)
     if_ c (do e1 e2 ...)
           (do e3 e4 e5 ...)

to be

p4' = do
      c <- bernN (Val 0.5)
      if_ c e_then e_else

That is, we merge (at least for the purpose of presentation) e1, e2 ...
into one node e_then, whose ll is the sum of ll of constituents. Ditto
for e_else.

When we switch the branches, we remove the ll of one branch from
llcum and add ll of the other branch. In our terms, we remove the ll
of the deactivated nodes and add ll of the activated nodes.

This algorithm is resilient with respect to adding replacing a node e
with e >>= diracN -- the desired behavior.

(module
  (aux hakaru10 exact)
  *
  (import
    scheme
    (chicken base)
    (only (chicken memory representation) record-instance? record->vector)
    (only srfi-1 filter-map count fold iota any)
    srfi-69
    (aux base)
    (only (aux hansei) probcc-τ probcc-unit probcc-explore probcc-normalize)
    (only (aux hakaru10 random) make-stdgen-source rng-source-copy)
    (aux hakaru10 distribution)
    (only (aux hakaru10 util) histogram-sort)
    (aux hakaru10 engine))
  (define (ll-close? a b #!optional (scale (abs a)))
    (or (= a b)
        (and (nan? a) (nan? b))
        (and (finite? a)
             (finite? b)
             (<= (abs (- a b)) (* 1e-09 (max 1.0 scale))))))
  (define (same-value? a b)
    (let1 (budget 100000)
          (define (spent?) (set! budget (sub1 budget)) (negative? budget))
          (let walk ((a a) (b b))
            (cond ((eq? a b) #t)
                  ((and (flonum? a) (flonum? b))
                   (or (eqv? a b) (and (nan? a) (nan? b))))
                  ((and (procedure? a) (procedure? b)) #t)
                  ((and (promise? a) (promise? b)) #t)
                  ((and (pair? a) (pair? b))
                   (or (spent?)
                       (and (walk (car a) (car b)) (walk (cdr a) (cdr b)))))
                  ((and (vector? a) (vector? b))
                   (and (= (vector-length a) (vector-length b))
                        (let loop ((i 0))
                          (or (= i (vector-length a))
                              (spent?)
                              (and (walk (vector-ref a i) (vector-ref b i))
                                   (loop (add1 i)))))))
                  ((or (deferred? a) (deferred? b)) (equal? a b))
                  ((and (record-instance? a) (record-instance? b))
                   (walk (record->vector a) (record->vector b)))
                  (else (equal? a b))))))
  (define (choice-support dk who)
    (or (distribution-support dk)
        (hakaru10-error
          'support
          who
          "a random choice without a finite support"
          dk)))
  (define (build model src semantics who)
    (build-chain (model->trace model who) src semantics))
  (define (static-graph model who)
    (let1 (src (make-stdgen-source 17))
          (let retry ((builds 1000))
            (if (zero? builds)
              (hakaru10-error
                'construction
                who
                "no complete build of the model in 1000 tries"
                model)
              (receive
                (ch result)
                (build model src 'extended who)
                (if result (values ch result) (retry (sub1 builds))))))))
  (define (walk-active! n)
    (let1 (k (scope-inactivity n)) (node-inactive-set! n k) (zero? k)))
  (define (branch-starts ch)
    (let1 (starts (make-vector (add1 (chain-count ch)) '()))
          (chain-fold
            (λ (n _)
               (when (eq? (node-role n) 'exit)
                     (let1 (b (node-branch n))
                           (for-each
                             (Λ ((,lo unquote hi) ,polarity)
                                (when (< lo hi)
                                      (vector-set!
                                        starts
                                        (add1 lo)
                                        (cons (list hi (branch-entry b) polarity)
                                              (vector-ref starts (add1 lo))))))
                             (list (branch-then-range b) (branch-else-range b))
                             '(#t #f)))))
            (void)
            ch)
          starts))
  (define (past-inactive-branches starts a)
    (let1 (hi (and (< a (vector-length starts))
                   (any (Λ ((,hi ,entry ,polarity))
                           (and (not (eq? (node-value entry) polarity)) hi))
                        (vector-ref starts a))))
          (if hi (past-inactive-branches starts (add1 hi)) a)))
  (define (graph->search-tree
           ch
           result
           #!key
           (max-paths #f)
           (who 'exact-posterior))
    (let ((last (chain-count ch))
          (now (chain-now ch))
          (starts (branch-starts ch))
          (paths 0))
      (define (restore! frame)
        (for-each
          (Λ ((,n unquote s)) (node-inactive-set! n 0) (node-state-set! n s))
          frame))
      (let walk ((a (past-inactive-branches starts 1)) (frame '()))
        (define (enter n s)
          (node-state-set! n s)
          (walk (past-inactive-branches starts (add1 a))
                (cons (cons n s) frame)))
        (if (> a last)
          (begin
            (add1! paths)
            (when (and max-paths (> paths max-paths))
                  (hakaru10-error
                    'support
                    who
                    "more paths than max-paths"
                    max-paths))
            (probcc-unit (force-value (node-value result))))
          (let1 (n (chain-node ch a))
                (cond ((not (walk-active! n))
                       (walk (past-inactive-branches starts (add1 a)) frame))
                      ((resampleable? n)
                       (filter-map
                         (Λ ((,x unquote p))
                            (and (positive? p)
                                 (probcc-τ
                                   p
                                   (restore! frame)
                                   (enter n (recomputed-state n now x)))))
                         (choice-support (node-distk n) who)))
                      ((eq? (node-kind n) 'observed)
                       (let* ((s (recomputed-state n now))
                              (w (exp (nstate-ll s))))
                         (if (or (not (finite? (nstate-ll s))) (zero? w))
                           '()
                           (list (probcc-τ w (restore! frame) (enter n s))))))
                      (else (enter n (recomputed-state n now)))))))))
  (define (model->search-tree model #!key (max-paths #f))
    (receive
      (ch result)
      (static-graph model 'model->search-tree)
      (graph->search-tree
        ch
        result
        #:max-paths
        max-paths
        #:who
        'model->search-tree)))
  (define (exact-posterior model #!key (max-paths #f))
    (receive
      (ch result)
      (static-graph model 'exact-posterior)
      (histogram-sort
        (map (Λ (((V ,v) ,p)) (cons v p))
             (probcc-normalize
               (probcc-explore
                 +inf.0
                 (graph->search-tree ch result #:max-paths max-paths)))))))
  (define (graph-forward ch result src #!optional (starts (branch-starts ch)))
    (let ((last (chain-count ch)) (now (chain-now ch)))
      (let walk ((a (past-inactive-branches starts 1)) (ll 0.0))
        (define (next) (past-inactive-branches starts (add1 a)))
        (if (> a last)
          (values (force-value (node-value result)) ll)
          (let1 (n (chain-node ch a))
                (cond ((not (walk-active! n)) (walk (next) ll))
                      ((memq (node-kind n) '(resampleable observed))
                       (receive
                         (x llx)
                         (distk-sample (node-distk n) src)
                         (node-state-set! n (recomputed-state n now x))
                         (walk (next) (if (resampleable? n) ll (+ ll llx)))))
                      (else
                       (node-state-set! n (recomputed-state n now))
                       (walk (next) ll))))))))
  (define (weighted-samples n model #!key (seed 17) (rng #f))
    (receive
      (ch result)
      (static-graph model 'weighted-samples)
      (let ((src (or rng (make-stdgen-source seed)))
            (starts (branch-starts ch)))
        (let loop ((i 0) (samples '()))
          (if (>= i n)
            (reverse samples)
            (receive
              (v ll)
              (graph-forward ch result src starts)
              (loop (add1 i) (cons (cons v ll) samples))))))))
  (define (likelihood-weighting n model #!key (seed 17) (rng #f))
    (let* ((samples
             (filter-map
               (Λ ((,v unquote ll)) (and (finite? ll) (cons v ll)))
               (weighted-samples n model #:seed seed #:rng rng)))
           (top (fold (λ (s m) (max (cdr s) m)) -inf.0 samples))
           (table (make-hash-table equal?)))
      (for-each
        (Λ ((,v unquote ll))
           (hash-table-update!/default
             table
             v
             (λ (w) (+ w (exp (- ll top))))
             0.0))
        samples)
      (let1 (total (hash-table-fold table (λ (v w acc) (+ acc w)) 0.0))
            (histogram-sort
              (map (Λ ((,v unquote w)) (cons v (/ w total)))
                   (hash-table->alist table))))))
  (define (node-signature n)
    (list (node-addr n)
          (node-kind n)
          (node-role n)
          (map node-addr (vector->list (node-parents n)))
          (map (Λ ((,entry unquote polarity))
                  (cons (node-addr entry) polarity))
               (node-scope n))))
  (define (same-built-value? a b)
    (cond ((and (deferred? a) (deferred? b))
           (same-value? (deferred-args a) (deferred-args b)))
          ((or (deferred? a) (deferred? b)) #f)
          (else (same-value? a b))))
  (define (first-mismatch same? ch ch*)
    (let loop ((a 1))
      (and (<= a (chain-count ch))
           (if (same? (chain-node ch a) (chain-node ch* a))
             (loop (add1 a))
             a))))
  (define (construction-check
           model
           #!key
           (seed 17)
           (rng #f)
           (semantics 'upstream))
    (let1 (src (or rng (make-stdgen-source seed)))
          (receive
            (ch result)
            (build model (rng-source-copy src) semantics 'check)
            (if (or (not result) (rejected-ll? (chain-llcum ch)))
              'skipped
              (receive
                (ch* result*)
                (build model (rng-source-copy src) semantics 'check)
                (rebuild-check ch result ch* result*)
                (redraw-check ch (rng-source-copy src))
                #t)))))
  (define (rebuild-check ch result ch* result*)
    (define (differ what . args)
      (apply hakaru10-error
             'internal
             'check
             (string-append "a rebuild and the initial trace differ in " what)
             args))
    (unless
      (and result*
           (= (chain-count ch) (chain-count ch*))
           (= (node-addr result) (node-addr result*)))
      (differ "graph" (chain-count ch) (and result* (chain-count ch*))))
    (cond ((first-mismatch
             (λ (n n*) (equal? (node-signature n) (node-signature n*)))
             ch
             ch*)
           =>
           (λ (a) (differ "graph" a)))
          ((first-mismatch
             (λ (n n*) (same-built-value? (node-value n) (node-value n*)))
             ch
             ch*)
           =>
           (λ (a)
              (differ
                "value"
                a
                (node-value (chain-node ch a))
                (node-value (chain-node ch* a)))))
          ((first-mismatch
             (λ (n n*)
                (ll-close?
                  (nstate-ll (node-state n))
                  (nstate-ll (node-state n*))))
             ch
             ch*)
           =>
           (λ (a) (differ "LL" a)))
          ((not (ll-close? (chain-llcum ch) (chain-llcum ch*)))
           (differ "LL" (chain-llcum ch) (chain-llcum ch*)))))
  (define (redraw-check ch src)
    (define (differ what . args)
      (apply hakaru10-error
             'internal
             'check
             (string-append
               "an address-order redraw and the initial trace differ in "
               what)
             args))
    (match/first
      (chain-fold
        (Λ (,n (,ll unquote scale))
           (cond ((resampleable? n)
                  (receive
                    (x llx)
                    (distk-sample (node-distk n) src)
                    (unless
                      (same-value? x (node-value n))
                      (differ "value" (node-addr n) x (node-value n)))
                    (cons (+ ll llx) (+ scale (abs llx)))))
                 ((eq? (node-kind n) 'observed)
                  (let ((llo (nstate-ll (node-state n)))
                        (inactive? (node-inactive? n)))
                    (cons (if inactive? ll (+ ll llo))
                          (if (and inactive? (extended? ch))
                            scale
                            (+ scale (abs llo))))))
                 (else (cons ll scale))))
        '(0.0 . 0.0)
        ch)
      ((,ll unquote scale)
       (let1 (llcum (chain-llcum ch))
             (unless
               (or (nan? llcum) (ll-close? ll llcum scale))
               (differ "LL" ll llcum))))))
  (define (resampleable? n) (eq? (node-kind n) 'resampleable))
  (define (node-distk n)
    (let1 (fn (node-fn n))
          (if (distk? fn) fn (fn (leaf-values (node-parents n))))))
  (define (recomputed-state n ts #!optional (v (node-value n)))
    (if (eq? (node-role n) 'exit)
      (make-nstate
        (let1 (v (node-value (branch-active-result (node-branch n))))
              (if (node-inactive? n) v (force-value v)))
        0.0
        #f
        ts)
      (let* ((dk (node-distk n)) (kind (distk-kind dk)) (rd (distk-rd dk)))
        (unless
          (eq? kind (node-kind n))
          (hakaru10-error
            'construction
            'check
            "a distribution constructor changed the kind of its node"
            n
            (node-kind n)
            kind))
        (case kind
          ((dirac)
           (make-nstate
             (let1 (x (distk-value dk))
                   (if (node-inactive? n) x (force-value x)))
             0.0
             #f
             ts))
          ((observed)
           (let1 (obs (distk-value dk))
                 (check-observation n obs)
                 (make-nstate obs ((rdistr-log-density rd) obs) rd ts)))
          (else (make-nstate v ((rdistr-log-density rd) v) rd ts))))))
  (define (scope-inactivity n)
    (count (Λ ((,entry unquote polarity))
              (not (eq? (node-value entry) polarity)))
           (node-scope n)))
  (define (same-rdistr? a b)
    (if (and a b)
      (and (eq? (rdistr-name a) (rdistr-name b))
           (same-value? (rdistr-params a) (rdistr-params b)))
      (eq? a b)))
  (define (reference-check ch)
    (define (fail msg n . args)
      (apply hakaru10-error 'internal 'check msg (node-addr n) args))
    (define (check-state n s)
      (let1 (s* (recomputed-state n (nstate-tstamp s) (nstate-val s)))
            (unless
              (same-value? (nstate-val s) (nstate-val s*))
              (fail "a value differs from its recomputation"
                    n
                    (nstate-val s)
                    (nstate-val s*)))
            (unless
              (same-rdistr? (nstate-dst s) (nstate-dst s*))
              (fail "a distribution differs from its recomputation"
                    n
                    (nstate-dst s)
                    (nstate-dst s*)))
            (unless
              (ll-close? (nstate-ll s) (nstate-ll s*))
              (fail "an LL differs from its recomputation"
                    n
                    (nstate-ll s)
                    (nstate-ll s*)))))
    (let ((now (chain-now ch)) (every-node? (extended? ch)))
      (chain-fold
        (λ (n _)
           (let ((s (node-state n)) (inactive (node-inactive n)))
             (unless
               (= inactive (scope-inactivity n))
               (fail "a deactivation count differs from its recomputation"
                     n
                     inactive
                     (scope-inactivity n)))
             (when (and (node-delayed n) (or every-node? (zero? inactive)))
                   (fail "an update is delayed on an active node or in 'extended mode"
                         n))
             (when (or (node-pending n) (node-queued n))
                   (fail "a node is still pending or queued after the step" n))
             (when (> (nstate-tstamp s) now)
                   (fail "a tstamp is after now" n (nstate-tstamp s) now))
             (when (or every-node? (zero? inactive)) (check-state n s))))
        (void)
        ch)
      (let1 (card (- (count node-inactive? (chain-sampleable ch))))
            (unless
              (= (chain-card-change ch) card)
              (hakaru10-error
                'internal
                'check
                "the card change differs from its recomputation"
                (chain-card-change ch)
                card)))
      #t))
  (define (resampleable-nodes ch) (reverse (chain-sampleable ch)))
  (define (load-state! ch xs)
    (let* ((now (chain-now ch))
           (rest (chain-fold
                   (λ (n xs)
                      (node-inactive-set! n (scope-inactivity n))
                      (node-state-set!
                        n
                        (if (resampleable? n)
                          (if (pair? xs)
                            (recomputed-state n now (car xs))
                            (hakaru10-error
                              'construction
                              'load-state!
                              "too few values"
                              n))
                          (recomputed-state n now)))
                      (node-delayed-set! n #f)
                      (node-pending-set! n #f)
                      (if (resampleable? n) (cdr xs) xs))
                   xs
                   ch)))
      (unless
        (null? rest)
        (hakaru10-error 'construction 'load-state! "too many values" rest))
      (chain-card-change-set!
        ch
        (- (count node-inactive? (chain-sampleable ch))))))
  (define (ext-states ch #!key (max-states 100000))
    (let ((now (chain-now ch)) (last (chain-count ch)) (found 0))
      (reverse
        (let walk ((a 1) (prefix '()) (states '()))
          (if (> a last)
            (begin
              (add1! found)
              (when (> found max-states)
                    (hakaru10-error
                      'support
                      'ext-states
                      "more states than max-states"
                      max-states))
              (cons (reverse prefix) states))
            (let1 (n (chain-node ch a))
                  (node-inactive-set! n (scope-inactivity n))
                  (if (resampleable? n)
                    (fold (Λ ((,x unquote p) ,states)
                             (if (positive? p)
                               (begin
                                 (node-state-set! n (recomputed-state n now x))
                                 (walk (add1 a) (cons x prefix) states))
                               states))
                          states
                          (choice-support (node-distk n) 'ext-states))
                    (begin
                      (node-state-set! n (recomputed-state n now))
                      (walk (add1 a) prefix states)))))))))
  (define (forced-step-log-alpha ch n x)
    (when (or (not (resampleable? n)) (node-inactive? n))
          (hakaru10-error
            'internal
            'forced-step-log-alpha
            "an active resampleable node is expected"
            n))
    (journal-begin! ch)
    (chain-now-set! ch (add1 (chain-now ch)))
    (chain-llcum-set! ch 0.0)
    (let* ((card-old (chain-card-change ch))
           (s (node-state n))
           (ll ((rdistr-log-density (nstate-dst s)) x))
           (fb (apply-proposal! n x ll (- ll (nstate-ll s))))
           (delta (begin (propagate! ch) (chain-llcum ch)))
           (ratio (if (rejected-ll? delta)
                    -inf.0
                    (log-acceptance-ratio
                      ch
                      (vector-length (chain-samp ch))
                      card-old
                      fb))))
      (journal-rollback! ch)
      ratio))
  (define (kernel-residual
           model
           #!key
           (seed 17)
           (max-states 100000)
           (semantics 'extended))
    (receive
      (ch result)
      (build model (make-stdgen-source seed) semantics 'kernel-residual)
      (unless
        result
        (hakaru10-error
          'construction
          'kernel-residual
          "the model's first build was rejected"))
      (let* ((nodes (resampleable-nodes ch))
             (states (ext-states ch #:max-states max-states))
             (log-weights
               (map (λ (xs) (load-state! ch xs) (chain-extended-ll ch))
                    states))
             (top (let1 (top (fold (λ (l top) (if (finite? l) (max l top) top))
                                   -inf.0
                                   log-weights))
                        (if (finite? top)
                          top
                          (hakaru10-error
                            'construction
                            'kernel-residual
                            "no state of the model has a positive finite weight"
                            top))))
             (weights
               (map (λ (l) (if (finite? l) (exp (- l top)) 0.0)) log-weights))
             (z (fold + 0.0 weights))
             (π (make-hash-table equal?))
             (πT (make-hash-table equal?)))
        (define (flow! xs mass)
          (hash-table-update!/default πT xs (λ (m) (+ m mass)) 0.0))
        (for-each
          (λ (xs w)
             (hash-table-update!/default π xs (λ (p) (+ p (/ w z))) 0.0))
          states
          weights)
        (hash-table-walk
          π
          (λ (xs p)
             (when (positive? p)
                   (load-state! ch xs)
                   (let* ((active
                            (filter-map
                              (λ (n i)
                                 (and (not (node-inactive? n)) (cons n i)))
                              nodes
                              (iota (length nodes))))
                          (pick (if (null? active) 0.0 (/ p (length active)))))
                     (when (null? active) (flow! xs p))
                     (for-each
                       (Λ ((,n unquote i))
                          (for-each
                            (Λ ((,x unquote q))
                               (when (positive? q)
                                     (let* ((ratio (forced-step-log-alpha ch n x))
                                            (accept
                                              (if (nan? ratio) 0.0 (min 1.0 (exp ratio))))
                                            (mass (* pick q)))
                                       (flow! (list-replace xs i x) (* mass accept))
                                       (flow! xs (* mass (- 1.0 accept))))))
                            (proposal-support n)))
                       active)))))
        (fold (λ (xs m)
                 (max m
                      (abs (- (hash-table-ref/default πT xs 0.0)
                              (hash-table-ref/default π xs 0.0)))))
              0.0
              (append (hash-table-keys π) (hash-table-keys πT))))))
  (define (proposal-support n)
    (let1 (dk (node-distk n))
          (unless
            (rdistr-independence? (distk-rd dk))
            (hakaru10-error
              'support
              'kernel-residual
              "a random choice whose proposal is not the independence one"
              dk))
          (choice-support dk 'kernel-residual)))
  (define (list-replace xs i x)
    (if (zero? i)
      (cons x (cdr xs))
      (cons (car xs) (list-replace (cdr xs) (sub1 i) x)))))

(define (nnot p x) (if x p 1.0))

(define ((nor strength-x strength-y noise) x y)
  (bern (- 1.0
           (* (nnot (- 1.0 strength-x) x)
              (nnot (- 1.0 strength-y) y)
              (- 1.0 noise)))))

(define grass
  (do/monad
    (← rain (dist bern 0.3))
    (← sprinkler (dist bern 0.5))
    (dist (condition #t (nor 0.9 0.8 0.1)) rain sprinkler)
    (return rain)))

grass = do
  rain        <- dist bern 0.3
  sprinkler   <- dist bern 0.5
  --  dist (nor 0.9 0.8 0.1) rain sprinkler
  grass_is_wet <- dist (True `condition` nor 0.9 0.8 0.1) rain sprinkler
  return rain

-- noisy-or function
nor :: Double -> Double -> Double -> Bool -> Bool -> Distribution Bool
nor strengthX strengthY noise x y =
  bern $ 1 - nnot (1-strengthX) x * nnot (1-strengthY) y * (1-noise)

-- noisy not function
nnot :: Num a => a -> Bool -> a
nnot p True  = p
nnot p False = 1

-- 0.46815 for 20000 -- with old resampling
-- 0.4684714427203697  -- theoretical

(define (test/grass _)
  (⊦ (equal-approx? 1e-12)
     '((#f . 0.53152855727963) (#t . 0.46847144272037))
     (exact-posterior grass)))

((eta 0.0) (memory #(8396592 3852176 1048576)) (stdout "") (stderr ""))

(define girl
  (let ()
    (define ((fuzzy2 pfuzz op) x y)
      (let1 (v (op x y))
            (categorical
              `((,v unquote (- 1.0 pfuzz)) (,(not v) unquote pfuzz)))))
    (do/monad
      (← girl1 (dist bern 0.5))
      (← girl2 (dist bern 0.5))
      (dist (condition #t (fuzzy2 0.01 (λ (x y) (or x y)))) girl1 girl2)
      (return girl1))))

girl = do
  girl1 <- dist bern 0.5
  girl2 <- dist bern 0.5
  dist (True `condition` fuzzy2 0.01 (||)) girl1 girl2
  return girl1
 where
   -- Creating a custom distribution
   fuzzy2 pfuzz op x y = let v = op x y in
     categorical [(v,1-pfuzz),(not v,pfuzz)]

(define (test/girl _)
  (⊦ (equal-approx? 1e-12)
     '((#f . 0.335570469798658) (#t . 0.664429530201342))
     (exact-posterior girl)))

((eta 0.0) (memory #(8396592 3857088 1048576)) (stdout "") (stderr ""))

(define (pcorr/pair pfuzz)
  (define (xor y x) (not (eq? x y)))
  (define ((fuzzy2 pfuzz op) x y)
    (let1 (v (op x y))
          (categorical
            `((,v unquote pfuzz) (,(not v) unquote (- 1.0 pfuzz))))))
  (do/monad
    (← x (dist bern 0.5))
    (← y (dist bern 0.5))
    (dist (condition #t (fuzzy2 pfuzz xor)) x y)
    (return (pair x y))))

pcorr pfuzz = do
  x <- dist bern 0.5
  y <- dist bern 0.5
  dist (True `condition` fuzzy2 pfuzz xor) x y
  return y
 where
   -- Creating a custom distribution
   fuzzy2 pfuzz op x y = let v = op x y in
     categorical [(v,pfuzz),(not v,1-pfuzz)]
   xor y x = not (x == y)

(define (test/pcorr _)
  (let1 (posterior (exact-posterior (pcorr/pair 0.1)))
        (⊦ (equal-approx? 1e-12)
           0.9
           (probability-of (Λ ((,x unquote y)) (eq? x y)) posterior))
        (⊦ (equal-approx? 1e-12) 0.5 (probability-of cdr posterior))
        (⊦ (equal-approx? 1e-12)
           '(((#f . #f) . 0.45)
             ((#f . #t) . 0.05)
             ((#t . #f) . 0.05)
             ((#t . #t) . 0.45))
           posterior)))

((eta 0.0) (memory #(8396592 3853560 1048576)) (stdout "") (stderr ""))

(define ph2
  (do/monad
    (← x (dist bern 0.2))
    (dist bern (lift (λ (xv) (if xv 0.1 0.7)) x))))

(define ph21
  (do/monad
    (← x (dist bern 0.2))
    (dist bern (lift (λ (xv) (if xv 0.1 0.7)) x))
    (return x)))

ph2 = do
  x <- dist bern 0.2
  dist bern ((\xv -> if xv then 0.1 else 0.7) <$> x)
-- 0.2 * 0.1 + 0.8*0.7 = 0.58

ph21 = do
  x <- dist bern 0.2
  y <- dist bern ((\xv -> if xv then 0.1 else 0.7) <$> x)
  return x

(define (test/ph2 _)
  (⊦ (equal-approx? 1e-12) 0.58 (probability #t (exact-posterior ph2)))
  (⊦ (equal-approx? 1e-12) 0.2 (probability #t (exact-posterior ph21))))

((eta 0.0) (memory #(8396592 3858080 1048576)) (stdout "") (stderr ""))

(define (one? x) (lift (cut = <> 1) x))

(define pcond1
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (one? x)
         (dist categorical '((10 . 0.5) (11 . 0.5)))
         (dist categorical '((20 . 0.5) (21 . 0.5))))))

(define pcond2
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (one? x) (return x) (dist categorical '((20 . 0.5) (21 . 0.5))))))

(define pcond4
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (one? x)
         (return x)
         (>> (dist (condition #t bern) 0.5)
             (dist categorical '((20 . 0.5) (21 . 0.5)))))))

pcond1 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (dist categorical (pure [(10,0.5), (11,0.5)]))
     (dist categorical (pure [(20,0.5), (21,0.5)]))

pcond2 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (return x)
     (dist categorical (pure [(20,0.5), (21,0.5)]))

pcond4 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (return x)
     (do
       dist (True `condition` bern) 0.5
       dist categorical (pure [(20,0.5), (21,0.5)]))

-- Exact:
--      [(0.166666666666666657, V 21); (0.166666666666666657, V 20);
--       (0.66666666666666663, V 1)])

(define (test/pcond _)
  (⊦ (equal-approx? 1e-12)
     '((10 . 0.25) (11 . 0.25) (20 . 0.25) (21 . 0.25))
     (exact-posterior pcond1))
  (⊦ (equal-approx? 1e-12)
     '((1 . 0.5) (20 . 0.25) (21 . 0.25))
     (exact-posterior pcond2))
  (⊦ (equal-approx? 1e-12)
     `((1 unquote (/ 2.0 3.0))
       (20 unquote (/ 1.0 6.0))
       (21 unquote (/ 1.0 6.0)))
     (exact-posterior pcond4)))

((eta 0.0) (memory #(8396592 3862472 1048576)) (stdout "") (stderr ""))

(define pcond3
  (do/monad
    (← x (dist categorical '((1 . 0.5) (2 . 0.5))))
    (if_ (one? x)
         (return x)
         (>> (dist normal 0.0 1.0)
             (dist categorical '((20 . 0.5) (21 . 0.5)))))))

pcond3 = do
  x <- dist categorical (pure [(1,0.5), (2,0.5)])
  if_ ((==1) <$> x)
     (return x)
     (do
       dist normal 0 1
       dist categorical (pure [(20,0.5), (21,0.5)]))

(define (test/pcond3-support _)
  (⊦= 'support (hakaru10-kind (τ (exact-posterior pcond3))))
  (⊦= 'none (hakaru10-kind (τ (mcmC 10 pcond3)))))

((eta 0.001) (memory #(8396592 3870880 1048576)) (stdout "") (stderr ""))

(define (post-bias c)
  (do/monad
    (← coin (dist bern 0.5))
    (if_ coin
         (dist (condition c normal) 0.0 1.0)
         (dist (condition c normal) 100.0 1.0))
    (return coin)))

(define nan-path
  (do/monad
    (← k (dist uniformly '(0 1 2)))
    (dist (condition 0.0 normal) 0.0 (lift exact->inexact k))
    (return k)))

(define infinite-path
  (do/monad
    (← b (dist bern 0.5))
    (dist (condition 0.0 beta) (lift (λ (bv) (if bv 0.5 2.0)) b) 1.0)
    (return b)))

post_bias c = do
  coin <- dist bern 0.5
  if_ coin (dist (c `condition` normal) 0 1)
           (dist (c `condition` normal) 100 1)
  return coin

_ = mcmC 10 (post_bias 1)
-- All true

(define (test/post-bias _)
  (⊦ equal? '((#t . 1.0)) (exact-posterior (post-bias 1.0)))
  (⊦ equal? '() (exact-posterior (dist (condition 5.0 uniform) 0.0 1.0)))
  (⊦ (equal-approx? 1e-12)
     `((1 unquote (/ 2.0 3.0)) (2 unquote (/ 1.0 3.0)))
     (exact-posterior nan-path))
  (⊦ equal?
     '((1 2) (1 2) ((0 . 30)))
     (list (map car (likelihood-weighting 1000 nan-path))
           (map car (hist (mcmC 1000 nan-path)))
           (hist (mcmC 30 nan-path #:seed 110866))))
  (⊦ equal?
     '(() () ())
     (list (exact-posterior infinite-path)
           (likelihood-weighting 100 infinite-path)
           (mcmC 5 infinite-path))))

((eta 0.007) (memory #(8396592 3870664 1048576)) (stdout "") (stderr ""))

((1 . 0.41196368095257)
 (2 . 0.209728624754566)
 (3 . 0.120691976674053)
 (4 . 0.0801847685132436)
 (5 . 0.0590318166419537)
 (6 . 0.046603790739908)
 (7 . 0.0386298621005105)
 (8 . 0.0331654796231948))

let
    [(0.031942788306915014, V 8); (0.0422088248548995876, V 7);
     (0.0481489259210673634, V 6); (0.0611501366775565255, V 5);
     (0.0843776503618006435, V 4); (0.123159075103266638, V 3);
     (0.203898291338526577, V 2); (0.405114307435967658, V 1)] =
 snd (Inference.normalize t1);;

let
  [(0.0318386565368576138, V 8); (0.0384083436381660176, V 7);
   (0.0445344709122082669, V 6); (0.055953476873295431, V 5);
   (0.0777907061915841, V 4); (0.123751456856916756, V 3);
   (0.211354336837841883, V 2); (0.416368552153129956, V 1)]
 = t2;;

(define (opposite-color color) (if (eq? color 'blue) 'green 'blue))

(define (observed-color color)
  (categorical `((,color . 0.8) (,(opposite-color color) . 0.2))))

(define (balls-prior max-balls)
  (do/monad
    (← balls (replicate/monad max-balls (dist uniformly '(blue green))))
    (return (map cons (iota max-balls 1) balls))))

(define (balls-observe max-balls balls obs nballs)
  (>> (for-each/monad
        (λ (color)
           (do/monad
             (← b (dist uniformly (iota nballs 1)))
             (fold-right
               (λ (i rest)
                  (if_ (lift (cut = <> i) b)
                       (dist (condition color observed-color)
                             (alist-ref i balls))
                       rest))
               (return color)
               (iota max-balls 1))))
        obs)
      (return (void))))

(define (cballs-model max-balls obs)
  (do/monad
    (← balls (balls-prior max-balls))
    (← nballs (dist uniformly (iota max-balls 1)))
    (fold-right
      (λ (i rest)
         (if_ (lift (cut = <> i) nballs)
              (balls-observe max-balls balls obs i)
              rest))
      (return (void))
      (iota max-balls 1))
    (return nballs)))

(define (cballs-closed-form max-balls n)
  (define (binomial m k)
    (if (zero? k) 1 (/ (* m (binomial (sub1 m) (sub1 k))) k)))
  (define (likelihood m)
    (fold (λ (k acc)
             (+ acc
                (* (binomial m k)
                   (expt 2.0 (- m))
                   (expt (/ (+ (* 0.8 k) (* 0.2 (- m k))) m) n))))
          0.0
          (iota (add1 m))))
  (let* ((ls (map likelihood (iota max-balls 1))) (z (fold + 0.0 ls)))
    (map (cut / <> z) ls)))

-- Distribution for the observing of a ball, with the failure rate 20%
observed_color color =
  categorical [(color,0.8),(opposite_color color,0.2)]

maxBalls :: Int
maxBalls = 8

-- Create the prior for the max amount of balls that may be in the urn
balls_prior n = do
  balls <- sequence . replicate n $ dist uniformly (pure [Blue,Green])
  return $ M.fromList $ zip [1..] balls

-- Submodel: Given balls color and the observations, perform conditioning
balls_observe :: M.IntMap (SExp Color) -> [Color] -> Int -> Model ()
balls_observe balls obs nballs = mapM_ observe_ball obs >> return (pure ())
 where
  -- The same ball can be drawn twice as we return the drawn balls
  -- into the urn. A ball keeps its true color, no matter how many
  -- times it is drawn.
   observe_ball color = do
        b <- dist uniformly (pure [1..nballs])
        -- It is more efficient to create if statements, rather than
        -- create an SExp(IntMap Color).
        let observe i =
              let true_color = balls ! i in
              if_ ((== i) <$> b) $
                dist (color `condition` observed_color) true_color
        foldr observe (return (pure color)) [1..maxBalls]

-- The model takes a list of observations as an argument
cballs_model :: [Color] -> Model Int
cballs_model obs = do
  balls  <- balls_prior maxBalls
  nballs <- dist uniformly (pure [1..maxBalls])
  let obs_number i =
        if_ ((== i) <$> nballs) $ balls_observe balls obs i
  foldr obs_number (return (pure())) [1..maxBalls]
  return nballs

(define (test/colored-balls _)
  (let1 (expected '(0.45765164162493 0.298386199220924 0.243962159154146))
        (⊦ (equal-approx? 1e-12)
           (map cons '(1 2 3) expected)
           (exact-posterior (cballs-model 3 (make-list 4 'blue))))
        (⊦ (equal-approx? 1e-12) expected (cballs-closed-form 3 4)))
  (⊦ (equal-approx? 1e-12)
     '(0.41196368095257
       0.209728624754566
       0.120691976674053
       0.0801847685132436
       0.0590318166419537
       0.046603790739908
       0.0386298621005105
       0.0331654796231948)
     (cballs-closed-form 8 10)))

((eta 0.013) (memory #(8396592 3867096 1048576)) (stdout "") (stderr ""))

(define pswitch
  (do/monad
    (← c (dist bern 0.5))
    (← m (dist uniformly '(0 1)))
    (if_ c
         (dist categorical (lift (λ (mv) `((,mv . 0.9) (,(- 1 mv) . 0.1))) m))
         (diracN (lift identity m)))
    (return c)))

-- (new, plan §6.8; no upstream source)
c ~ bern 0.5
m ~ uniformly [0,1]
if_ c (y ~ categorical [(m,.9),(1-m,.1)])
      (diracN (id <$> m))
return c

(define (test/pswitch _)
  (⊦ (equal-approx? 1e-12) '((#f . 0.5) (#t . 0.5)) (exact-posterior pswitch)))

((eta 0.0) (memory #(8396592 3869328 1048576)) (stdout "") (stderr ""))

(define (hansei-grass)
  (let* ((rain (probcc-coin 0.3))
         (sprinkler (probcc-coin 0.5))
         (grass-is-wet
           (or (and (probcc-coin 0.9) rain)
               (and (probcc-coin 0.8) sprinkler)
               (probcc-coin 0.1))))
    (probcc-when grass-is-wet rain)))

(define (hansei-girl)
  (let* ((girl1 (probcc-coin 0.5)) (girl2 (probcc-coin 0.5)))
    (probcc-when (probcc-coin (if (or girl1 girl2) 0.99 0.01)) girl1)))

(define (hansei-pcond4)
  (let1 (x (probcc-distribution '((1 0.5) (2 0.5))))
        (if (= x 1)
          x
          (probcc-when
            (probcc-coin 0.5)
            (probcc-distribution '((20 0.5) (21 0.5)))))))

(define (test/hansei _)
  (⊦ (equal-approx? 1e-12)
     (exact-posterior grass)
     (hansei->histogram (probcc-reify/exact/a hansei-grass)))
  (⊦ (equal-approx? 1e-12)
     (exact-posterior girl)
     (hansei->histogram (probcc-reify/exact/a hansei-girl)))
  (⊦ (equal-approx? 1e-12)
     (exact-posterior pcond4)
     (hansei->histogram (probcc-reify/exact/a hansei-pcond4))))

((eta 0.0) (memory #(8396592 3871216 1048576)) (stdout "") (stderr ""))

(define nested-dispatch
  (do/monad
    (← x (dist uniformly '(1 2 3)))
    (if_ (lift (cut > <> 1) x)
         (if_ (lift (cut > <> 2) x)
              (dist (condition #t bern) 0.1)
              (dist (condition #t bern) 0.3))
         (dist (condition #t bern) 0.6))
    (return x)))

(define (liar)
  (make-distribution
    'liar
    #:log-density
    (λ (x) (log (if x 0.9 0.1)))
    #:sample
    (λ (src)
       (let1 (x (< (rng-double! src 0.0 1.0) 0.9))
             (values x (log (if x 0.9 0.1)))))
    #:support
    (τ '((#t . 0.5) (#f . 0.5)))))

(define liar-or
  (do/monad
    (← a (dist liar))
    (← b (dist bern 0.3))
    (dist (condition #t bern) (lift (λ (av bv) (if (or av bv) 0.8 0.2)) a b))
    (return a)))

(define liar-or/nan
  (do/monad
    (← a (dist liar))
    (← b (dist bern 0.3))
    (dist (condition #t bern) (lift (λ (av bv) (if (or av bv) 0.8 0.2)) a b))
    (← k (dist uniformly '(0 1 2)))
    (dist (condition 0.0 normal) 0.0 (lift exact->inexact k))
    (return a)))

(define plus-inf-state
  (do/monad
    (← b (dist uniformly '(1 2 3)))
    (dist (condition 0.0 uniform)
          0.0
          (lift (λ (bv) (case bv ((1) 0.0) ((2) 1.0) (else 2.0))) b))
    (return b)))

(define (test/kernel-residual _)
  (define models
    `((grass unquote grass)
      (girl unquote girl)
      (pcorr unquote (pcorr/pair 0.1))
      (pcond1 unquote pcond1)
      (pcond2 unquote pcond2)
      (pcond4 unquote pcond4)
      (post-bias unquote (post-bias 1.0))
      (colored-balls-3x2 unquote (cballs-model 3 '(blue blue)))
      (pswitch unquote pswitch)
      (nested-dispatch unquote nested-dispatch)))
  (⊦ equal?
     '()
     (filter-map
       (Λ ((,name unquote model))
          (let1 (residual (kernel-residual model))
                (and (not (< residual 1e-12)) (cons name residual))))
       models))
  (⊦ = 0.0 (kernel-residual (if_ #f (dist bern 0.5) (return 0))))
  (⊦ > 1e-12 (kernel-residual (dist (uniformly '(1 1 2)))))
  (⊦ < 0.001 (kernel-residual liar-or))
  (⊦ < 0.001 (kernel-residual liar-or/nan))
  (⊦ > 1e-12 (kernel-residual nan-path))
  (⊦ > 1e-12 (kernel-residual plus-inf-state))
  (⊦ (equal-approx? 1e-12)
     `((2 unquote (/ 2.0 3.0)) (3 unquote (/ 1.0 3.0)))
     (exact-posterior plus-inf-state))
  (⊦ < 0.01 (kernel-residual pswitch #:semantics 'upstream))
  (⊦ > 1e-12 (kernel-residual pswitch #:semantics 'extended))
  (⊦= 'support (hakaru10-kind (τ (kernel-residual (dist bern/flip 0.5)))))
  (⊦= 'construction
      (hakaru10-kind
        (τ (kernel-residual
             (do/monad
               (← b (dist bern 0.5))
               (dist (condition #t bern) 0.0)
               (return b)))))))

((eta 0.448) (memory #(16793184 3888576 1048576)) (stdout "") (stderr ""))

(do/monad
  (← x (dist d args))
  (diracN e)
  (if_ test then-model else-model)
  (return x))

(define (test/three-interpretations _)
  (define models
    `((grass ,grass)
      (girl ,girl)
      (pcorr ,(pcorr/pair 0.1) #:chain-samples 80000)
      (ph2 ,ph2)
      (ph21 ,ph21)
      (pcond1 ,pcond1)
      (pcond2 ,pcond2)
      (pcond4 ,pcond4)
      (post-bias ,(post-bias 1.0))
      (colored-balls-3x4
        ,(cballs-model 3 (make-list 4 'blue))
        #:chain-samples
        80000)
      (pswitch ,pswitch #:semantics extended)))
  (define (disagreement
           name
           model
           #!key
           (semantics 'upstream)
           (chain-samples 20000))
    (let ((exact (exact-posterior model))
          (weighted (likelihood-weighting 20000 model))
          (chain (mcmc-posterior chain-samples model #:semantics semantics)))
      (and (not (and ((histograms-within 0.02) exact weighted)
                     ((histograms-within 0.03) exact chain)))
           (list name exact weighted chain))))
  (⊦ equal? '() (filter-map (cut apply disagreement <>) models)))

((eta 3.154) (memory #(16793184 3865240 1048576)) (stdout "") (stderr ""))

(define nested
  (do/monad
    (← c (dist bern 0.5))
    (← d (dist bern 0.4))
    (← z (dist normal 0.0 1.0))
    (if_ c
         (do/monad
           (← e (dist bern (lift (λ (zv) (if (> zv 0.0) 0.7 0.2)) z)))
           (if_ (lift (λ (ev dv) (and ev dv)) e d)
                (dist (condition 0.5 normal) z 1.0)
                (>> (dist (condition #t bern) 0.3)
                    (diracN (lift (cut * 2.0 <>) z)))))
         (if_ d (dist normal z 2.0) (dist gamma 2.0 1.0)))))

(define (phier n)
  (fold (λ (_ m) (do/monad (← x m) (dist normal x 3.0)))
        (dist normal 0.0 1.0)
        (iota n)))

(define (normals means sd) (traverse/monad (cut dist normal <> sd) means))

(define (dot xs ws)
  (reduce-right sexp+ 0.0 (map (λ (x w) (lift (cut * <> x) w)) xs ws)))

(define (regression k data)
  (do/monad
    (← w-mean (normals (make-list k 0.0) 2.0))
    (← w (normals w-mean 1.0))
    (← noise-sd (<$> (cut lift (cut / 1.0 <>) <>) (dist gamma 0.5 0.5)))
    (traverse/monad
      (Λ ((,xs unquote y)) (dist (condition y normal) (dot xs w) noise-sd))
      data)
    (return (collect w))))

phier n =
  (iterate (\m -> do {x<-m; dist normal x 3}) $ dist normal 0 1) !! n

normals :: [SExp Double] -> SExp Double -> MCMCM [SExp Double]
normals means std = mapM (\m -> dist normal m std) means


dot :: [Double] -> [SExp Double] -> SExp Double
dot xs ays = foldr1 (+) $ zipWith (\x ay -> (*x) <$> ay) xs ays

(define (test/construction-order _)
  (define models
    `((grass unquote grass)
      (pcond3 unquote pcond3)
      (pcond4 unquote pcond4)
      (pswitch unquote pswitch)
      (nested unquote nested)
      (colored-balls-8x10 unquote (cballs-model 8 (make-list 10 'blue)))
      (phier-10 unquote (phier 10))
      (regression-5x500 unquote (regression 5 (regression-data 5 500)))))
  (⊦ equal?
     (map (λ (entry) (list (car entry) #t #t)) models)
     (map (Λ ((,name unquote model))
             (list name
                   (construction-check model #:semantics 'upstream)
                   (construction-check model #:semantics 'extended)))
          models)))

((eta 0.132) (memory #(16793184 5489272 1048576)) (stdout "") (stderr ""))

(((C #<procedure (a3773)>) 0.35)
 ((C #<procedure (a3773)>) 0.35)
 ((C #<procedure (a3773)>) 0.15)
 ((C #<procedure (a3773)>) 0.15))

(define (test/hansei-format _)
  (let1 (shallow (probcc-explore 0 (model->search-tree grass)))
        (⊨ (any (λ (entry) (eq? 'C (caar entry))) shallow)))
  (⊦ (equal-approx? 1e-12)
     (exact-posterior grass)
     (hansei->histogram (probcc-explore +inf.0 (model->search-tree grass))))
  (for-each
    (Λ ((,name unquote m))
       (⊦ (equal-approx? 1e-12)
          (list name (exact-posterior m))
          (list name
                (hansei->histogram
                  (probcc-explore
                    0
                    (explore/reversed-levels (model->search-tree m)))))))
    `((grass unquote grass) (pcond4 unquote pcond4))))

((eta 0.001) (memory #(16793184 3865928 1048576)) (stdout "") (stderr ""))

(define (test/errors _)
  (⊦ equal?
     '((builder . construction)
       (exact . construction)
       (forward . construction))
     (kinds-under-interpretations
       (do/monad
         (← x (dist bern 0.5))
         (return (lift (λ (xv) (dist bern 0.5)) x)))))
  (⊦ equal?
     '((builder . branch) (exact . branch) (forward . branch))
     (kinds-under-interpretations
       (do/monad
         (← x (dist uniformly '(0 1)))
         (if_ x (return 1) (return 2))))))

((eta 0.0) (memory #(16793184 3873232 1048576)) (stdout "") (stderr ""))

{-
biased_coin_ill_typed_too c1 c2 = do
  b <- dist bern 0.5
  dist (b  `condition` bern) 0.4
  return b
-}

(define observes-its-argument
  (do/monad
    (← x (dist bern 0.5))
    (dist (λ (xv) (condition xv (bern 0.8))) x)
    (return x)))

(define (nested-uniformly choices)
  (do/monad
    (← xs (dist uniformly choices))
    (if_ (lift pair? xs) (dist uniformly xs) (return 'none))))

(define (nested-uniformly/ill-typed choices)
  (do/monad
    (← xs (dist uniformly choices))
    (if_ (lift pair? xs) (dist uniformly xs) (return 'none))
    (← y (dist bern 0.5))
    (dist (condition y bern) 0.5)
    (return xs)))

(define always-rejected
  (do/monad
    (← c (dist bern 0.5))
    (if_ #t
         (return c)
         (do/monad
           (← h (diracN (lift car '())))
           (← mu (diracN (lift cdr (pair h 1.0))))
           (dist normal mu 1.0)))))

(define guarded-result
  (do/monad
    (← c (dist bern 0.5))
    (← r (if_ c (return 'ok) (diracN (lift car '()))))
    (return c)))

(define guarded-result/mirror
  (do/monad
    (← c (dist bern 0.5))
    (← r (if_ (lift not c) (return 'ok) (diracN (lift car '()))))
    (return c)))

(define guarded-head
  (do/monad
    (← c (dist bern 0.5))
    (let xs (lift (λ (cv) (if cv '(1) '())) c))
    (if_ (lift not c)
         (>> (diracN (lift car xs)) (return 'then))
         (return 'else))))

(define guarded-head/mirror
  (do/monad
    (← c (dist bern 0.5))
    (let xs (lift (λ (cv) (if cv '() '(1))) c))
    (if_ c (>> (diracN (lift car xs)) (return 'then)) (return 'else))))

(define guarded-head/constant
  (do/monad
    (← c (dist bern 0.5))
    (← xs (diracN '()))
    (if_ (lift not c)
         (>> (diracN (lift car xs)) (return 'then))
         (return 'else))))

(define guarded-exit/constant
  (do/monad
    (← c (dist bern 0.5))
    (if_ (lift not c)
         (do/monad
           (← inner (if_ #t (diracN (lift car '())) (return 'x)))
           (return 'then))
         (return 'else))))

(define guarded-root/constant
  (do/monad
    (← c (dist bern 0.5))
    (if_ (lift not c)
         (>> (diracN (lift car '())) (return 'then))
         (return 'else))))

(define (test/static-graph _)
  (define observes-a-choice
    (do/monad (← x (dist bern 0.5)) (dist (condition x bern) 0.5)))
  (⊦= 'conditioning (hakaru10-kind (τ (exact-posterior observes-a-choice))))
  (⊦ equal?
     (make-list 4 'conditioning)
     (map (λ (run) (hakaru10-kind (τ (run observes-its-argument))))
          (list exact-posterior
                (cut likelihood-weighting 100 <>)
                (cut mcmC 50 <> #:check #t)
                kernel-residual)))
  (for-each
    (λ (choices)
       (⊦ (equal-approx? 1e-12)
          (list choices '((1 . 0.25) (2 . 0.25) (none . 0.5)))
          (list choices (exact-posterior (nested-uniformly choices))))
       (⊦ equal?
          (list choices 'conditioning)
          (list choices
                (hakaru10-kind
                  (τ (exact-posterior
                       (nested-uniformly/ill-typed choices)))))))
    '(((1 2) ()) (() (1 2))))
  (⊦= 'construction (hakaru10-kind (τ (exact-posterior always-rejected))))
  (⊦= 'support (hakaru10-kind (τ (exact-posterior grass #:max-paths 3))))
  (⊦ (equal-approx? 1e-12)
     (exact-posterior grass)
     (exact-posterior grass #:max-paths 4))
  (let1 (tree (model->search-tree grass #:max-paths 4))
        (⊦ (equal-approx? 1e-12)
           (exact-posterior grass)
           (hansei->histogram (probcc-explore +inf.0 tree)))
        (⊦= 'support (hakaru10-kind (τ (probcc-explore +inf.0 tree)))))
  (for-each
    (Λ ((,name unquote m))
       (⊦ equal?
          (list name 'exn 'exn 'exn)
          (list name
                (error-kind (τ (exact-posterior m)))
                (error-kind (τ (likelihood-weighting 100 m)))
                (error-kind (τ (mcmC 100 m))))))
    `((guarded-result unquote guarded-result)
      (guarded-result/mirror unquote guarded-result/mirror)
      (guarded-head unquote guarded-head)
      (guarded-head/mirror unquote guarded-head/mirror)
      (guarded-head/constant unquote guarded-head/constant)
      (guarded-exit/constant unquote guarded-exit/constant)
      (guarded-root/constant unquote guarded-root/constant)))
  (for-each
    (Λ ((,name unquote m))
       (⊦ equal?
          (list name 'else 'exn 'exn 'exn)
          (list name
                (car (mcmC 1 m))
                (error-kind (τ (mcmC 100 m #:semantics 'extended)))
                (error-kind (τ (mcmC 100 m #:check #t)))
                (error-kind (τ (kernel-residual m))))))
    `((guarded-head/constant unquote guarded-head/constant)
      (guarded-exit/constant unquote guarded-exit/constant)
      (guarded-root/constant unquote guarded-root/constant))))

((eta 0.022) (memory #(16793184 3873704 1048576)) (stdout "") (stderr ""))

(define two-levels
  (do/monad
    (← c (dist bern 0.4))
    (← d (dist uniformly '(0 1 2)))
    (if_ c
         (do/monad
           (← e (dist bern 0.7))
           (if_ e (diracN d) (diracN (lift (cut + 10 <>) d))))
         (do/monad
           (← y (dist categorical '((5 . 0.5) (6 . 0.5))))
           (dist (condition #t bern)
                 (lift (λ (yv dv) (if (= yv 5) 0.8 (/ (+ dv 1.0) 4.0))) y d))
           (return y)))))

((0 . 0.118143459915612)
 (1 . 0.118143459915612)
 (2 . 0.118143459915612)
 (5 . 0.30379746835443)
 (6 . 0.189873417721519)
 (10 . 0.0506329113924051)
 (11 . 0.0506329113924051)
 (12 . 0.0506329113924051))

(define (test/nested-branches _)
  (let1 (exact (exact-posterior two-levels))
        (⊦ (histograms-within 0.02) exact (mcmc-posterior 20000 two-levels))
        (⊦ (histograms-within 0.02)
           exact
           (mcmc-posterior 20000 two-levels #:semantics 'extended))
        (⊦ (histograms-within 0.02)
           exact
           (likelihood-weighting 20000 two-levels))))

((eta 0.289) (memory #(16793184 3877088 1048576)) (stdout "") (stderr ""))

(define mixng
  (do/monad
    (← x (dist normal 0.0 1.0))
    (if_ (lift (cut > <> 0.0) x)
         (dist normal 10.0 2.0)
         (dist gamma 3.0 (/ 1.0 3.0)))))

mixng = do
  x <- dist normal 0 1
  if_ ((>0) <$> x)
           (dist normal 10 2)
           (dist gamma 3 (1/3))

(define (test/continuous _)
  (⊦ (histograms-within 0.02)
     '((1 . 0.5) (20 . 0.25) (21 . 0.25))
     (likelihood-weighting 20000 pcond3))
  (⊦ (equal-approx? 0.4)
     5.5
     (fold (Λ ((,x unquote p) ,mean) (+ mean (* x p)))
           0.0
           (likelihood-weighting 20000 mixng)))
  (⊦ equal?
     '()
     (likelihood-weighting 100 (dist (condition 5.0 uniform) 0.0 1.0))))

((eta 0.291) (memory #(16793184 3876320 1048576)) (stdout "") (stderr ""))

(define (test/pre-allocation _)
  (⊦ (equal-approx? 1e-12)
     '(((#f . #f) . 0.25)
       ((#f . #t) . 0.25)
       ((#t . #f) . 0.25)
       ((#t . #t) . 0.25))
     (exact-posterior
       (do/monad
         (← coins (replicate/monad 2 (dist bern 0.5)))
         (let table (map cons '(a b) coins))
         (return (pair (alist-ref 'a table) (alist-ref 'b table))))))
  (⊦ (equal-approx? 1e-12)
     (map cons '(1 2 3) (cballs-closed-form 3 4))
     (exact-posterior (cballs-model 3 (make-list 4 'blue)))))

((eta 0.014) (memory #(16793184 3992488 1048576)) (stdout "") (stderr ""))

data Color = Blue | Green deriving (Eq, Show)

opposite_color :: Color -> Color
opposite_color Blue  = Green
opposite_color Green = Blue

-- Distribution for the observing of a ball, with the failure rate 20%
observed_color color =
  categorical [(color,0.8),(opposite_color color,0.2)]

maxBalls :: Int
maxBalls = 8

-- Create the prior for the max amount of balls that may be in the urn
balls_prior n = do
  balls <- sequence . replicate n $ dist uniformly (pure [Blue,Green])
  return $ M.fromList $ zip [1..] balls

-- Submodel: Given balls color and the observations, perform conditioning
balls_observe :: M.IntMap (SExp Color) -> [Color] -> Int -> Model ()
balls_observe balls obs nballs = mapM_ observe_ball obs >> return (pure ())
 where
  -- The same ball can be drawn twice as we return the drawn balls
  -- into the urn. A ball keeps its true color, no matter how many
  -- times it is drawn.
   observe_ball color = do
        b <- dist uniformly (pure [1..nballs])
        -- It is more efficient to create if statements, rather than
        -- create an SExp(IntMap Color).
        let observe i =
              let true_color = balls ! i in
              if_ ((== i) <$> b) $
                dist (color `condition` observed_color) true_color
        foldr observe (return (pure color)) [1..maxBalls]

-- The model takes a list of observations as an argument
cballs_model :: [Color] -> Model Int
cballs_model obs = do
  balls  <- balls_prior maxBalls
  nballs <- dist uniformly (pure [1..maxBalls])
  let obs_number i =
        if_ ((== i) <$> nballs) $ balls_observe balls obs i
  foldr obs_number (return (pure())) [1..maxBalls]
  return nballs

cbr = hist $ mcmC 10000 (cballs_model (replicate 10 Blue))

-- fromList [(1,4265),(2,2203),(3,1244),(4,805),(5,486),(6,291),(7,565),(8,141)]
-- Good agreement

(begin
  (define (opposite-color color) (case color ((blue) 'green) ((green) 'blue)))
  (define (observed-color color)
    (categorical `((,color . 0.8) (,(opposite-color color) . 0.2))))
  (define (balls-prior n)
    (do/monad
      (← balls (replicate/monad n (dist uniformly '(blue green))))
      (return (map cons (iota n 1) balls))))
  (define (dispatch x keys branch otherwise)
    (fold-right
      (λ (key rest) (if_ (lift (cut equal? <> key) x) (branch key) rest))
      otherwise
      keys))
  (define (balls-observe balls obs nballs max-balls)
    (>> (for-each/monad
          (λ (color)
             (do/monad
               (← b (dist uniformly (iota nballs 1)))
               (dispatch
                 b
                 (iota max-balls 1)
                 (λ (i)
                    (dist (condition color observed-color)
                          (alist-ref i balls)))
                 (return color))))
          obs)
        (return (void))))
  (define (cballs-model obs #!optional (max-balls 8))
    (do/monad
      (← balls (balls-prior max-balls))
      (← nballs (dist uniformly (iota max-balls 1)))
      (dispatch
        nballs
        (iota max-balls 1)
        (λ (n) (balls-observe balls obs n max-balls))
        (return (void)))
      (return nballs))))

(define (cballs-closed-form max-balls obs)
  (define (binomial n k)
    (fold (λ (i c) (/ (* c (- n i)) (add1 i))) 1 (iota k)))
  (define (seen color n k)
    (let1 (blue (/ (+ (* 0.8 k) (* 0.2 (- n k))) n))
          (if (eq? color 'blue) blue (- 1.0 blue))))
  (define (likelihood n)
    (fold (λ (k acc)
             (+ acc
                (* (binomial n k)
                   (expt 2.0 (- n))
                   (fold (λ (color p) (* p (seen color n k))) 1.0 obs))))
          0.0
          (iota (add1 n))))
  (let* ((ls (map likelihood (iota max-balls 1))) (z (fold + 0.0 ls)))
    (map (λ (n l) (cons n (/ l z))) (iota max-balls 1) ls)))

((1 0.41196368095257)
 (2 0.209728624754566)
 (3 0.120691976674053)
 (4 0.0801847685132436)
 (5 0.0590318166419537)
 (6 0.046603790739908)
 (7 0.0386298621005105)
 (8 0.0331654796231948))

((1 . 0.457651641624931) (2 . 0.298386199220923) (3 . 0.243962159154146))

(define (test/colored-balls _))

((eta 2.51) (memory #(12582912 3243232 1048576)) (stdout "") (stderr ""))

{-
Problem with the following model: nballs and b below are strongly correlated.
Hence if nballs is drawn as 5 and b is drawn as, say, 3, all the
proposals to resample nbalss as 1 or 2 will be rejected.

-- The model takes an array of observations as an argument
cballs_model :: [Color] -> Model Int
cballs_model obs = do
  balls  <- balls_prior maxBalls
  nballs <- dist uniformly (pure [1..maxBalls])
  -- The same ball can be drawn twice as we return the drawn balls
  -- into the urn. A ball keeps its true color, no matter how many
  -- times it is drawn.
  let observe_ball color = do
        b <- dist uniformly ((\n -> [1..n]) <$> nballs)
        -- It is more efficient to create if statements, rather than
        -- create an SExp(IntMap Color).
        let observe i =
              let true_color = balls ! i in
              if_ ((== i) <$> b) $
                dist (color `condition` categorical)
                     ((\bc -> [(bc,0.8),(opposite_color bc,0.2)]) <$>
                      true_color)
        foldr observe (return (pure color)) [1..maxBalls]
  mapM_ observe_ball obs
  return nballs
-}

(define (cballs-model/direct obs #!optional (max-balls 8))
  (do/monad
    (← balls (balls-prior max-balls))
    (← nballs (dist uniformly (iota max-balls 1)))
    (for-each/monad
      (λ (color)
         (do/monad
           (← b (dist uniformly (lift (λ (n) (iota n 1)) nballs)))
           (dispatch
             b
             (iota max-balls 1)
             (λ (i)
                (dist (condition color observed-color) (alist-ref i balls)))
             (return color))))
      obs)
    (return nballs)))

((1 . 0.457651641624931) (2 . 0.298386199220923) (3 . 0.243962159154146))

26

302

((2 0.41845)
 (3 0.01495)
 (4 0.1611)
 (5 0.22275)
 (6 0.0668)
 (7 0.08425)
 (8 0.0317))

(define (test/colored-balls/direct _))

((eta 1.719) (memory #(12582912 4205872 1048576)) (stdout "") (stderr ""))

neighbors :: CellNum -> [CellNum]
neighbors n = filter (\i -> dist n i <= d_max) [1..nCells]
 where
  dist n m = let (ix,iy) = divMod (n-1) gridDim
                 (jx,jy) = divMod (m-1) gridDim
             in (ix-jx)^2+(iy-jy)^2
  d_max = 18

(begin
  (define (cell->xy n #!optional (dim grid-dim))
    (let1 (m (sub1 n)) (cons (quotient m dim) (remainder m dim))))
  (define squared-distance
    (Λ ((,ix unquote iy) (,jx unquote jy))
       (+ (* (- ix jx) (- ix jx)) (* (- iy jy) (- iy jy)))))
  (define (neighbors n #!optional (dim grid-dim))
    (let1 (from (cell->xy n dim))
          (filter
            (λ (m) (<= (squared-distance from (cell->xy m dim)) 18))
            (iota (* dim dim) 1)))))

(1 2 3 4 5 11 12 13 14 15 21 22 23 24 31 32 33 34 41 42)

(define (test/birds/neighbors _)
  (⊦ equal? '((0 . 0) (0 . 3) (3 . 0) (3 . 3)) (map cell->xy '(1 4 13 16)))
  (⊦ equal?
     '((0 . 0) (0 . 9) (9 . 0) (9 . 9))
     (map (cut cell->xy <> 10) '(1 10 91 100)))
  (⊨ (every (λ (n) (equal? (iota 16 1) (neighbors n))) (iota 16 1)))
  (⊦ equal?
     '(1 2 3 4 5 11 12 13 14 15 21 22 23 24 31 32 33 34 41 42)
     (neighbors 1 10))
  (⊦ = 20 (length (neighbors 1 10))))

((eta 0.002) (memory #(12582912 3250680 1048576)) (stdout "") (stderr ""))

-- where is the single bird on the day d of the year y
birdpos :: ObsData -> Year -> Day -> Maybe CellNum
birdpos obs y d = case [i | i <- [1..nCells], obs ! (y,d,i) > 0] of
  [i] -> Just i
  _   -> Nothing

phi :: Year -> Day -> CellNum -> CellNum -> Features -> [Double] -> Double
phi y d from to features bs = sum $ zipWith (*) fs bs
  where
    fs = [ realToFrac $ features ! (y, d, from, to, i) | i <- [1..nFeatures]]

-- multinomial (0; a,b,c) --> 0
-- multinomial (1,theta(i,1),,,theta(i,16))  -->
--     categorical [(1,theta(i,1)),(2,(theta(i,2))),... (16,theta(i,16))]
theta :: Year -> Day -> CellNum -> CellNum -> Features -> [Double] -> Double
theta y d from to features bs = exp $ phi y d from to features bs

(begin
  (define (phi fs bs) (fold + 0.0 (map * fs bs)))
  (define (theta fs bs) (exp (phi fs bs)))
  (define ((transitions features year day) n bs)
    (map (λ (j) (cons j (theta (features year day n j) bs))) (neighbors n)))
  (define (birdpos obs year day)
    (match/first
      (filter (λ (i) (positive? (obs year day i))) (iota n-cells 1))
      ((,i) i)
      (else #f))))

  logDensity (Poisson lambda expl) x
    | x == 0      = -lambda
    | lambda == 0 = minusInf
    | x == 1      = log lambda - lambda
    | x == 2      = log (lambda*lambda/2) - lambda

(define (test/birds/transitions _))

((eta 0.001) (memory #(12582912 3263336 1048576)) (stdout "") (stderr ""))

model :: Features -> ObsData -> Model [Double]
model features obs = do
  b1  <- abs <$> dist normal 0 10 -- see Statement of Problem, p9
  b2  <- abs <$> dist normal 0 10 -- see Statement of Problem, p9
  b3  <- abs <$> dist normal 0 10 -- see Statement of Problem, p9
  b4  <- abs <$> dist normal 0 10 -- see Statement of Problem, p9
  let bs = collect [b1,b2,b3,b4]
  forM_ [1..nYears] $ \year -> do -- each year
    foldM_ (\n day -> do -- one bird exists in cell n on each day(1..nDays-1)
      let trans n bs = [ (j,theta year day n j features bs) | j<-neighbors n]
      -- n' cell where the bird is on day+1
      let day' = day + 1
      {-
      n' <- case birdpos obs year day' of
          Just nreal -> dist (nreal `condition` categorical) $
                        liftA2 trans n bs
          Nothing    -> dist categorical $ liftA2 trans n bs
      -}
      n' <- dist categorical $ liftA2 trans n bs
      forM_ [1..nCells] $ \i ->
        dist (fromIntegral (obs ! (year,day',i)) `condition` poisson)
           ((\nv -> if i==nv then 1 else 0) <$> n')
      return n'
      ) (1 :: SExp Int) [1..nDays-1]
  return bs

(define (birds features obs #!key (years 30) (days 20) (pin #t))
  (do/monad
    (← betas
       (replicate/monad
         n-features
         (<$> (cut lift abs <>) (dist normal 0.0 10.0))))
    (let bs (collect betas))
    (for-each/monad
      (λ (year)
         (fold/monad
           (λ (n day)
              (do/monad
                (let day′ (add1 day))
                (let trans (lift (transitions features year day) n bs))
                (let nreal (and pin (birdpos obs year day′)))
                (← n′
                   (dist (if nreal (condition nreal categorical) categorical)
                         trans))
                (for-each/monad
                  (λ (i)
                     (dist (condition (obs year day′ i) poisson)
                           (lift (λ (nv) (if (= i nv) 1 0)) n′)))
                  (iota n-cells 1))
                (return n′)))
           1
           (iota (sub1 days) 1)))
      (iota years 1))
    (return bs)))

(begin
  (define (synthetic-features years days #!key (seed 17))
    (define (lognpdf x mu sigma)
      (/ (exp (- (/ (expt (- (log x) mu) 2) (* 2.0 sigma sigma))))
         (* x sigma (sqrt (* 2.0 pi)))))
    (define (cosine u v)
      (let ((nu (sqrt (squared-distance u '(0 . 0))))
            (nv (sqrt (squared-distance v '(0 . 0)))))
        (if (or (zero? nu) (zero? nv))
          0.0
          (/ (+ (* (car u) (car v)) (* (cdr u) (cdr v))) (* nu nv)))))
    (define (vector-to p q) (cons (- (car q) (car p)) (- (cdr q) (cdr p))))
    (let* ((src (make-stdgen-source seed))
           (winds (list->vector
                    (map-in-order
                      (λ (_)
                         (receive
                           (angle ll)
                           (distribution-sample (uniform 0.0 (* 2.0 pi)) src)
                           (cons (cos angle) (sin angle))))
                      (iota (* years days)))))
           (top-right (cell->xy n-cells)))
      (memoize
        (λ (year day from to)
           (if (= from to)
             '(0.0 0.0 0.0 2.0)
             (let* ((p (cell->xy from)) (v (vector-to p (cell->xy to))))
               (list (* 4.0
                        (lognpdf (sqrt (squared-distance v '(0 . 0))) 1.0 1.0))
                     (cosine v (vector-to p top-right))
                     (cosine
                       v
                       (vector-ref winds (+ (* (sub1 year) days) (sub1 day))))
                     0.0)))))))
  (define (bird-flights features bs #!key (years 30) (days 20))
    (<$> collect
         (traverse/monad
           (λ (year)
              (<$> (λ (path) (collect (reverse path)))
                   (fold/monad
                     (λ (path day)
                        (<$> (cut cons <> path)
                             (dist categorical
                                   (lift (transitions features year day)
                                         (car path)
                                         bs))))
                     '(1)
                     (iota (sub1 days) 1))))
           (iota years 1))))
  (define ((indicator-counts paths missed?) year day cell)
    (if (and (not (missed? day))
             (= cell (list-ref (list-ref paths (sub1 year)) (sub1 day))))
      1
      0)))

((1 6 8 10 12 12) (1 11 8 15 16 13))

2605

(4.12945104050349 9.82328450023311 1.3251745427761 2.08709026263613)

(7.87600232652835 8.15640125395766 8.22387046673009 8.15729590109516)

(define (test/birds/synthetic _))

((eta 1.466) (memory #(12582912 3609088 1048576)) (stdout "") (stderr ""))

type RegrDatum = ([Double],Double) -- Xs and Y

model :: [RegrDatum] -> Model [Double]
model xsy = do
  let mu = replicate dimK 0
  w_mean   <- normals mu 2
  w        <- normals w_mean 1
  noise_sd <- (1/) <$> dist gamma 0.5 0.5
  let make_cond (xs,y) = dist (y `condition` normal) (dot xs w) noise_sd
  mapM make_cond xsy
  return $ collect w

-- normals means std
-- returns several independently distributed normal random variables
-- whose means are given by the list means, with the same std
normals :: [SExp Double] -> SExp Double -> MCMCM [SExp Double]
normals means std = mapM (\m -> dist normal m std) means


dot :: [Double] -> [SExp Double] -> SExp Double
dot xs ays = foldr1 (+) $ zipWith (\x ay -> (*x) <$> ay) xs ays

collect :: [SExp a] -> SExp [a]
collect = sequenceA

(begin
  (define (normals means sd) (traverse/monad (λ (m) (dist normal m sd)) means))
  (define (dot xs ays) (reduce-right sexp+ 0.0 (map sexp* ays xs)))
  (define (regression data)
    (do/monad
      (← w-mean (normals (make-list (length (car (car data))) 0.0) 2.0))
      (← w (normals w-mean 1.0))
      (← inverse-noise-sd (dist gamma 0.5 0.5))
      (let noise-sd (lift (cut / 1.0 <>) inverse-noise-sd))
      (for-each/monad
        (Λ ((,xs unquote y)) (dist (condition y normal) (dot xs w) noise-sd))
        data)
      (return (collect w)))))

(define (regression-data xss w sd)
  (<$> collect
       (traverse/monad
         (λ (xs)
            (<$> (cut lift (cut cons xs <>) <>) (dist normal (dot xs w) sd)))
         xss)))

(1.60859970223401 -0.932442673713944)

(1.62928672635352 -0.926548718746949)

(define (least-squares/2 data)
  (define (sum f) (fold (Λ ((,xs unquote y) ,acc) (+ acc (f xs y))) 0.0 data))
  (let* ((a (sum (λ (xs y) (* (first xs) (first xs)))))
         (b (sum (λ (xs y) (* (first xs) (second xs)))))
         (d (sum (λ (xs y) (* (second xs) (second xs)))))
         (u (sum (λ (xs y) (* (first xs) y))))
         (v (sum (λ (xs y) (* (second xs) y))))
         (det (- (* a d) (* b b))))
    (list (/ (- (* d u) (* b v)) det) (/ (- (* a v) (* b u)) det))))

(define (test/regression _))

((eta 0.708) (memory #(12582912 3745624 1048576)) (stdout "") (stderr ""))

_ =
 let n = 10000 in plot (tabulate 0.2 $ mcmC (n `div` 2) (dist normal 10 2) ++ mcmC (n `div` 2) (dist gamma 3 (1/3)))

86

5.42921999999999

(define (test/plot-mixture _))

((eta 0.024) (memory #(12582912 3278520 1048576)) (stdout "") (stderr ""))