add some comments and tidy up a little