Refactors and tidy up