/* 

  Urn puzzle in Picat.

  From "Perplexing the Web, One Probability Puzzle at a Time" 
  https://www.quantamagazine.org/perplexing-the-web-one-probability-puzzle-at-a-time-20240829/?mc_cid=94caee8978
  (about Daniel Litt)
  """
  Puzzle 1

  You have 100 urns, each contain 99 ball. In 99 of the urns, one ball is red,
  and the rest are green. In the last, all 99 balls are red. You pick an urn 
  at random and draw a ball. It's red. Now throw it away. 
  The next ball you draw from the SAME urn is probably 
  * Red
  * Green
  * Equally likely
  """

  This prpblem was tweeted by Daniel Litt (Jan 29, 2024)
  https://x.com/littmath/status/1751971133723656585
  """
  ... 
  Red: 25.7%
  Green: 28.1%
  Equally likely: 32.3%
  Don’t know/see results: 13.9%
  """

  It's equally likely that ball2 is red or green.
  Also, note that it's even chance that we drew from the last (100'th) urn with 
  the 99 red balls.

  Cf my Gamble model gamble_urn_puzzle3.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 : ball2
  Probabilities:
  green: 0.5060449050086355
  red: 0.4939550949913644
  mean = [green = 0.506045,red = 0.493955]

  var : urn
  Probabilities (truncated):
  100: 0.4939550949913644
  27: 0.0138169257340242
  99: 0.0103626943005181
  85: 0.0103626943005181
  .........
  18: 0.0017271157167530
  5: 0.0017271157167530
  4: 0.0017271157167530
  3: 0.0017271157167530
  mean = 75.4974

*/
go ?=>
  reset_store,
  run_model(100_000,$model,[show_probs_trunc,mean]),
  nl,
  show_store_lengths,
  % fail,
  nl.
go => true.
  
  
model() =>
  NumUrns = 100,
  Colors = [red,green],

  Urns = [cond(U==100,
          categorical([99/99,0/99],Colors), % last urn
          categorical([1/99,98/99],Colors)) % the rest of the urns
          : U in 1..NumUrns],

  Urn = random_integer1(NumUrns),

  Ball1 = Urns[Urn],

  % You pick a ball. It's red.
  %
  % Without this important condition, the probabilities are
  %   green: 49/50 (0.98)
  %   red: 1/50 (0.02)
  observe(Ball1 == red),
  
  % Throw away the first ball and pick a new ball from the SAME urn
  Ball2 = cond(Urn==100,
               red, % There are only red balls in the last urn
               cond(Ball1==red, 
                    green, % We drew the only red from that urn so it must be green
                    categorical([1/98,97/98],Colors))),
  if observed_ok then

    add("urn",Urn),
    add("ball2",Ball2)
  end.