/* 

  Bugs book example 3.4.1 in Picat.

  Example 3.4.1 (b) Three coins, with a prediction.
  """
  Suppose I have 3 coins in my pocket. The coins may be either fair, biased 3:1 in
  favour of heads, or 3:1 in favour of tails, but I don't know how many of
  each type there are among the 3 coins. I randomly select 1 coin att toss it once,
  observing a head. What is the posterior distribution of the probability of a head?
  """

  Cf
  - ppl_gamble_bugs_book_3_4_1.pi 
  - my Gamble model gamble_bugs_book_3_4_1b.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.

/*

  var : y diff
  Probabilities:
  0.25: 0.3710616094722055
  -0.5: 0.1735902067027895
  0.5: 0.1659642785470600
  -0.75: 0.1254264499297612
  -0.25: 0.1236203090507726
  0.75: 0.0403371462974112
  mean = -0.00576962

  var : y pred
  Probabilities:
  1: 0.5773630343166767
  0: 0.4226369656833233
  mean = 0.577363

  var : coin
  Probabilities:
  3: 0.4964880594019667
  2: 0.3395544852498495
  1: 0.1639574553481838
  mean = 2.33253

  var : theta true
  Probabilities:
  0.75: 0.4964880594019667
  0.5: 0.3395544852498495
  0.25: 0.1639574553481838
  mean = 0.583133

  var : coin_prob 1
  Probabilities:
  0: 0.8360425446518162
  1: 0.1639574553481838
  mean = 0.163957

  var : coin_prob 2
  Probabilities:
  0: 0.6604455147501506
  1: 0.3395544852498495
  mean = 0.339554

  var : coin_prob 3
  Probabilities:
  0: 0.5035119405980333
  1: 0.4964880594019667
  mean = 0.496488

*/
go ?=>
  reset_store,
  run_model(10_000,$model,[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,nl,
  % fail,
  nl.
go => true.
  
model() =>
  N = 3,
  Coin = random_integer1(3),
  Theta = [ I * 0.25 : I in 1..N],
  CoinProb = [cond(Coin==I,1,0) : I in 1..N],
  ThetaTrue = Theta[Coin],
  Y = bern(ThetaTrue),
  YPred = bern(ThetaTrue),
  YDiff = YPred - ThetaTrue,

  observe(Y == 1), % We observe a head
  if observed_ok then
    add("y diff",YDiff),
    add("y pred",YPred),
    add("coin",Coin),
    add("theta true",ThetaTrue),
    add("coin_prob 1",CoinProb[1]),
    add("coin_prob 2",CoinProb[2]),
    add("coin_prob 3",CoinProb[3]),    
  end.