cleaning up of the code as advised