/* 

  Rumor in Picat.

  From https://math.stackexchange.com/questions/12689/probability-on-spreading-of-rumors
  """
  A little help here. Exercise 21, Ch. 2 from Feller's book reads

  In a town [of] n+1 inhabitants, a person tells a rumor to a second person, who in 
  turn repeats it to a third person, etc. At each step, the recipient of the rumor 
  is chosen at random from the n people available. Find the probability that the rumor 
  will be told r times without:
  a) returning to the originator, 
  b) being repeated to any person. 
  Do the same problem when at each step the rumor is told by one person 
  to a gathering of N randomly chosen people. (The first question is the special case N=1).
  """


  Cf my Gamble model gamble_rumor.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 with some different number of people (n), and different number
  of people to spread the rumor to at the same time (rep). The number
  of steps (num_steps) is constant (10).


  [n = 10,rep = 1,num_steps = 10]
  var : repeat start
  Probabilities:
  10: 0.3940000000000000
  2: 0.1170000000000000
  3: 0.1006666666666667
  4: 0.0966666666666667
  5: 0.0783333333333333
  6: 0.0620000000000000
  8: 0.0560000000000000
  7: 0.0510000000000000
  9: 0.0443333333333333
  mean = 6.83033

  var : repeat any
  Probabilities:
  4: 0.2343333333333333
  3: 0.2056666666666667
  5: 0.2013333333333333
  6: 0.1300000000000000
  2: 0.1170000000000000
  7: 0.0716666666666667
  8: 0.0326666666666667
  9: 0.0063333333333333
  10: 0.0010000000000000
  mean = 4.405

  var : num distinct
  Probabilities:
  7: 0.4120000000000000
  6: 0.2680000000000000
  8: 0.2136666666666667
  5: 0.0640000000000000
  9: 0.0363333333333333
  4: 0.0050000000000000
  10: 0.0010000000000000
  mean = 6.87833

  var : p start
  Probabilities:
  true: 0.6060000000000000
  false: 0.3940000000000000
  mean = 0.606

  var : p any
  Probabilities:
  true: 0.9990000000000000
  false: 0.0010000000000000
  mean = 0.999

  [p_start = 0.606,p_any = 0.999]

  [n = 2,rep = 1,num_steps = 10]
  [p_start = 1.0,p_any = 1.0]

  [n = 20,rep = 1,num_steps = 10]
  [p_start = 0.363333,p_any = 0.901667]

  [n = 50,rep = 1,num_steps = 10]
  [p_start = 0.158333,p_any = 0.544]

  [n = 100,rep = 1,num_steps = 10]
  [p_start = 0.0736667,p_any = 0.307333]

  [n = 10,rep = 2,num_steps = 10]
  var : repeat start
  Probabilities:
  10: 0.4383333333333334
  4: 0.1096666666666667
  3: 0.1063333333333333
  6: 0.0856666666666667
  5: 0.0776666666666667
  7: 0.0653333333333333
  8: 0.0630000000000000
  9: 0.0540000000000000
  mean = 7.49067

  var : repeat any
  Probabilities:
  5: 0.2636666666666667
  3: 0.2196666666666667
  4: 0.1880000000000000
  6: 0.1593333333333333
  7: 0.1086666666666667
  8: 0.0473333333333333
  9: 0.0116666666666667
  10: 0.0016666666666667
  mean = 4.94633

  var : num distinct
  Probabilities:
  7: 0.3656666666666666
  8: 0.3643333333333333
  6: 0.1303333333333333
  9: 0.1133333333333333
  5: 0.0186666666666667
  10: 0.0076666666666667
  mean = 7.44633

  var : p start
  Probabilities:
  true: 0.5616666666666666
  false: 0.4383333333333334
  mean = 0.561667

  var : p any
  Probabilities:
  true: 0.9983333333333333
  false: 0.0016666666666667
  mean = 0.998333

  [p_start = 0.561667,p_any = 0.998333]

  [n = 2,rep = 2,num_steps = 10]
  [p_start = 1.0,p_any = 1.0]

  [n = 20,rep = 2,num_steps = 10]
  [p_start = 0.322,p_any = 0.866667]

  [n = 50,rep = 2,num_steps = 10]
  [p_start = 0.132,p_any = 0.488667]

  [n = 100,rep = 2,num_steps = 10]
  [p_start = 0.075,p_any = 0.296]

  [n = 10,rep = 3,num_steps = 10]
  var : repeat start
  Probabilities:
  10: 0.4550000000000000
  5: 0.1086666666666667
  6: 0.1046666666666667
  4: 0.1030000000000000
  9: 0.0840000000000000
  7: 0.0726666666666667
  8: 0.0720000000000000
  mean = 7.974

  var : repeat any
  Probabilities:
  4: 0.3310000000000000
  5: 0.2466666666666667
  6: 0.1723333333333333
  7: 0.1616666666666667
  8: 0.0683333333333333
  9: 0.0180000000000000
  10: 0.0020000000000000
  mean = 5.45167

  var : num distinct
  Probabilities:
  7: 0.4236666666666667
  8: 0.3370000000000000
  6: 0.1486666666666667
  9: 0.0770000000000000
  5: 0.0113333333333333
  10: 0.0020000000000000
  4: 0.0003333333333333
  mean = 7.32467

  var : p start
  Probabilities:
  true: 0.5450000000000000
  false: 0.4550000000000000
  mean = 0.545

  var : p any
  Probabilities:
  true: 0.9980000000000000
  false: 0.0020000000000000
  mean = 0.998

  [p_start = 0.545,p_any = 0.998]

  [n = 2,rep = 3,num_steps = 10]
  [p_start = 1.0,p_any = 1.0]

  [n = 20,rep = 3,num_steps = 10]
  [p_start = 0.291,p_any = 0.822333]

  [n = 50,rep = 3,num_steps = 10]
  [p_start = 0.112,p_any = 0.451667]

  [n = 100,rep = 3,num_steps = 10]
  [p_start = 0.062,p_any = 0.240333]

*/
go ?=>
  NumSteps = 10, % That's the 'r' in the problem statement
  member(Rep,1..3), % This is the number of people to spread the rumor to ('N' above)
  member(N, [10,2,20,50,100]),
  println([n=N,rep=Rep,num_steps=NumSteps]),
  reset_store,
  Opts = cond( (N == 10,NumSteps == 10), [show_probs,mean],[presentation=[]]),
  run_model(3000,$model(N,Rep,NumSteps), Opts),
  Store = get_store(),
  PStart = Store.get("p start").mean,
  PAny = Store.get("p any").mean,
  println([p_start=PStart,p_any=PAny]),
  nl,
  % show_store_lengths,nl,
  fail,
  nl.
go => true.
  
f(A,Rep,Ps,NumSteps) = Res =>
  if A.len >= NumSteps then
    Res = A
  else
    % One does not tell a rumor to themselves
    Ps1 = Ps.delete(A.last),
    Told = draw_without_replacement(Rep,Ps1),
    Res = f(A ++ Told,Rep,Ps,NumSteps)
  end.

model(N,Rep,NumSteps) =>
  Ps = 1..N,
  Start = 1,

  A = f([Start],Rep,Ps,NumSteps),
  Len = A.len,

  NumDistinct = A.remove_dups.len,
  RepeatedMap = new_map([P=[I : I in 1..Len, A[I] == P] : P in Ps]),

  % Note: We want the number of steps, not the step itself: so step - 1
  RepeatStart = cond(RepeatedMap.get(Start).len >= 2, RepeatedMap.get(Start).second-1,NumSteps),
  RepeatedAnyL = [ R.second : R in RepeatedMap.values, R.len > 1 ],
  RepeatedAny = cond(RepeatedAnyL.len > 0, RepeatedAnyL.min-1, NumSteps),

  PStart = check(RepeatStart < NumSteps),
  PAny = check(RepeatedAny < NumSteps),

  if N == 10, NumSteps == 10 then
    add("repeat start",RepeatStart),
    add("repeat any",RepeatedAny),
    add("num distinct",NumDistinct)
  end,
  add("p start",PStart),
  add("p any",PAny).