Code Tidy Up

Just a tidy up of code layout