moved some code in to the main class