tidy up / refactor