///fold:
var getProbs = function(vector) {
  return map(function(i) {return T.get(vector,i)}, _.range(vector.length))
}

///
var actionPrior = function() {
  return flip(.7) ? ['a'] : ['a'].concat(actionPrior());
}

var goalPosterior = Infer({method: 'rejection', samples: 5000}, function() {
  var buttonsToOutcomeProbs = mem(function(buttons) {return getProbs(dirichlet(ones([2,1])))});

  var vendingMachine = function(state, action) {
    return categorical({vs: ['bagel', 'cookie'], ps: buttonsToOutcomeProbs(action)});
  };

 //  Original inner inference
//   var chooseAction = function(goalSatisfied, transition, state) {
//   return Infer({method: 'rejection', samples: 1}, function() {
//     var action = actionPrior()
//     condition(goalSatisfied(transition(state, action)))
//     return action;
//   })
// }

 //Just generate a sample recursively, without using infer.
  var chooseAction = function(goalSatisfied, transition, state) {
    var action = actionPrior()
    return (goalSatisfied(transition(state, action)) ? action : chooseAction(goalSatisfied, transition, state))
  }

  var goal = categorical({vs: ['bagel', 'cookie'], ps: [.5, .5]});
  var goalSatisfied = function(outcome) {return outcome == goal;};
  //var chosenAction = sample(chooseAction(goalSatisfied, vendingMachine, 'state'));
  var chosenAction = chooseAction(goalSatisfied, vendingMachine, 'state');

  condition(goal == 'cookie' && _.isEqual(chosenAction, ['a', 'a']));

  return {once: buttonsToOutcomeProbs(['a'])[1],
          twice: buttonsToOutcomeProbs(['a', 'a'])[1]};
});

print("probability of actions giving a cookie")
viz.marginals(goalPosterior);