Tidying up and removing non-necessary code