#include "HGProbabilityCalc.h"

// LICENSE: This code is freeware.  You may use this code for any purpose you
// like so long as this license is attached.  If this code is used in a
// commercial application, you must send me an email so I feel uber-geeky.  No
// other compensation is required.  Email me at jechols@nerdbucket.com

// Sets up our basic urn situation: number of marbles, number we'll pick, and
// how many of the marbles are white.
//
// See http://en.wikipedia.org/wiki/Hypergeometric_distribution for more
// information about hg dist.
HGProbabilityCalc::HGProbabilityCalc(int marbles, int picks, int white):
  m_total_marbles(marbles),
  m_picks(picks),
  m_white_marbles(white)
{
  calculate_permutations();
}

// Computes the number of ways to draw all our picks.  Useful for information
// as well as the final probability calc.
void HGProbabilityCalc::calculate_permutations() {
  m_pick_permutations = JEStatistics::combin(m_total_marbles, m_picks);
}

// Computes the odds of a certain number of white marbles being picked
// out of our already-built urn.
mpf_class HGProbabilityCalc::odds_against(int matched) {
  // Ways to arrange white marble pulls to win
  mpz_class white_wins = JEStatistics::combin(m_white_marbles, matched);
  // Ways to pull the black marbles to win
  mpz_class black_wins = JEStatistics::combin(m_total_marbles - m_white_marbles, m_picks - matched);

  // Probability is the multiplication of white * black permutations,
  // divided by total permutations.  To win you must have one of the white
  // marble combos and one of the black marble combos.  So the number of
  // permutations that give you both are just the multiplication of the
  // white and black permutations.  It's almost easy to understand!

  // Get total permutations for getting this number of matches
  // If we have a nil value, we know something went "wrong" and
  // probability is 0
  mpf_class probability = 0.0;

  // If either calculation was bad, we know probability for this case is
  // nonexistent
  if (white_wins > 0 && black_wins > 0) {
    // Two-step process to ensure our calculation is done a float.
    // Using our cache for pick permutations appears unnecessary,
    // but for many statistical problems this class solves, you'll call
    // odds_against more than once.  Think of state lotteries or keno where
    // you may want odds for all possible matches.
    probability = (white_wins * black_wins);
    probability /= m_pick_permutations;
  }

  return probability;
}

// Computes the odds of a all possible numbers of white marbles being picked
// out of our already-built urn.  Uses the callback function to let the caller
// deal with the results of each iteration - like yield in Ruby.
void HGProbabilityCalc::get_all_odds(ProbabilityOddsCallback& poc) {
  // Normally we iterate from white marbles down to 0, but if this is a poker
  // situation there may be more white marbles than picks, which would make
  // for some illogical computations.  Pick whichever is lower.
  const int lower = m_picks > m_white_marbles ? m_white_marbles : m_picks;
  for (int matches = lower; matches >= 0; --matches)
  {
    poc(matches, odds_against(matches));
  }
}