Tidy up
Tidy up of code