/* 

  Ehrenfest urn scheme in Picat.

  From Polya Urn Models, page 21
  """
  An Ehrenfest urn starts out with one white and one blue ball. Balls are sampled from 
  the urn. Whenever a ball of a certain color is picked, we discard that ball and replace 
  in the urn a ball of the opposite color. The urn always returns to the initial state 
  after an even number of draws, and is always out of that state after an odd number 
  of draws. When it is out of the initial state, it can be in one of two equally 
  probable states: A state with two white balls or a state with two blue balls.
  """

  Here we randomize the number of even and odd draws: 
  (0..2..8) and (1..2..9) draws, respectively.

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

/*

  [num_white_init = 1,num_blue_init = 1]
  var : num white even
  Probabilities:
  1: 1.0000000000000000
  mean = 1.0

  var : num blue even
  Probabilities:
  1: 1.0000000000000000
  mean = 1.0

  var : num white odd
  Probabilities:
  2: 0.5007000000000000
  0: 0.4993000000000000
  mean = 1.0014

  var : num blue odd
  Probabilities:
  0: 0.5007000000000000
  2: 0.4993000000000000
  mean = 0.9986

  var : white even, blue even
  Probabilities:
  [1,1]: 1.0000000000000000
  mean = [[1,1] = 1.0]

  var : white odd, blue odd
  Probabilities:
  [2,0]: 0.5007000000000000
  [0,2]: 0.4993000000000000
  mean = [[2,0] = 0.5007,[0,2] = 0.4993]

  var : even draws
  Probabilities:
  8: 0.2033000000000000
  4: 0.2017000000000000
  6: 0.2005000000000000
  2: 0.1989000000000000
  0: 0.1956000000000000
  mean = 4.034

  var : odd draws
  Probabilities:
  9: 0.2073000000000000
  5: 0.2067000000000000
  1: 0.2002000000000000
  3: 0.1972000000000000
  7: 0.1886000000000000
  mean = 5.0112


  [num_white_init = 1,num_blue_init = 2]
  var : num white even
  Probabilities:
  1: 0.8101000000000000
  3: 0.1899000000000000
  mean = 1.3798

  var : num blue even
  Probabilities:
  2: 0.8101000000000000
  0: 0.1899000000000000
  mean = 1.6202

  var : num white odd
  Probabilities:
  2: 0.7291000000000000
  0: 0.2709000000000000
  mean = 1.4582

  var : num blue odd
  Probabilities:
  1: 0.7291000000000000
  3: 0.2709000000000000
  mean = 1.5418

  var : white even, blue even
  Probabilities:
  [1,2]: 0.8101000000000000
  [3,0]: 0.1899000000000000
  mean = [[1,2] = 0.8101,[3,0] = 0.1899]

  var : white odd, blue odd
  Probabilities:
  [2,1]: 0.7291000000000000
  [0,3]: 0.2709000000000000
  mean = [[2,1] = 0.7291,[0,3] = 0.2709]

  var : even draws
  Probabilities:
  4: 0.2058000000000000
  2: 0.2051000000000000
  8: 0.1966000000000000
  0: 0.1966000000000000
  6: 0.1959000000000000
  mean = 3.9816

  var : odd draws
  Probabilities:
  1: 0.2075000000000000
  3: 0.2048000000000000
  5: 0.2037000000000000
  9: 0.1945000000000000
  7: 0.1895000000000000
  mean = 4.9174


  [num_white_init = 1,num_blue_init = 10]
  var : num white even
  Probabilities:
  3: 0.3984000000000000
  1: 0.2853000000000000
  5: 0.2646000000000000
  7: 0.0493000000000000
  9: 0.0024000000000000
  mean = 3.1702

  var : num blue even
  Probabilities:
  8: 0.3984000000000000
  10: 0.2853000000000000
  6: 0.2646000000000000
  4: 0.0493000000000000
  2: 0.0024000000000000
  mean = 7.8298

  var : num white odd
  Probabilities:
  4: 0.4113000000000000
  2: 0.3765000000000000
  6: 0.1655000000000000
  0: 0.0262000000000000
  8: 0.0203000000000000
  10: 0.0002000000000000
  mean = 3.5556

  var : num blue odd
  Probabilities:
  7: 0.4113000000000000
  9: 0.3765000000000000
  5: 0.1655000000000000
  11: 0.0262000000000000
  3: 0.0203000000000000
  1: 0.0002000000000000
  mean = 7.4444

  var : white even, blue even
  Probabilities:
  [3,8]: 0.3984000000000000
  [1,10]: 0.2853000000000000
  [5,6]: 0.2646000000000000
  [7,4]: 0.0493000000000000
  [9,2]: 0.0024000000000000
  mean = [[3,8] = 0.3984,[1,10] = 0.2853,[5,6] = 0.2646,[7,4] = 0.0493,[9,2] = 0.0024]

  var : white odd, blue odd
  Probabilities:
  [4,7]: 0.4113000000000000
  [2,9]: 0.3765000000000000
  [6,5]: 0.1655000000000000
  [0,11]: 0.0262000000000000
  [8,3]: 0.0203000000000000
  [10,1]: 0.0002000000000000
  mean = [[4,7] = 0.4113,[2,9] = 0.3765,[6,5] = 0.1655,[0,11] = 0.0262,[8,3] = 0.0203,[10,1] = 0.0002]

  var : even draws
  Probabilities:
  6: 0.2063000000000000
  2: 0.2029000000000000
  4: 0.2001000000000000
  0: 0.1963000000000000
  8: 0.1944000000000000
  mean = 3.9992

  var : odd draws
  Probabilities:
  1: 0.2039000000000000
  5: 0.2027000000000000
  7: 0.2014000000000000
  3: 0.1975000000000000
  9: 0.1945000000000000
  mean = 4.9702

*/
go ?=>
  member([NumWhiteInit,NumBlueInit],[
                                     [1,1],
                                     [1,2],
                                     [1,10]
                                     ]),  
  println([num_white_init=NumWhiteInit,num_blue_init=NumBlueInit]),
  reset_store,
  run_model(10_000,$model(NumWhiteInit,NumBlueInit),[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,nl,
  fail,
  nl.
go => true.
  

f(N, NumWhite, NumBlue) = Res =>
  % Adjust for negative number of balls
  NumWhite2 = cond(NumWhite <= 0,0,NumWhite),
  NumBlue2 = cond(NumBlue <= 0,0,NumBlue),
  if N == 0 ; NumWhite2 + NumBlue2 == 0 then
    Res = [NumWhite2,NumBlue2]     
  else
    if flip( NumWhite2 / (NumWhite2 + NumBlue2) ) == true then
      % White ball
      Res = f(N-1,NumWhite2-1,NumBlue2+1)
    else
      % Blue ball
      Res = f(N-1,NumWhite2+1,NumBlue2-1)    
    end
  end.
  
model(NumWhiteInit,NumBlueInit) =>
  % Generate even and odd number of draws
  EvenDraws = 2 * random_integer(5),
  OddDraws = 2 * random_integer(5) + 1,

  [NumWhiteEven,NumBlueEven] = f(EvenDraws,NumWhiteInit,NumBlueInit),

  [NumWhiteOdd,NumBlueOdd] = f(OddDraws,NumWhiteInit,NumBlueInit),
  
  add_all([ ["num white even",NumWhiteEven],
            ["num blue even",NumBlueEven],
            ["num white odd",NumWhiteOdd],
            ["num blue odd",NumBlueOdd],
            ["white even, blue even",[NumWhiteEven,NumBlueEven]],
            ["white odd, blue odd",[NumWhiteOdd,NumBlueOdd]],            
            ["even draws",EvenDraws],
            ["odd draws",OddDraws]]).