tidy up of code

tidy up of code