/* 

  The car and the goats in Picat.

  From  Gunnar Blom, Lars Holst, Dennis Sandell:
  "Problems and Snapshots from the World of Probability"
  Page 3f, Problem 1.3 The car and the goats

  This first part of the problem is a "traditional" analysis of 
  the Monty Hall problem. See ppl_monty_hall.pi for a model 
  of this problem.

  The second part "A generalization" is a urn approach which
  this PPL model implements.

  Which of these strategies is best for A:
  * strategy 1
    """
    A draws a ball at random. If it is white, he wins the game,
    otherwise he loses.
    """
    This corresponds to the Stay strategy.

  * strategy 2
    """
    A draws a ball and throws it away. B removes a black ball.
    A again draws a ball. If it is white, he wins the game,
    otherwise he loses.
    """
    This corresponds to the Switch strategy.

  Here we test some different number of white (a) and black (b) balls.
  The first (a=1, b=2) is the traditional Monty Hall problem.

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

main => go.

% The theoretical probabilities for the two strategies
theoretical_prob(A,B) = [Strat1,Strat2] =>
  Strat1 = A / (A+B),
  Strat2 = (A/(A+B))  * (1 + 1/(A+B-2)).
  
/*
[a = 1,b = 2,theoretical = [0.333333,0.666667]]
var : 1 ball
Probabilities:
b: 0.6690000000000000
w: 0.3310000000000000
mean = [b = 0.669,w = 0.331]

var : 1 win
Probabilities:
0: 0.6690000000000000
1: 0.3310000000000000
mean = 0.331

var : 2 ball
Probabilities:
w: 0.6701000000000000
b: 0.3299000000000000
mean = [w = 0.6701,b = 0.3299]

var : 2 win
Probabilities:
1: 0.6701000000000000
0: 0.3299000000000000
mean = 0.6701

var : s2 beats s1
Probabilities:
0: 0.5528000000000000
1: 0.4472000000000000
mean = 0.4472


[a = 1,b = 3,theoretical = [0.25,0.375]]
var : 1 ball
Probabilities:
b: 0.7508000000000000
w: 0.2492000000000000
mean = [b = 0.7508,w = 0.2492]

var : 1 win
Probabilities:
0: 0.7508000000000000
1: 0.2492000000000000
mean = 0.2492

var : 2 ball
Probabilities:
b: 0.6198000000000000
w: 0.3802000000000000
mean = [b = 0.6198,w = 0.3802]

var : 2 win
Probabilities:
0: 0.6198000000000000
1: 0.3802000000000000
mean = 0.3802

var : s2 beats s1
Probabilities:
0: 0.7154000000000000
1: 0.2846000000000000
mean = 0.2846


[a = 1,b = 10,theoretical = [0.0909091,0.10101]]
var : 1 ball
Probabilities:
b: 0.9103000000000000
w: 0.0897000000000000
mean = [b = 0.9103,w = 0.0897]

var : 1 win
Probabilities:
0: 0.9103000000000000
1: 0.0897000000000000
mean = 0.0897

var : 2 ball
Probabilities:
b: 0.8960000000000000
w: 0.1040000000000000
mean = [b = 0.896,w = 0.104]

var : 2 win
Probabilities:
0: 0.8960000000000000
1: 0.1040000000000000
mean = 0.104

var : s2 beats s1
Probabilities:
0: 0.9052000000000000
1: 0.0948000000000000
mean = 0.0948


[a = 2,b = 10,theoretical = [0.166667,0.183333]]
var : 1 ball
Probabilities:
b: 0.8309000000000000
w: 0.1691000000000000
mean = [b = 0.8309,w = 0.1691]

var : 1 win
Probabilities:
0: 0.8309000000000000
1: 0.1691000000000000
mean = 0.1691

var : 2 ball
Probabilities:
b: 0.8144000000000000
w: 0.1856000000000000
mean = [b = 0.8144,w = 0.1856]

var : 2 win
Probabilities:
0: 0.8144000000000000
1: 0.1856000000000000
mean = 0.1856

var : s2 beats s1
Probabilities:
0: 0.8449000000000000
1: 0.1551000000000000
mean = 0.1551


[a = 2,b = 2,theoretical = [0.5,0.75]]
var : 1 ball
Probabilities:
w: 0.5043000000000000
b: 0.4957000000000000
mean = [w = 0.5043,b = 0.4957]

var : 1 win
Probabilities:
1: 0.5043000000000000
0: 0.4957000000000000
mean = 0.5043

var : 2 ball
Probabilities:
w: 0.7506000000000000
b: 0.2494000000000000
mean = [w = 0.7506,b = 0.2494]

var : 2 win
Probabilities:
1: 0.7506000000000000
0: 0.2494000000000000
mean = 0.7506

var : s2 beats s1
Probabilities:
0: 0.6304000000000000
1: 0.3696000000000000
mean = 0.3696


[a = 10,b = 10,theoretical = [0.5,0.527778]]
var : 1 ball
Probabilities:
b: 0.5032000000000000
w: 0.4968000000000000
mean = [b = 0.5032,w = 0.4968]

var : 1 win
Probabilities:
0: 0.5032000000000000
1: 0.4968000000000000
mean = 0.4968

var : 2 ball
Probabilities:
w: 0.5305000000000000
b: 0.4695000000000000
mean = [w = 0.5305,b = 0.4695]

var : 2 win
Probabilities:
1: 0.5305000000000000
0: 0.4695000000000000
mean = 0.5305

var : s2 beats s1
Probabilities:
0: 0.7366000000000000
1: 0.2634000000000000
mean = 0.2634


[a = 10,b = 2,theoretical = [0.833333,0.916667]]
var : 1 ball
Probabilities:
w: 0.8326000000000000
b: 0.1674000000000000
mean = [w = 0.8326,b = 0.1674]

var : 1 win
Probabilities:
1: 0.8326000000000000
0: 0.1674000000000000
mean = 0.8326

var : 2 ball
Probabilities:
w: 0.9179000000000000
b: 0.0821000000000000
mean = [w = 0.9179,b = 0.0821]

var : 2 win
Probabilities:
1: 0.9179000000000000
0: 0.0821000000000000
mean = 0.9179

var : s2 beats s1
Probabilities:
0: 0.8455000000000000
1: 0.1545000000000000
mean = 0.1545

*/
go ?=>
  member([A,B], [[1,2],
                 [1,3],
                 [1,10],
                 [2,10],
                 [2,2],
                 [10,10],
                 [10,2]
             ]),
  println([a=A,b=B,theoretical=theoretical_prob(A,B)]),
  reset_store,
  run_model(10_000,$model(A,B),[show_probs_trunc,mean]),
  nl,
  % show_store_lengths,
  fail,
  nl.
go => true.
  
strategy1(A,B) = uniform_draw(ones(A,w) ++ ones(B,b)).

strategy2(NumWhite,NumBlack) = Res =>
  Total = NumWhite + NumBlack,
  % A draws one ball and remove it
  Ball1 = categorical([NumWhite / Total, NumBlack / Total],[w,b]),
  NumWhite2 = cond(Ball1 == w,NumWhite-1,NumWhite),
  NumBlack2 = cond(Ball1 == b,NumBlack-1,NumBlack),

  % B remove one black ball
  NumWhite3 = NumWhite2,
  NumBlack3 = NumBlack2-1,

  Total3 = Total - 2, % We have now removed 2 balls
  % A draw again, If white: win else loss
  Ball = categorical([NumWhite3 / Total3, NumBlack3/Total3],[w,b]),
  Res = Ball.
  

model(A,B) =>
  Ball1 = strategy1(A,B),
  Win1 = cond(Ball1 == w,1,0),
  
  Ball2 = strategy2(A,B),
  Win2  = cond(Ball2 == w,1,0),

  S2BeatsS1 = cond(Win2 > Win1,1,0),

  add("1 ball",Ball1),
  add("1 win",Win1),  
  add("2 ball",Ball2),
  add("2 win",Win2),
  add("s2 beats s1",S2BeatsS1).