/* 

  Dice with re-roll in Picat.

  From https://www.reddit.com/r/Probability/comments/1duf4na/dice_probability_involving_rerolls/
  """
  Dice probability involving re-rolls

  Hi, In a scenario where you need to roll a 6 on a d6 dice for a game, 
  the chance is X/6 of success, where X is the number of d6 dice rolled.

  However you are allowed to re-roll results of 1 or 2. There can only be one 
  re-roll, further results of 1 or 2 do not generate re-rolls.

  What is the probability of success?

  My guess is (X/6) + ((1/3)(X/6)) but I'm uncertain if I've missed something. 
  """

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

/*

  Here we check with 1..6 d6's, i.e. the probability of
  getting at least one 6:

  num_dice = 1
  var : num 6s
  Probabilities:
  0: 0.7755000000000000
  1: 0.2245000000000000
  mean = 0.2245

  var : p
  Probabilities:
  false: 0.7755000000000000
  true: 0.2245000000000000
  mean = [false = 0.7755,true = 0.2245]


  num_dice = 2
  var : num 6s
  Probabilities:
  0: 0.6001000000000000
  1: 0.3466000000000000
  2: 0.0533000000000000
  mean = 0.4532

  var : p
  Probabilities:
  false: 0.6001000000000000
  true: 0.3999000000000000
  mean = [false = 0.6001,true = 0.3999]


  num_dice = 3
  var : num 6s
  Probabilities:
  0: 0.4779000000000000
  1: 0.3937000000000000
  2: 0.1156000000000000
  3: 0.0128000000000000
  mean = 0.6633

  var : p
  Probabilities:
  true: 0.5221000000000000
  false: 0.4779000000000000
  mean = [true = 0.5221,false = 0.4779]


  num_dice = 4
  var : num 6s
  Probabilities:
  1: 0.4159000000000000
  0: 0.3675000000000000
  2: 0.1806000000000000
  3: 0.0336000000000000
  4: 0.0024000000000000
  mean = 0.8875

  var : p
  Probabilities:
  true: 0.6325000000000000
  false: 0.3675000000000000
  mean = [true = 0.6325,false = 0.3675]


  num_dice = 5
  var : num 6s
  Probabilities:
  1: 0.4043000000000000
  0: 0.2861000000000000
  2: 0.2356000000000000
  3: 0.0650000000000000
  4: 0.0087000000000000
  5: 0.0003000000000000
  mean = 1.1068

  var : p
  Probabilities:
  true: 0.7139000000000000
  false: 0.2861000000000000
  mean = [true = 0.7139,false = 0.2861]


  num_dice = 6
  var : num 6s
  Probabilities:
  1: 0.3769000000000000
  2: 0.2768000000000000
  0: 0.2149000000000000
  3: 0.1061000000000000
  4: 0.0228000000000000
  5: 0.0025000000000000
  mean = 1.3525

  var : p
  Probabilities:
  true: 0.7851000000000000
  false: 0.2149000000000000
  mean = [true = 0.7851,false = 0.2149]

*/
go ?=>
  member(NumDice,1..6),
  println(num_dice=NumDice),
  reset_store,
  run_model(10_000,$model(NumDice),[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,
  fail,
  nl.
go => true.
  

dice2(I) = Res =>
  D1 = random_integer1(6),
  if D1 <= 2 then
    D2 = random_integer1(6),
    % Note: in the Gamble model one add D2 to D1
    Res = D2
  else
    Res = D1
  end.
model(NumDice) =>
  AllDice = [dice2(I) : I in 1..NumDice],

  Num6s = [cond(D == 6,1,0) : D in AllDice].sum,
  P = check(Num6s > 0),
  
  add("num 6s",Num6s),
  add("p",P).