Added some code to main based on Oscar's suggestions