Refactor and optimisations