Tidied up code for merge with master