tidied up code
tidied up code