Tidied up code for assignment two