Fixes and tidying up