Bug fix with probability calculations