code tidy up

code tidy up