Tidied up last night's code