bug fixes, code tidy up