/* 

  Ball box in Picat.

  From https://bitbucket.org/pedrozudo/hal_problog/src/master/examples/ball_box.pl
  """
  beta(2,2)~b.

  P::p:- b~=P.

  box(1):-p.
  box(2):- \+p.

  1/4::ball(X, red);3/4::ball(X, white):- box(1).
  3/4::ball(X, red);1/4::ball(X, white):- box(2).


  evidence(ball(1,red)).
  evidence(ball(2,red)).

  :-free(b).
  query(density(b)).
  """

  What I understand it's a common pick-a-ball-from-a-random-box-problem:
  - There are two boxes (box 1 and box 2)
  - box 1 contains 1 red ball and 3 white balls
  - box 2 contains 3 red balls and 1 white ball
  - one random box is selected 
  - one ball is picked from that box (and returned I suppose): it's a red ball
  - again one ball is picked from the same box: it's also red
  - Question: which box was selected?
  - Answer: It's box 2 with a probability of about 90%

  Here we test for 2..4 number of (red) balls.

  Cf my Gamble model gamble_ball_box.rkt

  This program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Picat page: http://www.hakank.org/picat/

*/

import ppl_distributions, ppl_utils.
import util.

main => go.

/*

  num_balls = 2
  var : b
  Probabilities (truncated):
  0.990413359579809: 0.0001073652566030
  0.989016663137676: 0.0001073652566030
  0.985635958348219: 0.0001073652566030
  0.97956221108587: 0.0001073652566030
  .........
  0.011109509821323: 0.0001073652566030
  0.010125761905214: 0.0001073652566030
  0.009415933317507: 0.0001073652566030
  0.001191874912461: 0.0001073652566030
  mean = 0.420696

  var : box
  Probabilities:
  2: 0.8978956409705819
  1: 0.1021043590294181
  mean = 1.8979

  var : p
  Probabilities:
  false: 0.8978956409705819
  true: 0.1021043590294181
  mean = [false = 0.897896,true = 0.102104]


  num_balls = 3
  var : b
  Probabilities (truncated):
  0.976917680146724: 0.0001520912547529
  0.969349870525874: 0.0001520912547529
  0.968125009439539: 0.0001520912547529
  0.966003409974784: 0.0001520912547529
  .........
  0.012192826458689: 0.0001520912547529
  0.010156212301358: 0.0001520912547529
  0.008468844347431: 0.0001520912547529
  0.002656953686673: 0.0001520912547529
  mean = 0.409226

  var : box
  Probabilities:
  2: 0.9680608365019011
  1: 0.0319391634980989
  mean = 1.96806

  var : p
  Probabilities:
  false: 0.9680608365019011
  true: 0.0319391634980989
  mean = [false = 0.968061,true = 0.0319392]


  num_balls = 4
  var : b
  Probabilities (truncated):
  0.961474727105557: 0.0002099958000840
  0.957882259287197: 0.0002099958000840
  0.957463117937563: 0.0002099958000840
  0.953762626232381: 0.0002099958000840
  .........
  0.009719601282469: 0.0002099958000840
  0.00948718762693: 0.0002099958000840
  0.008895821618289: 0.0002099958000840
  0.006597693527563: 0.0002099958000840
  mean = 0.401015

  var : box
  Probabilities:
  2: 0.9865602687946241
  1: 0.0134397312053759
  mean = 1.98656

  var : p
  Probabilities:
  false: 0.9865602687946241
  true: 0.0134397312053759
  mean = [false = 0.98656,true = 0.0134397]

*/
go ?=>
  member(NumBalls,2..4),
  println(num_balls=NumBalls),
  reset_store,
  run_model(30_000,$model(NumBalls),[show_probs_trunc,mean]),
  nl,
  fail,
  nl.
go => true.
  
  
model(NumBalls) =>
  B = beta_dist(2,2),
  P = flip(B),
  Box = condt(P,1,2),

  Balls = [ cond(Box == 1,
                    categorical([1/4,3/4],["red","white"]),
                    categorical([3/4,1/4],["red","white"]))
               : _ in 1..NumBalls],

  % We observe that all balls are red
  foreach(Ball in 1..NumBalls)
    observe(Balls[Ball] == "red"),
  end,
 
  if observed_ok then
    add("b",B),
    add("p",P),
    add("box",Box)
  end.