Changed some code to look nicer