# Little class to help us with hypergeometric probabilities.
#
# 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
class HypergeometricProbabilities
  attr_accessor :marbles, :white, :picks

  def initialize(options)
    @marbles  = options[:marbles]
    @white    = options[:white]
    @picks    = options[:picks]
  end

  # Computes the number of ways you can represent all the white marbles -
  # since marble order doesn't matter, this tells us how many ways we could
  # draw all the white marbles and still get a perfect "win".  In a
  # lottery-specific situation this is also useful for the probability calc,
  # but since this has expanded to all hypergeometric distributions, this
  # function is really just useful for informational purposes.
  def get_white_permutations
    return combin(@marbles, @white)
  end

  # Computes the number of ways to draw all our picks.  Useful for information
  # as well as the final probability calc.
  def get_pick_permutations
    return combin(@marbles, @picks)
  end

  # Yields # of matches, probability of occurring for each possibility of
  # matches, from [white, picks].min down to zero.  Useful for letting our
  # caller spew data, store it in a hash, array, whatever.
  def get_all_odds
    # Get total permutations once.  Yeah, optimizing slow Ruby code is silly,
    # but this is an easy optimization.
    picked_perms = get_pick_permutations

    # Calculate matches down to 0
    [@white, @picks].min.downto 0 do |matched|
      # Ways to arrange white marble pulls to win
      white_wins  = combin(@white, matched)
      # Ways to pull the black marbles to win
      black_wins  = combin(@marbles - @white, @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
      probability = 0
      unless white_wins.nil? || black_wins.nil?
        probability = (white_wins * black_wins) / picked_perms.to_f
      end

      yield(matched, probability)
    end
  end

  private

  # Slow but effective recursive factorial function.  This could use a lookup
  # table and be MUCH faster, but I don't really care much until I actually
  # need this class to be fast (in which case I'll probably not use Ruby)
  def factorial(num)
    return nil if num < 0
    return 1 if num <= 1
    return num * factorial(num - 1)
  end

  # Basic combination function.  Very slow as it uses the recursive factorial
  # function.  But oh well.
  def combin(n, k)
    return 0 if (n < k || n < 0 || k < 0)
    return factorial(n) / (factorial(n - k) * factorial(k) )
  end

end
