Basic refactoring and tidy-up