/* 

  Coin flips in Picat.

  https://pedrozudo.github.io/docs/20_thesis/thesis_final_online_compressed.pdf
  page 112
  """
  Example 6.14. We model a coin flip scenario where the prior probability of the coin
  turning up heads is distributed according to a mixture of two beta distributions. DC-
  ProbLog then allows us to learn the posterior distribution by taking into account the
  data in Lines 6 to 8.
  1 0.2::a.
  2 b~beta(1,1):- a.
  3 b~beta(1,2):- \+a.
  4 B::coin_(flip N):- B is b.
  5
  6 evidence(coin_(flip 1), true).
  7 evidence(coin_(flip 2), false).
  8 evidence(coin_(flip 3), true)  
  """


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

/*
  
  obs = [true,false,true]
  var : a
  Probabilities:
  false: 0.7741007194244605
  true: 0.2258992805755396
  mean = [false = 0.774101,true = 0.225899]

  var : b
  Probabilities (truncated):
  0.97701469116553: 0.0014388489208633
  0.956196715129553: 0.0014388489208633
  0.953198354983884: 0.0014388489208633
  0.952511605797993: 0.0014388489208633
  .........
  0.11375084432707: 0.0014388489208633
  0.098148246555909: 0.0014388489208633
  0.093722070359515: 0.0014388489208633
  0.092905917397437: 0.0014388489208633
  mean = 0.52544


  obs = [true,false,true,true]
  var : a
  Probabilities:
  false: 0.6956521739130435
  true: 0.3043478260869565
  mean = [false = 0.695652,true = 0.304348]

  var : b
  Probabilities (truncated):
  0.982999087136628: 0.0025575447570332
  0.97652945566699: 0.0025575447570332
  0.952412719369843: 0.0025575447570332
  0.949316905950588: 0.0025575447570332
  .........
  0.201836636726502: 0.0025575447570332
  0.19462809374637: 0.0025575447570332
  0.159304290234198: 0.0025575447570332
  0.136902263325793: 0.0025575447570332
  mean = 0.599196


  obs = [true,false,true,true,true]
  var : a
  Probabilities:
  false: 0.7050000000000000
  true: 0.2950000000000000
  mean = [false = 0.705,true = 0.295]

  var : b
  Probabilities (truncated):
  0.968330445936587: 0.0050000000000000
  0.94534182935141: 0.0050000000000000
  0.933348819068863: 0.0050000000000000
  0.928706981949481: 0.0050000000000000
  .........
  0.321153069777179: 0.0050000000000000
  0.320226031799502: 0.0050000000000000
  0.204955117173529: 0.0050000000000000
  0.203609453827779: 0.0050000000000000
  mean = 0.648824

*/
go ?=>
  member(Obs,[ [true,false,true],[true,false,true,true],[true,false,true,true,true]]),
  println(obs=Obs),
  reset_store,
  run_model(10_000,$model(Obs),[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,nl,
  fail,
  nl.
go => true.

 
model(Obs) =>
  N = Obs.len,
  A = flip(0.2),
  B = cond(A == true, beta_dist(1,1), beta_dist(1,2)),

  CoinFlips = [ flip(B) : _ in 1..N],

  foreach({I,C} in zip(1..Obs.len,Obs))
    observe(CoinFlips[I] == C),
  end,
  
  if observed_ok then
    add("a",A),
    add("b",B)
  end.