/* 

  Urn puzzle in Picat.

  From Daniel Litt  (Jan 28, 2024)
  https://x.com/littmath/status/1751648838501224790
  """
  You are given an urn containing 100 balls; n of them are red, and 100-n are green, 
  where n is chosen uniformly at random in [0, 100]. You take a random ball out of 
  the urn—it’s red—and discard it. The next ball you pick (out of the 99 remaining) is:
    More likely to be red: 22.6%
    More likely to be green: 37.1%
    Equally likely: 20.9%
    Don’t know/see results: 19.5%
  """

  Note: The "official" answer 
  (in https://www.quantamagazine.org/perplexing-the-web-one-probability-puzzle-at-a-time-20240829/?mc_cid=94caee8978) is that probability of the next ball red is 2/3, and blue 1/3.

  However, I got a slightly different answer by restricting to be at least one red ball.

  Cf my Gamble model gamble_urn_puzzle2.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 : num red
  Probabilities (truncated):
  101: 0.0207833733013589
  99: 0.0201838529176659
  95: 0.0191846522781775
  93: 0.0191846522781775
  .........
  5: 0.0005995203836930
  4: 0.0005995203836930
  3: 0.0005995203836930
  2: 0.0005995203836930
  mean = 67.9982

  var : num green
  Probabilities (truncated):
  -1: 0.0207833733013589
  1: 0.0201838529176659
  7: 0.0191846522781775
  5: 0.0191846522781775
  .........
  98: 0.0005995203836930
  97: 0.0005995203836930
  96: 0.0005995203836930
  95: 0.0005995203836930
  mean = 32.0018

  var : ball 1
  Probabilities (truncated):
  3: 0.0223820943245404
  7: 0.0213828936850520
  1: 0.0199840127897682
  6: 0.0197841726618705
  .........
  97: 0.0007993605115907
  101: 0.0005995203836930
  98: 0.0005995203836930
  100: 0.0003996802557954
  mean = 34.1105

  var : color ball 1
  Probabilities:
  red: 1.0000000000000000
  mean = [red = 1.0]

  var : ball 2
  Probabilities (truncated):
  42: 0.0143884892086331
  64: 0.0129896083133493
  4: 0.0127897681854516
  57: 0.0123900879296563
  .........
  93: 0.0073940847322142
  6: 0.0071942446043165
  40: 0.0067945643485212
  61: 0.0053956834532374
  mean = 51.3545

  var : color ball 2
  Probabilities:
  red: 0.6652677857713829
  green: 0.3347322142286171
  mean = [red = 0.665268,green = 0.334732]



  I.e. the probabiity is about 2/3 that the next ball is red, and about 1/3 that it's green.

  Note from my Gamble model (which calculates exact probabilities: 
  * without the constraint (num-red > 0), then the probabilities are exactly 2/3 and 1/3:
      red: 2/3 (0.6666666666666666)
      green: 1/3 (0.3333333333333333)
  
  See https://colab.research.google.com/drive/1vK8sl3ZecjBTjrcNcee6hmlHvrp3Zh4_?usp=sharing
  for another (Bayesian simulation) approach. It gives 2/3 and 1/3. 

  This problem was found via the Quanta Magazine article on Daniel Litt:
  "Perplexing the Web, One Probability Puzzle at a Time" by Erica Klarreich:
  https://www.quantamagazine.org/perplexing-the-web-one-probability-puzzle-at-a-time-20240829/?mc_cid=94caee8978

*/
go ?=>
  reset_store,
  run_model(10_000,$model,[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,
  % fail,
  nl.
go => true.
  
  
model() =>
  Total = 100,
  Total1 = Total+1,
  NumRed = random_integer1(Total1), % number of balls: 0..100 -> 1..101


  % There must be at least one red ball, i.e. the one we first select.
  % Without this, the probabilities are exactly 2/3 red and 1/3 green,
  % which is the official answer
  observe(NumRed > 0),

  NumGreen = Total - NumRed,

  % What color is the ball
  Color = [cond(Ball <= NumRed,red,green) : Ball in 1..Total1],

  Ball1 = random_integer1(Total1),
  ColorBall1 = Color[Ball1],

  % The first ball is red
  observe(Color[Ball1] == red),

  Ball2 = random_integer1(Total1),
  % Discard first ball since it cannot be selected.
  observe(Ball1 != Ball2),

  ColorBall2 = Color[Ball2],
  
  if observed_ok then
    add("num red",NumRed),
    add("num green",NumGreen),
    add("ball 1",Ball1),
    add("color ball 1",ColorBall1),
    add("ball 2",Ball2),
    add("color ball 2",ColorBall2),
  end.