Update to code to tidy it up