/* 

  Consecutive heads Picat.

  https://brainstellar.com/puzzles/probability/116
  """
  What is the expected number of coin tosses required to get N consecutive heads?
  Denote the term by E[N]

  Answer: Recurrence form: E[N] = 2 E[N](N-1) + 2. Closed form: 2^(N+1) + 2.
  """

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

/*

  [num_heads = 1,theoretical = 2]
  var : len
  Probabilities (truncated):
  1: 0.5024999999999999
  2: 0.2503000000000000
  3: 0.1229000000000000
  4: 0.0622000000000000
  .........
  10: 0.0011000000000000
  11: 0.0005000000000000
  14: 0.0001000000000000
  12: 0.0001000000000000
  mean = 1.9925


  [num_heads = 2,theoretical = 6]
  var : len
  Probabilities (truncated):
  2: 0.2697000000000000
  3: 0.1235000000000000
  4: 0.1217000000000000
  5: 0.0938000000000000
  .........
  34: 0.0002000000000000
  31: 0.0002000000000000
  40: 0.0001000000000000
  37: 0.0001000000000000
  mean = 5.8631


  [num_heads = 3,theoretical = 14]
  var : len
  Probabilities (truncated):
  3: 0.1244000000000000
  5: 0.0673000000000000
  6: 0.0656000000000000
  4: 0.0577000000000000
  .........
  82: 0.0001000000000000
  77: 0.0001000000000000
  71: 0.0001000000000000
  70: 0.0001000000000000
  mean = 14.1641


  [num_heads = 4,theoretical = 30]
  var : len
  Probabilities (truncated):
  4: 0.0667000000000000
  8: 0.0315000000000000
  9: 0.0305000000000000
  6: 0.0304000000000000
  .........
  145: 0.0001000000000000
  143: 0.0001000000000000
  138: 0.0001000000000000
  135: 0.0001000000000000
  mean = 29.9536


  [num_heads = 5,theoretical = 62]
  var : len
  Probabilities (truncated):
  5: 0.0333000000000000
  9: 0.0182000000000000
  10: 0.0163000000000000
  11: 0.0155000000000000
  .........
  253: 0.0001000000000000
  241: 0.0001000000000000
  239: 0.0001000000000000
  216: 0.0001000000000000
  mean = 61.808


*/
go ?=>
  member(NumHeads,1..5),
  println([num_heads=NumHeads,theoretical=(2**(NumHeads+1)-2)]),
  reset_store,
  run_model(10_000,$model(NumHeads),[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,nl,
  fail,
  nl.
go => true.

take_last(A,N) = drop(A,A.len-N).
  
f(A,NumHeads) = Res =>
  if (A.len >= NumHeads, take_last(A,NumHeads).sum == NumHeads) then
    Res = A
  else
    Res = f(A++[bern(1/2)],NumHeads)
  end.
  
model(NumHeads) =>
  % Heads = ones(NumHeads,1),

  A = f([],NumHeads),
  Len = A.len,
  % println(A==Len),
  add("len",Len).