/* 

  Multivariate hypergeometric distribution in Picat.

  From Mathematica MultivariateHypergeometricDistribution
  """
  An urn contains 12 red balls, 23 blue balls, and 9 green balls. 

  Find the distribution of a sample of 5 balls drawn without replacement:

  dist = MultivariateHypergeometricDistribution[5, {12, 23, 9}];  
  Probability[x==2 && z==3],{x,y,} in dist)
  -> 
  9/1763
  N@%
  -> 
  0.00510493

  Find the average number of balls of each color in a sample:
  Mean[dist] // N
  -> {1.36364, 2.61364, 1.02273}
  """

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



/*
  random = [[0,3,2],[1,4,0],[0,4,1],[1,3,1],[4,0,1]]
  pdf = 0.00510493
  cdf = 0.00510493
  mean = [1.36364,2.61364,1.02273]

  PDF for all combinations:
  [1,3,1] 0.176120249574589 
  [2,2,1] 0.13838019609432 
  [2,3,0] 0.107629041406693 
  [1,2,2] 0.100640142614051 
  [1,4,0] 0.097844583096994 
  [0,4,1] 0.073383437322745 
  [0,3,2] 0.058706749858196 
  [3,2,0] 0.051251924479378 
  [2,1,2] 0.050320071307025 
  [3,1,1] 0.041933392755854 
  [0,5,0] 0.030984117980715 
  [1,1,3] 0.021347909039344 
  [0,2,3] 0.019568916619399 
  [4,1,0] 0.010483348188964 
  [3,0,2] 0.00729276395754 
  [2,0,3] 0.005104934770278 
  [4,0,1] 0.004102179726116 
  [0,1,4] 0.002668488629918 
  [1,0,4] 0.001392254937349 
  [5,0,0] 0.000729276395754 
  [0,0,5] 0.000116021244779 

*/
go ?=>
  N = 5,
  NumBalls = [12,23,9],
  Check = [2,0,3],
  println(random=multivariate_hypergeometric_dist_n(N,NumBalls,5)),
  println(pdf=multivariate_hypergeometric_dist_pdf(N,NumBalls,Check)),
  println(cdf=multivariate_hypergeometric_dist_cdf(N,NumBalls,Check)),
  println(mean=multivariate_hypergeometric_dist_mean(N,NumBalls)),
  nl,
  println("PDF for all combinations:"),
  Cs = comb_sum(NumBalls.len,N,N),
  [C=multivariate_hypergeometric_dist_pdf(N,NumBalls,C) : C in Cs].sort_down(2).printf_list,
  nl.
go => true.

/*
  [n = 5,num_balls = [12,23,9],check = [2,0,3]]
  var : draw balls
  Probabilities:
  [1,3,1]: 0.1780000000000000
  [2,2,1]: 0.1380000000000000
  [2,3,0]: 0.1066000000000000
  [1,2,2]: 0.1047000000000000
  [1,4,0]: 0.0964000000000000
  [0,4,1]: 0.0763000000000000
  [0,3,2]: 0.0568000000000000
  [2,1,2]: 0.0513000000000000
  [3,2,0]: 0.0496000000000000
  [3,1,1]: 0.0370000000000000
  [0,5,0]: 0.0311000000000000
  [1,1,3]: 0.0228000000000000
  [0,2,3]: 0.0180000000000000
  [4,1,0]: 0.0108000000000000
  [3,0,2]: 0.0067000000000000
  [4,0,1]: 0.0061000000000000
  [2,0,3]: 0.0052000000000000
  [0,1,4]: 0.0028000000000000
  [1,0,4]: 0.0011000000000000
  [5,0,0]: 0.0007000000000000

  var : red
  Probabilities:
  1: 0.4030000000000000
  2: 0.3011000000000000
  0: 0.1850000000000000
  3: 0.0933000000000000
  4: 0.0169000000000000
  5: 0.0007000000000000
  mean = 1.3562

  var : blue
  Probabilities:
  3: 0.3414000000000000
  2: 0.3103000000000000
  4: 0.1727000000000000
  1: 0.1247000000000000
  5: 0.0311000000000000
  0: 0.0198000000000000
  mean = 2.6158

  var : green
  Probabilities:
  1: 0.4354000000000000
  0: 0.2952000000000000
  2: 0.2195000000000000
  3: 0.0460000000000000
  4: 0.0039000000000000
  mean = 1.028

  var : p
  Probabilities:
  false: 0.9948000000000000
  true: 0.0052000000000000
  mean = 0.0052

  theoretical_pdf = 0.00510493
  theoretical_mean = [1.36364,2.61364,1.02273]

*/
go2 ?=>
  N = 5,
  NumBalls = [12,23,9],
  Check = [2,0,3],
  println([n=N,num_balls=NumBalls,check=Check]),
  reset_store,
  run_model(10_000,$model2(N,NumBalls,Check),[show_probs,mean % ,
                           % show_percentiles,
                           % show_hpd_intervals,hpd_intervals=[0.94],
                           % show_histogram,
                           % min_accepted_samples=1000,show_accepted_samples=true
                          ]),
  println(theoretical_pdf=multivariate_hypergeometric_dist_pdf(N,NumBalls,Check)),
  println(theoretical_mean=multivariate_hypergeometric_dist_mean(N,NumBalls)),
  nl,
  % show_store_lengths,nl,
  % fail,
  nl.
go2 => true.
  
  
model2(N,NumBalls,Check) =>
  DrawnBalls = multivariate_hypergeometric_dist(N,NumBalls),
  [Red,Blue,Green] = DrawnBalls,

  P = check([Red,Blue,Green] == Check),
  
  add("draw balls",DrawnBalls),
  add("red",Red),
  add("blue",Blue),
  add("green",Green),
  add("p",P).