Code tidied up