/* 

  Urn puzzle in Picat.

  From Daniel Litt (Sep 2, 2024)
  https://x.com/littmath/status/1830404515788427575
  """
  You have 3 urns, each containing 100 balls. In the first two, 99 of the balls 
  are red and the last is green. In the third urn, all 100 balls are red. 
  You choose one of the urns at random and remove 99 randomly chosen balls from it; 
  they’re all red. The last is probably
  - Green 
  - Red
  - Equally likely
  """

  The last is probably red and the urn picked was probably the third urn.


  Cf my Gamble model gamble_urn_puzzle4.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.

/*

  var : left
  Probabilities:
  red: 0.9840283939662822
  green: 0.0159716060337178
  mean = [red = 0.984028,green = 0.0159716]

  var : urn
  Probabilities:
  3: 0.9840283939662822
  2: 0.0085773439810707
  1: 0.0073942620526471
  mean = 2.97663

*/
go ?=>
  reset_store,
  run_model(10_000,$model,[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,
  % fail,
  nl.
go => true.
  
model() =>
  NumUrns = 3,
  Urn = random_integer1(NumUrns),

  % Urn 1: 99 red balls, 1 green
  % Urn 2: 99 red balls, 1 green
  % Urn 3: 100 red balls   
  BallsInUrn = case(Urn,
                    [
                    [1, [red : _ in 1..99] ++ [green]],
                    [2, [red : _ in 1..99] ++ [green]],
                    [3, [red : _ in 1..100]]
                    ]),

  % Remove 99 balls and observe that all are red.
  % What color is the last ball?
  NinetynineRed = [red : _ in 1..99],
  Shuffle = shuffle(BallsInUrn),
  observe(Shuffle[1..99] == NinetynineRed),
  
  Left = Shuffle.last,
  if observed_ok then
    add("urn",Urn),
    add("left",Left)
  end.