/* 

  Run until blue ball in Picat.

  From https://community.wolfram.com/groups/-/m/t/3007279
  """
  Experiment
  1. Start with 2 balls in urn: 1 red and 1 blue
  2. If you pick a blue ball, terminate. If red, put the red ball back and add one red ball to the urn
  3. Repeat until you pick a blue ball
  Find the average length of trials.
  """


  Cf my Gamble model gamble_run_until_blue_ball.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.
% import ordset.

main => go.

/*

  Testing both num to add 1 and 2 (but skips the unlimited version of 2)

  [num_to_add = 1,limit = 1]
  var : len
  Probabilities:
  2: 1.0000000000000000
  mean = 2.0


  [num_to_add = 1,limit = 10]
  var : len
  Probabilities:
  2: 0.4972000000000000
  3: 0.1723000000000000
  11: 0.0973000000000000
  4: 0.0814000000000000
  5: 0.0512000000000000
  6: 0.0340000000000000
  7: 0.0236000000000000
  8: 0.0181000000000000
  9: 0.0147000000000000
  10: 0.0102000000000000
  mean = 3.9115


  [num_to_add = 1,limit = 100]
  var : len
  Probabilities (truncated):
  2: 0.4981000000000000
  3: 0.1725000000000000
  4: 0.0821000000000000
  5: 0.0460000000000000
  .........
  69: 0.0001000000000000
  66: 0.0001000000000000
  64: 0.0001000000000000
  50: 0.0001000000000000
  mean = 6.2161


  [num_to_add = 1,limit = 1000]
  var : len
  Probabilities (truncated):
  2: 0.4992000000000000
  3: 0.1683000000000000
  4: 0.0857000000000000
  5: 0.0483000000000000
  .........
  73: 0.0001000000000000
  58: 0.0001000000000000
  51: 0.0001000000000000
  49: 0.0001000000000000
  mean = 8.5394


  [num_to_add = 1,limit = 10000]
  var : len
  Probabilities (truncated):
  2: 0.5067000000000000
  3: 0.1627000000000000
  4: 0.0837000000000000
  5: 0.0476000000000000
  .........
  66: 0.0001000000000000
  63: 0.0001000000000000
  62: 0.0001000000000000
  52: 0.0001000000000000
  mean = 10.177


  [num_to_add = 1,limit = 100000]
  var : len
  Probabilities (truncated):
  2: 0.5006000000000000
  3: 0.1681000000000000
  4: 0.0818000000000000
  5: 0.0474000000000000
  .........
  76: 0.0001000000000000
  75: 0.0001000000000000
  64: 0.0001000000000000
  60: 0.0001000000000000
  mean = 11.3996


  [num_to_add = 1,limit = none]
  var : len
  Probabilities (truncated):
  2: 0.5041333333333333
  3: 0.1671666666666667
  4: 0.0844000000000000
  5: 0.0482333333333333
  .........
  119: 0.0000333333333333
  104: 0.0000333333333333
  80: 0.0000333333333333
  77: 0.0000333333333333
  mean = 13.2663


  [num_to_add = 2,limit = 1]
  var : len
  Probabilities:
  2: 1.0000000000000000
  mean = 2.0


  [num_to_add = 2,limit = 10]
  var : len
  Probabilities:
  2: 0.5002000000000000
  12: 0.2407000000000000
  4: 0.1259000000000000
  6: 0.0663000000000000
  8: 0.0378000000000000
  10: 0.0291000000000000
  mean = 5.3836


  [num_to_add = 2,limit = 100]
  var : len
  Probabilities (truncated):
  2: 0.4937000000000000
  4: 0.1278000000000000
  102: 0.0845000000000000
  6: 0.0633000000000000
  .........
  98: 0.0007000000000000
  94: 0.0005000000000000
  86: 0.0005000000000000
  74: 0.0005000000000000
  mean = 16.3896


  [num_to_add = 2,limit = 1000]
  var : len
  Probabilities (truncated):
  2: 0.4996000000000000
  4: 0.1290000000000000
  6: 0.0612000000000000
  8: 0.0384000000000000
  .........
  208: 0.0001000000000000
  204: 0.0001000000000000
  202: 0.0001000000000000
  180: 0.0001000000000000
  mean = 48.087


  [num_to_add = 2,limit = 10000]
  var : len
  Probabilities (truncated):
  2: 0.4998000000000000
  4: 0.1304000000000000
  6: 0.0607000000000000
  8: 0.0385000000000000
  .........
  210: 0.0001000000000000
  200: 0.0001000000000000
  188: 0.0001000000000000
  184: 0.0001000000000000
  mean = 151.017


  [num_to_add = 2,limit = 100000]
  var : len
  Probabilities (truncated):
  2: 0.5061000000000000
  4: 0.1240000000000000
  6: 0.0622000000000000
  8: 0.0413000000000000
  .........
  240: 0.0001000000000000
  236: 0.0001000000000000
  216: 0.0001000000000000
  174: 0.0001000000000000
  mean = 476.96


*/
go ?=>
  member(NumToAdd,[1,2]),
  member(Limit,[1,10,100,1000,10_000,100_000,none]),
  [NumToAdd,Limit] != [2,none],
  println([num_to_add=NumToAdd,limit=Limit]),
  reset_store,
  NumRuns = cond(Limit == none, 30_000, 10_000),
  run_model(NumRuns,$model(NumToAdd,Limit),[show_probs_trunc,mean % ,
                           % show_percentiles,
                           % show_hpd_intervals,hpd_intervals=[0.94],
                           % show_histogram,
                           % min_accepted_samples=1000,show_accepted_samples=true
                          ]),
  nl,
  % show_store_lengths,nl,
  fail,
  nl.
go => true.
  

f(NumRed,Limit,NumToAdd) = Res =>
  if NumRed >= Limit then
    Res = NumRed
  else
    Pct = NumRed / (1 + NumRed),
    if flip(Pct) == true then
       Res = f(NumRed+NumToAdd,Limit,NumToAdd)
    else
       Res = NumRed
    end
  end.

% No limit
f(NumRed,NumToAdd) = Res =>
  Pct = NumRed / (1 + NumRed),
  if flip(Pct) == true then
    Res = f(NumRed+NumToAdd,NumToAdd)
  else
    Res = NumRed
  end.

model(NumToAdd,Limit) =>

  % We start with one red (and one blue)
  if Limit == none then
    % No limit
    Len = f(1,NumToAdd) + 1, % add the blue ball  
  else
    Len = f(1,Limit,NumToAdd) + 1, % add the blue ball
  end,
  
  add("len",Len).