Code tidy up